PYTHON / ADVANCED PYTHON
Writing context managers with contextlib
Build generator-based context managers with @contextmanager, handle exceptions in the with-body correctly, and compose them using ExitStack and suppress.
What you will learn
- Split setup and teardown around a single yield in a @contextmanager generator
- Guarantee cleanup with try/finally so a raising with-body still unwinds
- Suppress an exception by catching it around the yield, not by returning True
- Enter a variable number of managers at runtime with contextlib.ExitStack
Understanding Writing context managers with contextlib
A context manager is just an object with __enter__ and __exit__. The @contextmanager decorator lets you write both halves as one generator: everything before the yield is __enter__, the value you yield becomes the target of `as`, and everything after the yield is __exit__. The decorator wraps your generator function so that calling it produces a fresh single-use manager object rather than running any of the body.
The important detail is how the two halves are resumed. On a clean exit, __exit__ calls next() on your generator, which continues after the yield. If the with-body raised, __exit__ instead calls gen.throw(), so the exception appears to originate at the yield statement itself. That is why cleanup written as a plain statement after yield silently disappears on failure, and why try/finally around the yield is the only reliable way to release a resource.
Because the exception is delivered into your generator, you control its fate with ordinary except clauses. Catch it and fall through, and the exception is suppressed for the caller; re-raise or let it escape, and it propagates. contextlib ships ready-made pieces built on the same protocol: suppress(*exceptions) discards matching errors, closing(obj) calls obj.close(), and ExitStack collects managers you decide on at runtime and unwinds them in reverse order.
from contextlib import contextmanager
contextmanager
def tag(name):
print(f"<{name}>")
try:
yield name.upper()
finally:
print(f"</{name}>")
with tag("p") as upper:
print(f" hello from {upper}")
contextmanager
def swallow(exc_type):
try:
yield
except exc_type as e:
print(f"caught {type(e).__name__}: {e}")
with swallow(ValueError):
raise ValueError("bad input")
print("still running")A @contextmanager generator's yield is the seam where control leaves for the with-body and returns either by next() on success or gen.throw() on failure.
Worked examples
ExitStack for a dynamic number of managers
Enters an unknown-at-write-time list of context managers and unwinds them in reverse order.
from contextlib import ExitStack, contextmanager
contextmanager
def resource(n):
print(f"open {n}")
try:
yield n
finally:
print(f"close {n}")
with ExitStack() as stack:
handles = [stack.enter_context(resource(i)) for i in range(3)]
print("using", handles)Example explained
Line 1stack.enter_context(cm) calls cm.__enter__() immediately and returns its value, so handles holds 0, 1, 2.
Line 2The stack records each manager it entered, which is what makes a runtime-sized `with` possible.
Line 3Leaving the ExitStack block calls the recorded __exit__ methods last-in-first-out, giving close 2, 1, 0.
Line 4If resource(1) had raised during setup, resource(0) would still be closed, because it was already on the stack.
Reusing a manager as a decorator
Shows that @contextmanager results are also decorators, and how suppress differs from a full try/except.
from contextlib import contextmanager, suppress
contextmanager
def announce(label):
print(f"start {label}")
try:
yield
finally:
print(f"end {label}")
announce("task")
def work():
print("working")
work()
work()
with suppress(ZeroDivisionError):
print(1 / 0)
print("survived")Example explained
Line 1The object returned by announce("task") inherits ContextDecorator, so it can wrap a function instead of a block.
Line 2It builds a fresh generator on every call, which is why the second work() prints the full start/end pair again.
Line 31 / 0 raises before print receives an argument, so nothing is printed from that line at all.
Line 4suppress ends the with-block at the point of the error and continues after it; it does not resume the remaining body.
Important notes
Returning True from a @contextmanager generator does not suppress anything; suppression happens only when an except clause around the yield catches the exception and does not re-raise.
A generator may yield exactly once per use; a second yield produces RuntimeError: generator didn't stop when the block exits.
Common mistakes
Writing teardown as a bare statement after yield with no try/finally: when the body raises, gen.throw() surfaces the error at the yield and the teardown line never executes, leaking the resource.
Storing the manager in a variable and using it in two with statements: the generator is already exhausted, so the second __enter__ raises RuntimeError: generator didn't yield.
Wrapping the yield in try/except Exception without re-raising: every error in the with-body is silently swallowed and the program continues with half-finished work.
Try it yourself
Change, predict, then run
Write a @contextmanager named override(d, key, value) that temporarily sets d[key] to value and restores the original on exit. Prove it works by raising inside the with-block and printing the dict afterwards.
Open the Python workspaceCheck your understanding
A @contextmanager generator prints its cleanup message on the line right after yield, with no try/finally. The with-body raises ValueError. What happens?
- The cleanup line is skipped and the ValueError propagates, because the exception is raised at the yield point inside the generator
- The cleanup line runs and then the ValueError propagates, because contextlib always resumes the generator on exit
- The cleanup line runs and the ValueError is suppressed, since leaving the generator normally signals a handled error
- The generator raises StopIteration, which replaces the original ValueError
Show answer
On failure __exit__ calls gen.throw(), so the ValueError appears at the yield statement; with no handler it immediately propagates out of the generator and the statements after yield are never reached. Option 2 is tempting because the success path really does resume the generator, but that resumption uses next() and only happens when no exception occurred; only a finally block runs in both cases.