open++

test_util.py

import unittest
import tempfile
from StringIO import StringIO

from util import open_filepath_or, _is_string_or_bytes


class TestFilePathOpening(unittest.TestCase):
    def test_is_string_or_bytes(self):
        self.assertTrue(_is_string_or_bytes('foo'))
        self.assertTrue(_is_string_or_bytes(u'foo'))
        self.assertTrue(_is_string_or_bytes(b'foo'))
        self.assertFalse(_is_string_or_bytes(StringIO('bar')))
        self.assertFalse(_is_string_or_bytes([1]))

    def test_file_closed(self):
        """File gets closed in decorator"""
        f = tempfile.NamedTemporaryFile('r')
        filepath = f.name
        with open_filepath_or(filepath) as fh:
            pass
        self.assertTrue(fh.closed)

    def test_file_closed_harder(self):
        """File gets closed in decorator, even if exceptions happen."""
        f = tempfile.NamedTemporaryFile('r')
        filepath = f.name
        try:
            with open_filepath_or(filepath) as fh:
                raise TypeError
        except TypeError:
            self.assertTrue(fh.closed)
        else:
            # If we're here, no exceptions have been raised inside the
            # try clause, so the context manager swallowed them. No
            # good.
            raise Exception("`open_filepath_or` didn't propagate exceptions")

    def test_filehandle(self):
        """Filehandles slip through untouched"""
        with tempfile.TemporaryFile('r') as fh:
            with open_filepath_or(fh) as ffh:
                self.assertTrue(fh is ffh)
            # And it doesn't close the file-handle
            self.assertFalse(fh.closed)

    def test_StringIO(self):
        """StringIO (useful e.g. for testing) slips through."""
        f = StringIO("File contents")
        with open_filepath_or(f) as fh:
            self.assertTrue(fh is f)

util.py

from contextlib import contextmanager
from future.builtins import bytes, str


def _is_string_or_bytes(s):
    """Returns True if input argument is string (unicode or not) or bytes.

    """
    return isinstance(s, str) or isinstance(s, bytes)


def _get_filehandle(filepath_or, *args, **kwargs):
    """Open file if `filepath_or` looks like a string/unicode/bytes, else
    pass through.

    """
    if _is_string_or_bytes(filepath_or):
        fh, own_fh = open(filepath_or, *args, **kwargs), True
    else:
        fh, own_fh = filepath_or, False
    return fh, own_fh

@contextmanager
def open_filepath_or(filepath_or, *args, **kwargs):
    """Context manager, like ``open``, but let's file handles and file
    like objects pass untouched.

    It is useful when implementing a function that can accept both
    strings and file-like objects (like numpy.loadtxt, etc).

    Parameters
    ----------
    filepath_or : str/bytes/unicode string or file-like
         If string, file to be opened using ``open``. Else, it is returned
         untouched.

    Other parameters
    ----------------
    When `filepath_or` is a string, any extra arguments are passed on to
    ``open``.

    Examples
    --------
    >>> with open_filepath_or('filename') as f:  # doctest: +SKIP
    ...     pass
    >>> fh = open('filename')                    # doctest: +SKIP
    >>> with open_filepath_or(fh) as f:          # doctest: +SKIP
    ...     pass
    >>> fh.closed                                # doctest: +SKIP
    False
    >>> fh.close()                               # doctest: +SKIP

    """
    fh, own_fh = _get_filehandle(filepath_or, *args, **kwargs)
    try:
        yield fh
    finally:
        if own_fh:
            fh.close()