A context manager for files or file-like objects

I usually design my Python programs so that if a program needs to read or write to a file, the functions will take a filename argument that can be either a path string or a file-like object already open for reading / writing.

(I think I picked up this habit from Mark Pilgrim’s Dive Into Python, in particular chapter 10 about scripts and streams.)

This has the great advantage of making tests easier to write. Instead of having to create dummy temporary files on disk I can wrap strings in StringIO() and pass that instead.

But the disadvantage is I then have a bit of boiler-plate at the top of the function:

def read_something(filename):
    # Tedious but not heinous boiler-plate
    if isinstance(filename, basestring):
        filename = open(filename)

    return filename.read()

The other drawback is that code doesn’t close the file it opened. You could have filename.close() before returning but that will also close file-like objects that were passed in, which may not be what the caller wants. I think the decision whether to close the file belongs to the caller when the argument is a file-like object.

You could set a flag when opening the file, and then close the file afterwards if the flag is set, but that is yet more boiler-plate and quite ugly.

So here is a context manager which behaves like open(). If the argument is a string it handles opening and closing the file cleanly. If the argument is anything else then it just reads the contents.

class open_filename(object):
    """Context manager that opens a filename and closes it on exit, but does
    nothing for file-like objects.
    """
    def __init__(self, filename, *args, **kwargs):
        self.closing = kwargs.pop('closing', False)
        if isinstance(filename, basestring):
            self.fh = open(filename, *args, **kwargs)
            self.closing = True
        else:
            self.fh = filename

    def __enter__(self):
        return self.fh

    def __exit__(self, exc_type, exc_val, exc_tb):
        if self.closing:
            self.fh.close()

        return False

And then you use it like this:

from io import StringIO

file1 = StringIO(u'The quick brown fox...')
file2 = 'The quick brown fox'

with open_filename(file1) as fh1, open_filename(file2) as fh2:
    foo, bar = fh1.read(), fh2.read()

If you always want the file to be closed on leaving the block you use the closing keyword argument set to True (the default of False means the file will only be closed if it was opened by the context manager).

file1 = StringIO(u'...jumps over the lazy dog.')
assert file1.closed == False

with open_filename(file1, closing=True) as fh:
    foo = fh.read()

assert file1.closed == True

Today is my brother’s birthday. If I had asked him what he wanted for a present I am pretty certain he would have asked for a blog post about closing files in a computer programming language.