Monthly Archives: June 2012

A context manager for files or file-like objects

I usually design my [Python][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][dive], in particular [chapter 10 about scripts and streams][chap10].)

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()`][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][ctxt] 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.

[dive]: http://www.diveintopython.net/
[python]: http://www.python.org/
[stringio]: http://docs.python.org/library/io.html
[ctxt]: http://docs.python.org/library/stdtypes.html#typecontextmanager
[chap10]: http://www.diveintopython.net/scripts_and_streams/index.html