contextlib: @contextmanager and Helpers
22 min · Intermediate
Need help? Ask a tutor on WhatsApp (opens in a new tab)A class with two methods is a lot of code for a simple setup and cleanup. The contextlib module lets you write the same thing as one short generator function.
@contextmanager turns a generator into a manager
contextlib.contextmanager is a decorator that builds a context manager from a generator function with one yield: code before yield is setup, code after it is cleanup.
Write a generator with a single yield. The code before yield is the setup, and the code after it is the cleanup. The @contextmanager decorator does the rest.
from contextlib import contextmanager
@contextmanagerdef section(name): print("start", name) try: yield finally: print("end", name)
with section("report"): print("body")Example 1 · A section manager
from contextlib import contextmanager
@contextmanagerdef section(name): print("start", name) try: yield finally: print("end", name)
with section("report"): print("body")
try: with section("export"): raise ValueError("disk full")except ValueError as err: print("caught:", err)start report body end report start export end export caught: disk full
What this code does
- The
finallyline printsend exportbefore the error reaches theexcept. - One
yieldsplits the function into setup and cleanup.
The with block runs while the generator is paused at yield. When the block ends, the generator is resumed and the cleanup runs.
Yield a value and protect the cleanup
The value you yield becomes the as name. If the block raises, the error is thrown into the generator at the yield line. Put the cleanup in finally, or it is skipped on errors.
from contextlib import contextmanager
@contextmanagerdef override(settings, key, value): old = settings[key] settings[key] = value try: yield settings finally: settings[key] = oldExample 2 · Temporarily change a setting
from contextlib import contextmanager
@contextmanagerdef override(settings, key, value): old = settings[key] settings[key] = value try: yield settings finally: settings[key] = old
config = {"mode": "live"}with override(config, "mode", "test") as c: print(c["mode"])print(config["mode"])
try: with override(config, "mode", "debug"): raise RuntimeErrorexcept RuntimeError: passprint(config["mode"])test live live
What this code does
- The yielded dictionary is the
asname. - Even after an error, the setting is restored to
live.
Without try/finally, an error in the block would leave the setting changed for good.
Ready-made helpers
contextlib already has managers for common jobs. suppress ignores chosen errors. redirect_stdout sends printed text elsewhere. nullcontext does nothing, which helps when a manager is optional.
import iofrom contextlib import redirect_stdout, suppress
with suppress(FileNotFoundError): open("missing.txt")
buffer = io.StringIO()with redirect_stdout(buffer): print("hidden")print(buffer.getvalue().strip())Example 3 · suppress and redirect_stdout
import iofrom contextlib import redirect_stdout, suppress
with suppress(KeyError): {}["missing"]print("no error")
buffer = io.StringIO()with redirect_stdout(buffer): print("this is captured")text = buffer.getvalue().strip()print(text.upper())no error THIS IS CAPTURED
What this code does
suppressis shorter and clearer than try/except/pass.redirect_stdoutis handy for testing code that prints.
suppress(FileNotFoundError) replaces a four-line try/except that only says pass.
ExitStack for a changing number of managers
ExitStack is a contextlib object that holds any number of cleanup steps and managers, and runs them in reverse order when its own with block ends.
A with line has a fixed number of managers. ExitStack collects any number at run time and closes them in reverse order when the block ends.
from contextlib import ExitStack
with ExitStack() as stack: for name in ("a", "b", "c"): stack.callback(print, "closing", name) print("working")Example 4 · Cleanup in reverse order
from contextlib import ExitStack
with ExitStack() as stack: for name in ("a", "b", "c"): stack.callback(print, "closing", name) print("working")working closing c closing b closing a
What this code does
- Callbacks run in the reverse order of registration: c, b, a.
- Each
callbackcall registers one cleanup step.
Example 5 · Open any number of files
import osimport tempfilefrom contextlib import ExitStack
with tempfile.TemporaryDirectory() as folder: names = ["a.txt", "b.txt", "c.txt"] with ExitStack() as stack: files = [stack.enter_context(open(os.path.join(folder, n), "w")) for n in names] for number, f in enumerate(files, start=1): f.write(str(number)) print([f.closed for f in files])[True, True, True]
What this code does
enter_contextenters each manager and remembers it for cleanup.- After the block, all three files are closed.
It also accepts managers with stack.enter_context(...), for example to open a list of files.
Next step
Check your understanding
12 quick questions