Skip to content

Context Managers

Some resources need cleanup.

  • A file handle has to be closed.
  • A database connection has to be released back to the pool.
  • A lock has to be unlocked, even if the code holding it raises halfway through.

The with statement is Python's tool for tying setup and cleanup together so the cleanup always runs.

The with statement

Without with, opening and closing a file looks like this:

Python
f = open("tasks.txt")
try:
    contents = f.read()
finally:
    f.close()

The try/finally is there so that close() runs even if read() raises.

Warning

Forget the finally and a raised exception leaves the file open until the garbage collector gets to it.

with collapses the same pattern into one block:

Python
with open("tasks.txt") as f:
    contents = f.read()

When the block exits, whether normally or because of an exception, Python calls the cleanup code on f, without needing to manually specify it.

The file is closed before the next line runs.

The __enter__ and __exit__ protocol

Anything with __enter__ and __exit__ methods can be used in a with block. __enter__ runs at the start, and whatever it returns gets bound to the name after as. __exit__ runs at the end and receives details about any exception that was raised inside the block.

Python
class Section:
    def __init__(self, label: str) -> None:
        self.label = label

    def __enter__(self) -> "Section":
        print(f"--- {self.label} ---")
        return self

    def __exit__(self, exc_type, exc_value, traceback) -> None:
        print(f"--- end {self.label} ---")


with Section("normalize tasks"):
    print("doing work")

Writing a class for every context manager is more ceremony than most cases need. The contextlib module has a shortcut.

contextlib.contextmanager

@contextmanager lets you write a context manager as a generator function. Everything before yield is the setup, the yielded value is what gets bound to the as name, and everything after yield is the cleanup.

Python
import time
from contextlib import contextmanager


@contextmanager
def timer(label: str):
    start = time.perf_counter()
    try:
        yield
    finally:
        elapsed_ms = (time.perf_counter() - start) * 1000
        print(f"{label}: {elapsed_ms:.2f}ms")


with timer("normalize tasks"):
    normalized = [t["title"].strip().title() for t in raw_tasks]

The try/finally around yield is important.

Without it, a raised exception inside the with block would skip the cleanup. With it, the elapsed time prints whether the block succeeded or blew up.

Info

When we work with database sessions later, we'll use the same pattern: a with block opens a session, runs a few queries, and the cleanup commits or rolls back the transaction before closing the session. That keeps every database interaction scoped to a single block, instead of trusting every code path to remember to clean up.