backend / python core / stdlib / 04_contextlib.md

contextlib — context manager helpers

2 min read source

contextlib — context manager helpers

Tools for writing context managers without writing classes, plus utilities that build on them.

@contextmanager — write a context manager as a generator

Instead of writing a class with __enter__/__exit__:

from contextlib import contextmanager

@contextmanager
def open_resource(name):
    print(f"opening {name}")
    resource = acquire(name)
    try:
        yield resource             # caller gets this in `as`
    finally:
        print(f"closing {name}")
        resource.release()

with open_resource("db") as r:
    r.query("...")

Code before yield is __enter__; code after is __exit__. The try/finally ensures cleanup runs even on exception.

To handle exceptions inside the with-block, catch around yield:

@contextmanager
def transactional():
    tx = begin()
    try:
        yield tx
        tx.commit()
    except:
        tx.rollback()
        raise

suppress — ignore specific exceptions

from contextlib import suppress

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

# Equivalent to:
try:
    os.remove("might_not_exist.txt")
except FileNotFoundError:
    pass

Cleaner than try/except/pass. Multiple types: suppress(FileNotFoundError, PermissionError).

closing — call .close() on exit

For objects that have a close() method but aren’t context managers:

from contextlib import closing
from urllib.request import urlopen

with closing(urlopen("https://example.com")) as page:
    data = page.read()
# page.close() called automatically

Modern stdlib types (files, sockets, db cursors) usually support with directly, so this is mostly for legacy or third-party objects.

ExitStack — combine N context managers dynamically

When you need to enter a variable number of context managers:

from contextlib import ExitStack

with ExitStack() as stack:
    files = [stack.enter_context(open(f)) for f in filenames]
    # all files open here; all close on stack exit, in reverse order
    for f in files:
        process(f)

Without ExitStack, nesting with for an unknown number of resources requires recursion or hacks. ExitStack makes it linear.

You can also push generic cleanup callbacks:

with ExitStack() as stack:
    stack.callback(print, "exiting")    # runs at end
    stack.callback(cleanup_temp, "/tmp/x")

Callbacks run in reverse order of registration.

redirect_stdout / redirect_stderr

Capture output of code that prints:

import io
from contextlib import redirect_stdout

buf = io.StringIO()
with redirect_stdout(buf):
    print("hello")
    third_party_thing.run()

captured = buf.getvalue()

Useful in tests, or when integrating with libraries that only print instead of returning.

Async variants (3.7+)

from contextlib import asynccontextmanager, AsyncExitStack

@asynccontextmanager
async def db_session():
    session = await create_session()
    try:
        yield session
    finally:
        await session.close()

async def main():
    async with db_session() as s:
        await s.query(...)

AsyncExitStack mirrors ExitStack for async contexts.

nullcontext — conditional context manager

from contextlib import nullcontext

cm = open("log.txt", "w") if log_file else nullcontext()

with cm as f:
    if f:
        f.write("...")

nullcontext is a no-op with — useful when you want to conditionally use a context manager but want to avoid duplicating the body.

When to write a class vs use @contextmanager

  • Class — when state matters across enter/exit (e.g. transaction, lock with metadata), when you need to override __init__ semantics, or when the protocol is reused.
  • @contextmanager — for one-shot resource-management functions; when the logic is naturally generator-like (setup, yield, cleanup).

Interview angle

“Implement a context manager that times its block and logs the duration.” Two ways: class with __enter__/__exit__, or generator with @contextmanager. Show both.

“How do you handle 100 dynamic resources with cleanup?” → ExitStack.