Python context managers

Context managers

In Python, a context manager is an object that can be used in a with statement. Here's a context manager that reports the total wall-clock time spent inside a with block:

import time

class Timer(object):
    def __init__(self, msg):
        self._msg = msg

    def __enter__(self):
        # time.monotonic() requires Python >= 3.3
        self._start = time.monotonic()

    def __exit__(self, exc_type, exc_value, exc_traceback):
        if exc_type:
            print('Failed: {}: {}'.format(self._msg, exc_value))
        else:
            print('{}: {} s'.format(self._msg, time.monotonic() - self._start))

We can use it like this:

with Timer("doing stuff"):
    for i in range(1000000):
        pass

which produces this output:

doing stuff: 0.04281306266784668 s

Handling exceptions

Context managers can also handle exceptions from inside the block they guard. Timer does this. If we do:

with Timer("doing stuff"):
    raise ValueError('ack')

we get:

Failed: doing stuff: ack
Traceback (most recent call last):
  File "test.py", line 20, in <module>
    raise ValueError('ack')
ValueError: ack

Because __exit__() doesn't return True, the exception continues to propagate, so we see both the Failed message and the traceback. If __exit__() returned True, the exception would be swallowed by the context manager and we would see only the Failed line.

Context from the context manager

Let's say we wanted the context manager to return some information, perhaps the granularity of the timer. We can modify __enter__() like this:

def __enter__(self):
    # time functions require Python >= 3.3
    self._start = time.monotonic()
    return time.clock_getres(time.CLOCK_MONOTONIC)

Now we can do:

with Timer("doing stuff") as resolution:
    print('Resolution: {}'.format(resolution))
    for i in range(1000000):
        pass

which produces:

Resolution: 1e-09
doing stuff: 0.043778783998277504 s

The `contextmanager` decorator

Writing a context manager involves some boilerplate: the __enter__() and __exit__() methods, and possibly an __init__() method to handle any arguments. There's an easier way: the contextlib.contextmanagerdecorator. We can rewrite Timer like this:

import contextlib
import time

@contextlib.contextmanager
def Timer(msg):
    start = time.monotonic()
    try:
        yield time.clock_getres(time.CLOCK_MONOTONIC)
    except BaseException as e:
        print('Failed: {}: {}'.format(msg, e))
        raise
    else:
        print('{}: {} s'.format(msg, time.monotonic() - start))

Now Timer is just a generator that yields once (yielding the value to be bound by the with statement). The yield expression raises any exception thrown from the block, which the context manager can then handle or not.

Self-management

One use for context managers is ensuring that an object (file, network connection, database handle) is closed when leaving a block. contextlib.closing() does this: its __exit__() method calls another object's close(). Use it like this:

import contextlib

with contextlib.closing(open("/etc/passwd")) as fh:
    lines = fh.readlines()

But that seems a bit verbose. A better approach: many objects implement __enter__() and __exit__(), so you can use them to manage themselves. The implementation of file (the type returned by open()) has something like this:

def __enter__(self):
    return self

def __exit__(self, exc_type, exc_value, exc_traceback):
    self.close()

So you can just do:

with open("/etc/passwd") as fh:
    lines = fh.readlines()

Exercise!

Reimplement the contextmanager decorator. You'll need to use some of the more obscure features of generators.

bgilbert/contextmanager.md