Context managers

Guarantee that cleanup happens — on success, on an exception and on an early return — without writing finally everywhere.

20of 24
≈ 12Min read
5Parts

with, and what it guarantees

with calls __enter__ on the way in and __exit__ on the way out — after a return, after a break, after an exception, always. It is the try/finally you would otherwise have to remember, written once inside the object instead of at every call site.

py
with open("notes.txt", encoding="utf-8") as f:
    data = f.read()
    return process(data)        # the file is still closed

with lock, open(path) as f:     # several at once, closed in reverse order
    ...

The protocol

__exit__ receives the exception if there was one, and its return value decides what happens next: a falsy value lets the exception continue, a truthy one swallows it. Returning True by accident silently hides every error inside the block, which is a very expensive typo.

py
class Timer:
    def __enter__(self):
        self.start = time.perf_counter()
        return self                    # whatever "as x" binds

    def __exit__(self, exc_type, exc, traceback):
        self.elapsed = time.perf_counter() - self.start
        return False                   # False (or None): let any exception continue

with Timer() as t:
    work()
print(t.elapsed)

@contextmanager — the short way

One generator, one yield, and the try/finally around it. The finally is not optional: without it, an exception in the block skips everything after the yield and your cleanup never runs.

py
from contextlib import contextmanager

@contextmanager
def timer():
    start = time.perf_counter()
    try:
        yield start              # everything before yield is __enter__
    finally:                     # everything after is __exit__
        print(time.perf_counter() - start)

with timer():
    work()

The contextlib shelf

ExitStack is the answer to "I need to open N files and close them all", and nullcontext is the answer to "this argument is sometimes a context manager and sometimes already open" — both replace code that is otherwise genuinely awkward.

py
from contextlib import suppress, redirect_stdout, ExitStack, closing, nullcontext

with suppress(FileNotFoundError):        # try/except/pass, said properly
    Path("cache.tmp").unlink()

with redirect_stdout(io.StringIO()) as out:
    noisy()                              # captured, not printed

with ExitStack() as stack:               # an unknown number of contexts
    files = [stack.enter_context(open(p)) for p in paths]

def read(path, opener=None):
    with (opener or nullcontext(sys.stdin)) as f:   # "sometimes a context"
        ...

Beyond files

The pattern is always the same: two things that must happen in pairs. Any time you write "and afterwards, always", you have found one.

  • Locks and transactions — acquire and release, begin and commit or roll back.
  • Temporary state — change a setting, restore it; decimal.localcontext, numpy.errstate.
  • Timing and tracing — start a span, end it however the block ends.
  • Test doubles — unittest.mock.patch is a context manager, and so is pytest.raises.

Practice