Writing Context Managers with contextlib

New words in this topic, in plain English. Come back here whenever one feels fuzzy.

Context manager

An object that does setup when a with block starts and cleanup when it ends.

with statement

Runs a block between the context manager's setup and cleanup.

__enter__ / __exit__

The two special methods behind with: one at the start, one at the end.

@contextmanager

A decorator that turns a one-yield generator function into a context manager.

try / finally

Code in finally runs whatever happens in try, even after an exception (Topic 7.1).

ExitStack

A container that can hold any number of context managers and clean them up in reverse order.

suppress

A context manager that quietly ignores the named exception types.

Buffer (StringIO)

An in-memory text object that behaves like a file.

1. 1

   Time a block without printing the time

   Write a @contextmanager named timed(label) that prints start label, yields, and prints done label in finally. Use it around a loop. In real code you would record time.perf_counter() before and after and log the difference; here we print fixed text so the output never changes.

2. 2

   Open three files at once

   Extend the ExitStack example to write and read three files. Check that all(f.closed ...) is still True.

Errors are the best teachers. Make each change, read the error, guess what went wrong, then reveal the answer.

from contextlib import contextmanager

@contextmanager
def announce(name):
    print(f"start {name}")
    yield name
    print(f"end {name}")

try:
    with announce("risky"):
        raise ValueError("boom")
except ValueError:
    print("caught")

1. 1

   You met context managers as a user in Topic 7.4: with open(path) as f:. The with statement calls __enter__ at the start, runs the block, and calls __exit__ at the end, even if an exception happened. Writing a class with those two methods works, but contextlib.contextmanager does it in a few lines with a generator.

2. 2

   Recall from Topic 9.2: a generator pauses at each yield and resumes later. @contextmanager turns a generator function with **exactly one yield** into a context manager. Code before the yield is the setup (__enter__). The value yielded is what as receives. Code after the yield is the cleanup (__exit__). So a generator pauses at the yield while the with body runs, then continues afterwards.

3. 3

   Put the yield inside try: and the cleanup in finally:. Without try/finally, an exception in the with body is raised at the yield, the generator ends early, and the cleanup lines after the yield never run. This is the most important rule of this topic, and the break-it example shows it.

4. 4

   suppress(ErrorType) ignores the named errors inside the block: with suppress(FileNotFoundError): os.remove(path). It reads better than an empty try/except. nullcontext() does nothing: it stands in when a context manager is optional (with lock or nullcontext():). closing(obj) calls obj.close() at the end for objects that have a close method but are not context managers themselves.

5. 5

   ExitStack manages a changing number of context managers or cleanup steps. You open files in a loop with stack.enter_context(open(name)), register extra cleanup with stack.callback(func, ...), and when the with ExitStack() block ends, everything is closed in the reverse order, even if one of them fails. Use it when you cannot write a fixed number of with items.

6. 6

   redirect_stdout(target) sends print output into a buffer such as io.StringIO(), which is useful in tests and for capturing what a noisy function prints. In AI and data code, custom context managers are common for timing a block, holding a database transaction, temporarily changing a setting, or opening and closing a connection to a model service or vector database.

• ↔

  Short generator form versus a full class. @contextmanager is short and clear for one-off setup/cleanup, but a class is better when the manager is reused, has extra methods or needs fine control over exceptions.

• ↔

  Guaranteed cleanup versus hidden control flow. with makes cleanup certain, but a context manager that swallows exceptions can hide bugs. Swallow errors only on purpose, and name the manager so that is obvious (suppress).

• Done when you can explain what __enter__ and __exit__ do in a with block.

• Done when you can write a context manager with @contextmanager using try/finally.

• Done when you can explain why the cleanup is skipped without finally.

• Done when you can use suppress, nullcontext and closing for their jobs.

• Done when you can use ExitStack to open a list of files.