Context Managers FTW

In their most basic form, Python's context managers make it easier to do the right thing by getting rid of some boiler-plate. For example, it was sometimes tempting to skip the try/finally to close a file object correctly, since Python would usually clean it up anyways when it went out of scope. But it's much easier to do it right this way:

with open('filename') as fp:
  contents = fp.read()

instead of this way:

fp = open('filename')
try:
  contents = fp.read()
finally:
  fp.close()

However, the more I work with context managers, the more they change the way I think about many different problems.

It's easy to extend this pattern into other places where you have a resource to clean up when you're done. I use this one a lot to create a temporary directory to work with, and then it's cleaned up at the end:

@contextlib.contextmanager
def temp_directory(*args, **kwargs):
  path = tempfile.mkdtemp(*args, **kwargs)
  try:
    yield path
  finally:
    shutil.rmtree(path)

with temp_directory() as tmp:
  # work with files in tmp, and they'll be cleaned up when you're done!

In many of these cases, the context manager isn't a revolutionary shift from just writing the try/finally code inline, but as you start building on the basic patterns, you can start to come up with some more interesting solutions.

I have a few places where I need to replace an existing file in a safe way. For example, I have a running web app that uses a file-based cache. I have a background process that updates the data for the cache, so I want to write the new data into a new file, and if it completes successfully, swap it in place of the old cache. I could do all this inline where I'm updating the cache, but it's much easier to abstract the details away so that it's easy to code once and reuse elsewhere.

@contextlib.contextmanager
def overwriting(path, prefix='', suffix='.tmp'):
  dirname, filename = os.path.split(path)
  tmp_path = os.path.join(dirname, prefix + filename + suffix)
  fd = os.open(tmp_path, os.O_EXCL | os.O_CREAT | os.O_WRONLY | os.O_TRUNC, 0644)
  f = os.fdopen(fd, 'w')
  try:
    yield f
    if not f.closed:
      f.flush()
    os.rename(tmp_path, path)
  except:
    os.remove(tmp_path)
    raise
  finally:
    f.close()

with overwriting('filename') as fp:
  fp.write('Hello world!')

In the case of the cache, it was critical to get the logic right so that I didn't swap out the cache until it was fully written, and didn't swap out the cache if there was an error somewhere in the middle, etc. But, this has also made it easy to take advantage of this in places I probably wouldn't have bothered with making as robust before. Soon after writing the overwriting context manager I needed to write some code to make some tweaks to an XML document and re-write it. This is what I ended up with:

@contextlib.contextmanager
def xml_mutator(filename):
  xml = ElementTree.parse(filename)
  yield xml
  with overwriting(filename) as fp:
    xml.write(fp)

It simply parses an XML file, gives you the ElementTree document to manipulate, and then writes back any changes you made to the document. Normally I would have just overwritten the original file directly, but now I get the nice behavior that if something goes wrong, I won't overwrite the original file with a broken file.

Other good examples I've encountered include mmapping files, or setting SIGALRM to run some code with a timeout. Being able to abstract some of the messy details of these APIs makes it easier to just get the pattern right once and reuse it. This has simplified my code quite a bit, and has made context managers one of my favorite language features the more I've started factoring them into the way I think about my code.