Kent's Korner

by Kent S Johnson

Introduction

A common need in programming is to guarantee some kind of cleanup after a block of code is run; for example, closing a file or database connection after use. The usual way to do this in Python prior to version 2.5 is with a try / except block. For example:

f = open('output.txt', 'w')
try:
  f.write('something')
finally:
  f.close()

In Python 2.5 and later, this can be written more simply as

with open('output.txt', 'w') as f:
  f.write('something')

(Python 2.5 requires that you add from __future__ import with_statement to the module that uses with.)

The general form of a with statement is

with <expression> [as <name>]:
  <block>

where as <name> is optional. The result of evaluating <expression> must be a context manager - an object that implements the context manager protocol (see below). In Python 2.5 and later, files, threading locks and Decimal contexts all implement the context manager protocol and may be used in with statements. Also it's not too hard to write your own context managers, by directly implementing the required protocol or using the contextlib.contextmanager decorator.

How does this work?

The plumbing underlying a with statement is slightly complicated. I will simplify a bit; for all the details see PEP 343.

• The expression in the with statement is evaluated. The result must be a context manager - an object with __enter__() and __exit__() methods.
• The context manager's __enter__() method is called. The result of calling __enter__() is assigned to the optional variable, if present.
• The nested block is executed.
• The context manager's __exit__() method is called. This method does any required cleanup, such as closing a file.

The __exit__() method is always called, even if the block raises an exception. In case of an exception, the exception details are passed to __exit__(), which may choose to process the exception in some way, to pass it to the enclosing block, or both.

Writing your own context manager

Fortunately you don't have to understand all of this to use context managers, or even to write your own; the standard library module contextlib contains a contextmanager decorator that converts a generator into a context manager. (If you don't understand decorators and generator functions, you should probably read up on them a little before continuing.) Here's how it works:

• Write a generator function which yields a single value.
• Everything up to the yield statement runs as the __enter__() method.
• The yielded value is returned from __enter__() and assigned to the user variable.
• Everything after the yield statement runs as the __exit__() method.
• The yield statement may be included in a try / catch / finally block for exception handling.

Examples

The What's New in Python 2.6 document includes a simple contextmanager example - a context manager to manage a database transaction. Below is a modified version.

Commonly when using a database, you get a cursor from the database connection and send commands to the database using the cursor. When you are done you close the cursor. If the commands are successful you commit the transaction; if any command fails the transaction is rolled back. You might write the code like this:

cursor = connection.cursor()
try:
  try:
    # do stuff with cursor
  finally:
    cursor.close()
except:
  connection.rollback()
  raise
else:
  connection.commit()

This quickly becomes tiresome to repeat throughout the client code. A context manager makes the client code much simpler:

from contextlib import contextmanager

@contextmanager
def db_transaction(connection):
  cursor = connection.cursor()
  try:
    try:
      yield cursor
    finally:
      cursor.close()
  except:
    connection.rollback()
    raise
  else:
    connection.commit()

Now the client code may be written very simply as

with db_transaction(connection) as cursor:
  # do stuff with cursor

This is simpler, more readable and more likely to be written correctly than the original client code.

Another example, based on blogmaker code, uses a context manager to trap and log exceptions:

from contextlib import contextmanager
import logging

@contextmanager
def error_trapping():
    ''' A context manager that traps and logs exception in its block.
        Usage:
        with error_trapping():
            might_raise_exception()
        this_will_always_be_called()
    '''
    try:
        yield None
    except Exception:
        logging.error('Error', exc_info=True)

Special bonus example - the error_trapping context manager can itself be wrapped in a decorator. When this decorator is applied to a function, and exceptions raised within the function will be trapped and logged:

from functools import wraps

def trap_errors(f):
    ''' A decorator to trap and log exceptions '''
    @wraps(f)
    def wrapper(*args, **kwds):
        with error_trapping():
            return f(*args, **kwds)
    return wrapper

References

• PEP 343 is the specification for the with statement.
• contextlib contains the contextmanager decorator
• What's New in Python 2.6 has a good explanation of the with statement and contextmanager decorator.