Backend / Python core / Stdlib / 04_contextlib.md

Context managers: the with protocol and contextlib

Updated 5 interview angles 7 min read source
On this page10
  1. The mechanism: what with expands to
  2. What __exit__ decides
  3. Generator-based managers: @contextmanager
  4. ExitStack: a dynamic number of resources
  5. The rest of contextlib, briefly
  6. Async context managers
  7. Where it breaks in production
  8. The trade not taken
  9. Related
  10. Interview angle

Context managers: the with protocol and contextlib

Article.

A context manager pairs setup with guaranteed cleanup: with open(p) as f: closes the file however the block ends. The question interviewers are really asking: what exactly does with guarantee, and where do those guarantees stop?

The cleanup runs on every exit from the block, but not if setup itself fails; an exception can be suppressed, replaced or swallowed by accident; and async cleanup can be cancelled halfway. Outputs below are from Python 3.14.5.

Decorators, which the old version of this topic was paired with, are in Decorators.

The mechanism: what with expands to

python 
with EXPR as VAR:
    BLOCK

# behaves like (PEP 343):
mgr = EXPR
exit = type(mgr).__exit__          # looked up on the type, before entering
value = type(mgr).__enter__(mgr)
try:
    VAR = value
    BLOCK
except BaseException as e:
    if not exit(mgr, type(e), e, e.__traceback__):
        raise                       # a falsy return lets it propagate
else:
    exit(mgr, None, None, None)

On 3.14 the bytecode is LOAD_SPECIAL __exit__, LOAD_SPECIAL __enter__, a CALL, the body, and an exception handler (WITH_EXCEPT_START) that calls __exit__ and re-raises unless it returned something truthy.

An object without the methods fails before anything runs: with 1: raises TypeError: 'int' object does not support the context manager protocol (missed __exit__ method); before 3.11 it was an AttributeError.

What __exit__ decides

__exit__(exc_type, exc, tb) gets the exception, or three Nones, and its return value is a decision:

  • Falsy (including no return): the exception propagates. A class that logs in __exit__ and returns None has handled nothing; the exception continues after the cleanup.
  • Truthy: the exception is suppressed. Execution continues after the with statement. Return True only for exceptions you mean to swallow, as contextlib.suppress does.
  • An exception raised inside __exit__ replaces the original. The original survives only as __context__ in the traceback, so a failing cleanup can hide the error that caused it:
python 
class BadExit:
    def __enter__(self): return self
    def __exit__(self, *exc): raise RuntimeError("cleanup failed")

with BadExit():
    raise ValueError("real error")
# RuntimeError('cleanup failed'), with __context__ = ValueError('real error')

Where the guarantee stops: __enter__

__exit__ runs on every exit from the block, but the block only starts once __enter__ has returned. If __enter__ raises, nothing is cleaned up:

python 
class Pair:
    def __enter__(self):
        self.a = connect("A")
        self.b = connect("B")       # raises: A is never closed
        return self
    def __exit__(self, *exc):
        self.b.close(); self.a.close()

A manager that acquires several resources in __enter__ must clean up the ones it already has before re-raising, or build on ExitStack, which does that for you (below).

Several managers, one statement

python 
with (
    open(src) as fin,
    open(dst, "w") as fout,      # parenthesized form: official since 3.10
):
    ...

They enter left to right and exit right to left, and if the second fails to enter, the first is still exited.

Generator-based managers: @contextmanager

python 
from contextlib import contextmanager

@contextmanager
def transaction(conn):
    tx = conn.begin()
    try:
        yield tx                 # the value bound by `as`
    except BaseException:
        tx.rollback()
        raise                    # re-raise, or the caller's error is swallowed
    else:
        tx.commit()

Under the hood the decorator calls next() on enter, and on an exception gen.throw(exc) into the paused yield (Generators). Three rules follow:

  • Re-raise what you catch. An except around the yield that does not re-raise suppresses the caller’s exception silently, the most common bug in hand-written managers.
  • except BaseException, not except Exception, for rollback: a KeyboardInterrupt or an asyncio CancelledError mid-transaction must roll back too, and both are BaseException subclasses.
  • Exactly one yield. A generator that never yields raises RuntimeError: generator didn't yield; one that yields twice, generator didn't stop.

A @contextmanager object is single-use: entering the same object a second time fails (on 3.14.5 with an internal AttributeError). Call the function again for a new one. Used as a decorator, @transaction() works because ContextDecorator creates a fresh manager per call.

Single-use, reusable, reentrant

The contextlib docs sort managers into three kinds. Most are single-use: a @contextmanager object, a file. Reusable ones can be used again after exiting, but not nested inside themselves: a threading.Lock. Reentrant ones can be nested in their own with: a threading.RLock. Entering a non-reentrant manager twice is a deadlock (a Lock) or an error.

ExitStack: a dynamic number of resources

python 
from contextlib import ExitStack

with ExitStack() as stack:
    files = [stack.enter_context(open(p)) for p in paths]
    ...       # if the 7th open fails, the first 6 are closed

Everything entered is exited in reverse order, including when an enter_context in the middle fails. stack.callback(fn, *args) registers plain cleanup functions, stack.push(exit_like) registers __exit__-style callbacks that may suppress, and since 3.11 enter_context() on a non-manager raises TypeError instead of failing later.

Handing ownership out: pop_all()

python 
class Pipeline:
    def __init__(self, paths):
        with ExitStack() as stack:
            self.files = [stack.enter_context(open(p)) for p in paths]
            self._stack = stack.pop_all()   # success: keep them open
        # failure: the with exits and closes whatever was opened

    def close(self):
        self._stack.close()

pop_all() moves the registered cleanups to a new stack and empties the old one, so a constructor can acquire N resources all-or-nothing and own them beyond the with block. This is the fix for the Pair leak above.

The rest of contextlib, briefly

Helper Use it for
suppress(*excs) ignoring an expected exception (FileNotFoundError on delete)
closing(obj) anything with close() but no __exit__, such as a generator
nullcontext(v) “maybe a manager”: the same with either way
redirect_stdout(buf) capturing output of code that prints
chdir(path) a temporary working directory (3.11; not thread-safe)
ContextDecorator a manager usable as a decorator too

Async context managers

async with calls await type(m).__aenter__(m) and await type(m).__aexit__(...). @asynccontextmanager, AsyncExitStack, aclosing() (3.10, for async generators) and AsyncContextDecorator (3.10) mirror the sync helpers, and asyncio.timeout() and asyncio.TaskGroup are themselves async context managers.

The trap is cancellation. If the task is cancelled while __aexit__ is awaiting its cleanup (closing a connection, flushing a buffer), the cleanup await raises CancelledError and stops halfway: half-closed connections and lost writes.

Keep cleanup short and bounded; shield it only when it is genuinely small and must complete, and accept that a shielded await still delays cancellation (Cancellation, timeouts and error handling in asyncio).

Where it breaks in production

Resources opened outside with

A Session, httpx.Client or file created with a plain assignment and closed at the end of the happy path leaks on every error path, until the connection pool runs dry under load. Run tests with python -X dev (or -W error::ResourceWarning): unclosed files and sockets then warn or fail, and -X tracemalloc=10 adds the line where each leaked object was created.

Cleanup that hides the real error

An __exit__ or finally that raises (a rollback on a dead connection) replaces the exception that triggered it. Log the original before attempting cleanup that can fail, or catch and suppress the cleanup error so the real one propagates.

return inside finally

A return, break or continue that leaves a finally block discards any exception in flight. Since 3.14 the compiler warns: SyntaxWarning: 'return' in a 'finally' block (PEP 765); Ruff flags it as B012 (Exception swallowing in finally).

Swallowed exceptions

The @contextmanager that catches without re-raising, and the __exit__ that returns True unconditionally, both turn failures into silent success. Grep for return True in __exit__ methods and for bare except around a yield.

The trade not taken

A class with __enter__/__exit__ versus @contextmanager: the generator form is shorter and keeps setup and teardown side by side, and is the default.

The class form earns its place when the manager carries state or methods beyond the block (a transaction you can also query), must be reusable or reentrant, or must decide about suppression with full access to the exception. And try/finally remains right for one-off cleanup in a single function: a context manager pays off when the same setup-and-cleanup appears twice.

Interview angle 5

  • “What does with guarantee?” — that __exit__ runs on every exit from the block, normal or exceptional. Not that it runs if __enter__ fails: a manager acquiring two resources in __enter__ leaks the first if the second fails.
    • Follow-up: “When does __exit__ suppress an exception, and what if it raises?” — only when it returns something truthy; None propagates. If it raises, its exception replaces the original, which survives as __context__.
  • “Write a timer context manager.” — @contextmanager with start = perf_counter(), try: yield finally: log(perf_counter() - start).
    • Follow-up: “Make it usable as a decorator too.” — subclass ContextDecorator, or use the @contextmanager one with @timer(), which creates a fresh manager per call.
    • Follow-up: “Make it async, and what if the task is cancelled inside the block?” — @asynccontextmanager; cancellation propagates through the yield and finally still runs, but an await in that finally can itself be cancelled.
  • “How can a @contextmanager suppress an exception by accident?” — an except around the yield that does not re-raise.
  • “Open N files, N known at runtime, all closed even if the seventh open fails.” — ExitStack.enter_context in a loop.
    • Follow-up: “Now your constructor must own them beyond the block.” — stack.pop_all() on success hands the cleanups to a new stack the object keeps; on failure the with closes what was opened.
  • “Class or @contextmanager?” — @contextmanager by default; a class when the manager has state or methods, must be reusable or reentrant, or decides suppression itself.
Previous
functools — higher-order function utilities