8 PEP 343: The 'with' statement

The 'with' statement clarifies code that previously would use try...finally blocks to ensure that clean-up code is executed. In this section, I'll discuss the statement as it will commonly be used. In the next section, I'll examine the implementation details and show how to write objects for use with this statement.

The 'with' statement is a new control-flow structure whose basic structure is:

with expression [as variable]:
    with-block

The expression is evaluated, and it should result in an object that supports the context management protocol. This object may return a value that can optionally be bound to the name variable. (Note carefully that variable is not assigned the result of expression.) The object can then run set-up code before with-block is executed and some clean-up code is executed after the block is done, even if the block raised an exception.

To enable the statement in Python 2.5, you need to add the following directive to your module:

from __future__ import with_statement

Some standard Python objects now support the context management protocol and can be used with the 'with' statement. File objects are one example:

with open('/etc/passwd', 'r') as f:
    for line in f:
        print line
        ... more processing code ...

After this statement has executed, the file object in f will have been automatically closed, even if the 'for' loop raised an exception part-way through the block.

The threading module's locks and condition variables also support the 'with' statement:

lock = threading.Lock()
with lock:
    # Critical section of code
    ...

The lock is acquired before the block is executed and always released once the block is complete.

The new localcontext() function in the decimal module makes it easy to save and restore the current decimal context, which encapsulates the desired precision and rounding characteristics for computations:

from decimal import Decimal, Context, localcontext

# Displays with default precision of 28 digits
v = Decimal('578')
print v.sqrt()

with localcontext(Context(prec=16)):
    # All code in this block uses a precision of 16 digits.
    # The original context is restored on exiting the block.
    print v.sqrt()

8.1 Writing Context Managers

Under the hood, the 'with' statement is fairly complicated. Most people will only use 'with' in company with existing objects and don't need to know these details, so you can skip the rest of this section if you like. Authors of new objects will need to understand the details of the underlying implementation and should keep reading.

Let's think through an example. I won't present detailed code but will only sketch the methods necessary for a database that supports transactions.

(For people unfamiliar with database terminology: a set of changes to the database are grouped into a transaction. Transactions can be either committed, meaning that all the changes are written into the database, or rolled back, meaning that the changes are all discarded and the database is unchanged. See any database textbook for more information.)

Let's assume there's an object representing a database connection. Our goal will be to let the user write code like this:

db_connection = DatabaseConnection()
with db_connection as cursor:
    cursor.execute('insert into ...')
    cursor.execute('delete from ...')
    # ... more operations ...

The transaction should be committed if the code in the block runs flawlessly or rolled back if there's an exception. Here's the basic interface for DatabaseConnection that I'll assume:

class DatabaseConnection:
    # Database interface
    def cursor (self):
        "Returns a cursor object and starts a new transaction"
    def commit (self):
        "Commits current transaction"
    def rollback (self):
        "Rolls back current transaction"

The __enter__() method is pretty easy, having only to start a new transaction. For this application the resulting cursor object would be a useful result, so the method will return it. The user can then add as cursor to their 'with' statement to bind the cursor to a variable name.

class DatabaseConnection:
    ...
    def __enter__ (self):
        # Code to start a new transaction
        cursor = self.cursor()
        return cursor

The __exit__() method is the most complicated because it's where most of the work has to be done. The method has to check if an exception occurred. If there was no exception, the transaction is committed. The transaction is rolled back if there was an exception.

In the code below, execution will just fall off the end of the function, returning the default value of None. None is false, so the exception will be re-raised automatically. If you wished, you could be more explicit and add a return statement at the marked location.

class DatabaseConnection:
    ...
    def __exit__ (self, type, value, tb):
        if tb is None:
            # No exception, so commit
            self.commit()
        else:
            # Exception occurred, so rollback.
            self.rollback()
            # return False

8.2 The contextlib module

The new contextlib module provides some functions and a decorator that are useful for writing objects for use with the 'with' statement.

The decorator is called contextmanager, and lets you write a single generator function instead of defining a new class. The generator should yield exactly one value. The code up to the yield will be executed as the __enter__() method, and the value yielded will be the method's return value that will get bound to the variable in the 'with' statement's as clause, if any. The code after the yield will be executed in the __exit__() method. Any exception raised in the block will be raised by the yield statement.

Our database example from the previous section could be written using this decorator as:

from contextlib import contextmanager

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

db = DatabaseConnection()
with db_transaction(db) as cursor:
    ...

The contextlib module also has a nested(mgr1, mgr2, ...) function that combines a number of context managers so you don't need to write nested 'with' statements. In this example, the single 'with' statement both starts a database transaction and acquires a thread lock:

lock = threading.Lock()
with nested (db_transaction(db), lock) as (cursor, locked):
    ...

Finally, the closing(object) function returns object so that it can be bound to a variable, and calls object.close() at the end of the block.

import urllib, sys
from contextlib import closing

with closing(urllib.urlopen('http://www.yahoo.com')) as f:
    for line in f:
        sys.stdout.write(line)