Python Context Managers: Beyond

Overview

Everyone learns with open(...) as f on day one and then stops thinking about it. The with statement is a general mechanism for "do setup, do work, do teardown, even if the work explodes." Once you start writing your own, you find uses for it constantly — database transactions, temporary state changes, timing blocks, and locking.

Here's the actual mechanism, and then the patterns I use.

What with really does

When Python executes:

with expr as target:
    body

It does roughly this:

manager = expr
value = manager.__enter__()
target = value
try:
    body
except:
    if not manager.__exit__(*sys.exc_info()):
        raise
else:
    manager.__exit__(None, None, None)

Two dunder methods. __enter__ runs before the block, __exit__ runs after — always, whether the block raised or not. The return value of __exit__ decides whether an exception propagates: return True to swallow it, anything falsy to let it through.

That last part is the feature almost nobody uses. A context manager can absorb specific exceptions, which is how some libraries implement "retry on failure" or "log and continue."

Class-based: the explicit version

class Timer:
    def __init__(self, label=""):
        self.label = label

    def __enter__(self):
        self.start = time.perf_counter()
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.elapsed = time.perf_counter() - self.start
        print(f"{self.label}: {self.elapsed:.3f}s")
        return False  # do not suppress exceptions

with Timer("db query"):
    results = run_expensive_query()

__exit__ returning False (or nothing) means exceptions propagate normally. That's almost always what you want. Returning True makes the exception disappear, which is a great way to introduce silent failures into a codebase.

contextlib.contextmanager: the shorter version

For simple cases, you don't need a class. Decorate a generator and everything before yield is the setup, everything after is the teardown:

from contextlib import contextmanager

@contextmanager
def timer(label=""):
    start = time.perf_counter()
    try:
        yield
    finally:
        elapsed = time.perf_counter() - start
        print(f"{label}: {elapsed:.3f}s")

Ten lines instead of fifteen, and easier to read.

The try/finally is not optional. If you skip it and the block raises, the teardown never runs. This is the most common bug in hand-rolled context managers.

The generator form can also return a value:

@contextmanager
def open_config(path):
    with open(path) as f:
        cfg = json.load(f)
    yield cfg

with open_config("settings.json") as cfg:
    print(cfg["debug"])

Real use: temporary directory change

@contextmanager
def chdir(path):
    prev = os.getcwd()
    os.chdir(path)
    try:
        yield
    finally:
        os.chdir(prev)

with chdir("/tmp/build"):
    subprocess.run(["make"])

This one saves real code. Any time you need to change global state temporarily — working directory, environment variables, locale, logging level — a context manager guarantees the restoration even if the code inside throws. Without it, you eventually ship a bug where the process runs the rest of its life in the wrong directory.

Real use: environment variables

@contextmanager
def env(**overrides):
    original = {k: os.environ.get(k) for k in overrides}
    os.environ.update(overrides)
    try:
        yield
    finally:
        for k, v in original.items():
            if v is None:
                os.environ.pop(k, None)
            else:
                os.environ[k] = v

with env(DATABASE_URL="sqlite://:memory:", DEBUG="1"):
    run_tests()

I use this in test suites constantly. Tests that mutate os.environ directly are a nightmare to debug when they leak into each other.

Suppressing exceptions

The contextlib.suppress helper is built in and handles the "I expect this to fail and don't care" case:

from contextlib import suppress

with suppress(FileNotFoundError):
    os.remove("tempfile.txt")

with suppress(KeyError, IndexError):
    value = data["config"]["missing"]

Use this sparingly. Catching an exception you didn't actually anticipate hides real bugs. But for genuinely expected failures — a temp file that might not exist — it's cleaner than a two-line try/except.

Reusable context managers

Two things that aren't obvious:

A class-based context manager can be reused. You can call it multiple times on the same instance because __enter__ resets state:

t = Timer("phase")

with t:
    phase_one()

with t:
    phase_two()

A generator-based one cannot. Once the generator is exhausted, calling it again raises RuntimeError: generator didn't yield. Each use needs a fresh call. This bites people who define one with @contextmanager and then try to reuse the variable.

# Broken
cm = timer("x")
with cm: pass
with cm: pass   # RuntimeError

# Fine
with timer("x"): pass
with timer("x"): pass

Nested and multiple contexts

Multiple context managers on one line run in order, and exit in reverse:

with open("input.txt") as src, open("output.txt", "w") as dst:
    dst.write(src.read())

Both __exit__ methods run, in reverse order, even if the body raises. That's why this is safe.

Async context managers

Async versions use __aenter__ and __aexit__, and the keyword is async with:

@asynccontextmanager
async def session_scope():
    session = AsyncSession()
    try:
        yield session
        await session.commit()
    except:
        await session.rollback()
        raise
    finally:
        await session.close()

async with session_scope() as s:
    s.add(User(name="alice"))

This is the standard pattern for async database sessions. The commit-on-success / rollback-on-failure behavior is exactly what __aexit__ is designed to express, and doing it manually in every function is worse in every way.

ExitStack: dynamic context management

Sometimes you don't know how many context managers you need at compile time. ExitStack handles it:

from contextlib import ExitStack

with ExitStack() as stack:
    files = [stack.enter_context(open(p)) for p in paths]
    # all files close when the block exits, even if this list is empty
    for f in files:
        process(f)

Less common, but when you need it, you really need it. The alternative is a hand-managed list of open files and a try/finally, which is exactly what ExitStack is abstracting.

When not to use one

Context managers are for setup and teardown that must happen as a pair. If you just need cleanup at the end of a function, try/finally is simpler. If you need to conditionally clean up based on success, a plain function with an explicit close() is often clearer than a context manager — the with syntax hides the operation, which is a bug when the operation matters.

But for anything resource-like — files, connections, locks, temp state — the with statement is the right tool, and writing your own is a five-minute investment that pays off for the rest of the codebase's life.