3. Context managers

Intermediate · 6 min read

A context manager runs setup code before a block and guaranteed cleanup after it — even if the block raises an error. with open(...) is the one you already know.

3.1 What with does

with open("demo.txt", "w", encoding="utf-8") as f:   # setup: open the file
    f.write("hello")
# cleanup: the file is closed here automatically — even if write() had failed
print(f.closed)                                      # → True

3.2 Writing one with a class

import time

class Timer:
    def __enter__(self):                 # runs at the start of the with-block
        self.start = time.perf_counter()
        return self                      # becomes the "as" variable

    def __exit__(self, exc_type, exc, tb):   # runs at the end, error or not
        self.ms = (time.perf_counter() - self.start) * 1000
        return False                     # False = don't swallow errors

with Timer() as t:
    sum(range(100_000))
print(t.ms < 1000)                       # → True

3.3 The easy way: @contextmanager

from contextlib import contextmanager
import os

@contextmanager
def temp_env(name, value):
    """Set an environment variable only inside the with-block."""
    old = os.environ.get(name)
    os.environ[name] = value             # setup
    try:
        yield                            # the with-block runs here
    finally:                             # cleanup always runs
        if old is None:
            del os.environ[name]
        else:
            os.environ[name] = old

with temp_env("LLM_MODEL", "test-model"):
    print(os.environ["LLM_MODEL"])       # → test-model
print("LLM_MODEL" in os.environ)         # → False

3.4 Handy built-ins

from contextlib import suppress
import os

with suppress(FileNotFoundError):        # ignore one specific, expected error
    os.remove("does-not-exist.txt")
print("carried on")                      # → carried on

Database connections, HTTP clients (with httpx.Client() as client:), locks and temporary directories (tempfile.TemporaryDirectory()) all work the same way.

Why it matters for GenAI

Timing a block of pipeline steps, swapping configuration in tests, and making sure HTTP clients and database connections always close are everyday context-manager jobs.

Practice

• Write a @contextmanager called section(title) that prints --- title --- before the block and --- end --- after it.