lesson 165 Decorators & Context Managers · advanced

contextlib: context managers from generators

contextlib.contextmanager turns a generator with one yield into a context manager: setup before the yield, cleanup after it in a finally. contextlib also ships ready-made managers: suppress, redirect_stdout, and ExitStack for a varying number of resources.

You'll be able to

Write a @contextmanager generator whose cleanup is in finally
Use contextlib.suppress and redirect_stdout instead of writing a class
Open a variable number of files safely with ExitStack

The idea

A class with __enter__ and __exit__ is a lot of code for "do this before, undo it after". `contextlib.contextmanager` builds the class for you from a generator function: the code before yield is the enter part, the yielded value is what as binds, and the code after yield is the exit part. If the block raises, the exception is raised at the `yield`, so the cleanup must be in a finally to run in that case.

Run each example to see what it prints. Running examples doesn't count toward progress.

example-1.py · A timing block from a generator

from contextlib import contextmanager

@contextmanager
def timed(label, clock, results):
    start = clock()
    try:
        yield
    finally:
        results[label] = clock() - start

ticks = iter([0.0, 1.5, 2.0, 2.25])
results = {}
with timed("download", lambda: next(ticks), results):
    pass
try:
    with timed("parse", lambda: next(ticks), results):
        raise ValueError("bad row 812")
except ValueError as err:
    print("error:", err)
print(results)

contextlib also has managers you don't need to write:

suppress(FileNotFoundError) ignores exactly those exceptions in its block;
redirect_stdout(buffer) sends everything printed in the block to buffer, for example an io.StringIO, which is how you capture a noisy function's output;
ExitStack() enters any number of context managers, chosen while the program runs, and exits all of them (in reverse order) when its own block ends: one with for a list of files.

example-2.py · suppress, redirect_stdout and ExitStack

import io
import os
from contextlib import ExitStack, redirect_stdout, suppress
from pathlib import Path

for name, text in [("ord.csv", "date,tmax_c\n2024-07-01,31.7\n"), ("lax.csv", "date,tmax_c\n2024-07-01,24.4\n")]:
    with open(name, "w", encoding="utf-8") as f:
        f.write(text)

with suppress(FileNotFoundError):
    os.remove("old-report.txt")
print("no old report, no problem")

def noisy_import():
    print("reading ord.csv")
    print("reading lax.csv")
    return 2

buffer = io.StringIO()
with redirect_stdout(buffer):
    count = noisy_import()
print(count, buffer.getvalue().splitlines())

paths = ["ord.csv", "lax.csv"]
with ExitStack() as stack:
    files = [stack.enter_context(Path(p).open(encoding="utf-8")) for p in paths]
    headers = [f.readline().strip() for f in files]
print(headers, [f.closed for f in files])

When to use it: @contextmanager for short setup/cleanup pairs; the ready-made managers wherever they fit; a class when the manager has several methods or lots of state. Mistakes to avoid: cleanup after yield without finally; yielding more than once (a context manager's generator must yield exactly once); wide suppress(Exception), which hides bugs; and opening files in a loop without ExitStack, which leaves the earlier ones open if a later one fails.

Review

@contextmanager: setup, then try: yield value and cleanup in finally:.
Exceptions from the block are raised at the yield; re-raise them unless handling them is the manager's job.
suppress, redirect_stdout and ExitStack cover common needs without writing a manager.
Next (see Functional & Dynamic Python): functions as building blocks: pure functions, partial, composition, introspection and registries.

Practise this