Context managers: the cleanup that runs even when the block raises

The problem with solves

f = open(path, "w")
for item in items:
    f.write(process(item) + "\n")
f.close()

If process raises on the fortieth item, f.close() never runs. What happens to the thirty-nine lines already written depends on buffering: Python collects writes in memory and sends them to disk when the buffer fills or the file is closed. A short run that never closes can leave you with an empty file and a log full of successes. The file is closed when the process exits, but if the interpreter is shutting down because of an unhandled exception, whether that flush happens is not something to rely on.

The fix you already know:

with open(path, "w") as f:
    for item in items:
        f.write(process(item) + "\n")

with promises that the file's cleanup runs whether the block finishes, raises, or is left by return or break. This lesson is about what that promise is made of and how to make it for your own resources.

The protocol

A context manager is any object with __enter__ and __exit__. with calls __enter__ first and binds its return value to the name after as. When the block ends — by any route — it calls __exit__ with three arguments describing the exception, or None, None, None if there was none.

class Timer:
    def __enter__(self):
        self.start = time.perf_counter()
        return self
    def __exit__(self, exc_type, exc, tb):
        self.elapsed = time.perf_counter() - self.start
        log.info("block took %.2fs%s", self.elapsed, " (failed)" if exc else "")
        return False

Returning False from __exit__ lets the exception propagate after cleanup. Returning True suppresses it, which is appropriate only for a manager whose purpose is to suppress — contextlib.suppress(FileNotFoundError) exists for that. Suppressing by accident turns a crash into silent wrong output.

The generator form

Writing two methods for every small resource is tedious. contextlib.contextmanager turns a generator with exactly one yield into a context manager:

from contextlib import contextmanager

@contextmanager
def timed(label):
    start = time.perf_counter()
    try:
        yield
    finally:
        log.info("%s: %.2fs", label, time.perf_counter() - start)

with timed("embedding 500 chunks"):
    vectors = embed(chunks)

Everything before yield is __enter__; everything after is __exit__. The try/finally is what makes the cleanup run on an exception — without it, an exception in the block would skip the code after yield. Whatever you yield becomes the as value.

A resource-shaped one:

@contextmanager
def api_session(key):
    s = requests.Session()
    s.headers["Authorization"] = f"Bearer {key}"
    try:
        yield s
    finally:
        s.close()

What belongs in one

Anything that must be undone: an open file, a network session, a database connection or transaction, a lock, a temporary directory, a changed working directory, a changed environment variable, a timer, a progress bar. The test is: "if the code in the middle blows up, is there something that must still happen?" If yes, it is a context manager.

The standard library is full of them. tempfile.TemporaryDirectory() creates a folder and deletes it at the end. threading.Lock() can be used as with lock: so the lock is released even on error, which is the difference between a bug and a deadlock. sqlite3 connections used with with conn: commit on success and roll back on exception — a transaction in one line. contextlib.redirect_stdout(buffer) captures prints, useful in tests. unittest.mock.patch from module 4 is one, and that is why its effect ends at the close of the block.

Several at once

with open(src) as fin, open(dst, "w") as fout:
    for line in fin:
        fout.write(transform(line))

Comma-separated managers are entered left to right and exited right to left. If the second open fails, the first is still closed properly.

When the number of resources is not known until runtime — one file per output category, say — contextlib.ExitStack collects them:

from contextlib import ExitStack

with ExitStack() as stack:
    files = {cat: stack.enter_context(open(f"{cat}.txt", "w")) for cat in categories}
    for item in items:
        files[item.category].write(item.text + "\n")

Every file is closed when the block ends, in reverse order, however it ends.

Async

async with is the same protocol with __aenter__ and __aexit__, for resources whose setup or teardown awaits — httpx.AsyncClient() in module 6 was one. contextlib.asynccontextmanager is the generator form. The rules do not change: cleanup after yield, try/finally around it.

The mistake to avoid

Do not put the body of the work inside the manager's own code. A context manager sets up, yields, and tears down; the block is the caller's. A contextmanager that does a loop of API calls before its yield is a function wearing a costume, and its cleanup will not cover the block that matters.

Try this now

Write timed(label) and wrap a model call in it. Then write a changed_env(name, value) manager that sets an environment variable and restores the old value — including when there was none — and prove with a deliberate raise inside the block that the variable is restored anyway.