Python 101Free
IDIOMATIC PYTHON

Context Managers

Guaranteeing cleanup with the with statement.

SECTION 01

with and the protocol

A context manager is an object that knows how to set up and tear down around a block of code. The with statement uses two methods: __enter__ runs at the start of the block, __exit__ runs at the end, even if the block raises an exception.

The canonical example is a file. with open(path) as f: ... opens the file at the start of the block and closes it at the end. The with statement promises the close runs no matter what happens inside.

Before with, the equivalent code was try / finally with explicit cleanup. The with form is shorter, harder to forget, and clearer about which lines are inside the resource and which are not. Reach for it whenever cleanup is part of using a thing.

python
with open("data.txt") as f:
    text = f.read()
# the file is closed here, even if read() raised
SECTION 02

Building one with contextlib

Writing __enter__ and __exit__ by hand on a class is verbose. The @contextmanager decorator from contextlib lets you write a context manager as a single generator function with one yield.

Code before the yield is the enter step. The yielded value is what the as clause binds to. Code after the yield is the exit step, run when the with block ends. A try / finally around the yield makes sure the cleanup happens even if the block raises.

This turns ten lines of class boilerplate into a five-line generator. For one-off context managers in your own code, this is the form to use. The class form is still appropriate when the context manager is a long-lived object with other state.

python
import os
from contextlib import contextmanager

@contextmanager
def cd(path):
    old = os.getcwd()
    os.chdir(path)
    try:
        yield               # the with-block runs here
    finally:
        os.chdir(old)

with cd("/tmp"):
    ...                     # work inside /tmp
← PREVIOUS
Decorators