Python Context Managers

Sometimes it's helpful to wrap some code in a given context. Either the code requires some resource from the context, or the context adds some aspect to the code.

Python provides with statement to create a context for the followed block of code.

with context_manager:
    # this block runs within the context
    ...

or if the context provides a resource to the enclosed code block:

with context_manager as ctx_resource:
    # this block can use the resource provided to the context using thee
    # ctx_resource object
    ...

The context is defined by the related context manager. In other words it's the context manager that gives meaning to the context, enclosing the block of code following the with statement.

It's common that the context manager is created by factory functions, accepting arguments to configure the context with.

def create_context(arg):
    # returns a context manager, possibly configured by provided arg
    ...

with create_context(arg):
    ...

Let's demonstrate some concrete examples for what's been discussed so far.

The standard library provides general purpose context managers in contextlib module.

contextlib.closing closes resources like files, sockets, etc (technically any object that has a close() method), and is a nice way to guarantee resources are always closed properly.

from contextlib import closing

with closing(obj):
    ...
    # obj.close() is called automatically when this code block is done,
    # even when exceptions are raised

contextlib.suppress provides an easy way to suppress exceptions that the code can afford to ignore.

from contextlib import suppress

with suppress(IOError, OSError, RuntimeError):
    ...
    # If this block raises any of the above exception types,
    # it'll be caught and ignored, execution continues after the context block

The built-in open function could be used as a context manager providing the file object to the context code block, while ensuring that the file is closed when the context is finished.

with open('/path/to/file', 'r') as file_obj:
    file_obj.read()

Context Managers

There's nothing special about the context manager, it can be any object that implements the __enter__ and __exit__ methods to initialize and finalize the context respectively. Python calls __enter__ method of the context manager before execution continues with the context code block (suite), and always calls __exit__ even if an exception was raised from the block.

The contextlib module provides a shortcut to create context managers, the @contextmanager decorator.

from contextlib import contextmanager

@contextmanager
def my_context():
    ...
    # the statements above run before the context, to initialize the context
    yield
    # the statements below run after the context, to finalize the context
    ...

with my_context():
    ...

Under the hood @contextmanager creates a context manager that wraps a generator. The context manager starts the generator on __enter__, and continues the generator on __exit__, So the generator defines the logic of initialization and finalization of the context. That's why the my_context function should not return but yield (it's a generator function).

If an unhandled exception happens in the block, the wrapping context manager's __exit__ method receives it (according to with statement contract) and re-raises it inside the generator, after the yield.

The context manager created by previous example did not handle exceptions, so if an exception is raised in the context block, the finalization statements won't get a chance to run. We should be explicit about such behavior when using @contextmanager decorator.

@contextmanager
def my_context():
    ...
    try:
        yield
    finally:
        # the statements below run after context code block, and if an exception
        # is raised within the context block
        ...

For example we could create a context manager that logs the beginning and end of the given code block. The factory accepts a name for the context as well, so that it's clear which log belongs to which suite.

from logging import getLogger
from contextlib import contextmanager

@contextmanager
def autolog_context(name):
    logger = getLogger('autolog')
    logger.info('"{}" started'.format(name))
    try:
        yield
    except:
        logger.exception('error occurred in "{}"'.format(name))
        raise  # re-raise the original exception
    logger.info('"{}" completed successfully'.format(name))


with autolog_context('page_generation'):
    ...

with autolog_context('unfortunate_task'):
    ...
    raise ValueError()  # this exception is auto logged then propogated
    # the statements below won't run
    ...

The Context Resource

Sometimes the context provides a resource to be used by the context block. The built-in open function is a good example:

with open('/tmp/somefile.txt', 'wt') as fh:
    # fh is the resource the context manager provided to the suite
    fh.write('')

To provide the resource to the context, the return value of __enter__ from the context manager is used.

When @contextmanager decorator is used to create the context manager, the value that the generator yields is the context resource.

For example we'd like to modify our autolog_context so it provides the logger object to the context block.

@contextmanager
def autolog_context(name):
    logger = getLogger('autolog')
    ...
    try:
        yield logger
    except:
        ...

with autolog_context('page_generation') as logger:
    ...
    logger.info('still in the block')
    ...

There maybe cases where the context populates some data which is still required after the context is finished. One approach is to use a mutable type (for example a list or a dict) for the context resource. After the context manager provided the resource to the block, it can still mutate it (append data) during finalization, and since it's the same object reference, the change is available to the calling scope.

For example this context manager records memory usage of current process before and after the block, logs the diff, and the data still stays in scope after the context exits.

from contextlib import contextmanager
from logging import getLogger
from resource import getrusage, RUSAGE_SELF

@contextmanager
def capture_mem_usage(name):
    mem_before = getrusage(RUSAGE_SELF).ru_maxrss
    mem_samples = []
    yield mem_samples
    mem_after = getrusage(RUSAGE_SELF).ru_maxrss
    mem_samples.insert(0, mem_before)  # the context may have modified the list
    mem_samples.append(mem_after)
    mem_diff = mem_after - mem_before
    getLogger('capture_mem_usage').info('"{}" used "{}" bytes of memory'.format(name, mem_diff))


with capture_mem_usage('process_reports') as mem_usages:
    ...

mem_before, mem_after = mem_usage[0], mem_usage.pop()