Context Managers: Beyond with-open
1 · The lesson
readYou've used with open(...) as f: since file-io-basics. The file closes whether the block succeeds, raises, or returns early. That guarantee is the reason with exists, and open isn't the only object that offers it. Locks, database transactions, timing blocks, redirected stdout, temporary working directories, a changing number of open files: anything with a "set up, do work, clean up" shape can be written as a context manager.
This lesson covers the underlying protocol, the two ways to build one, and the parts of contextlib that turn small try/finally patterns into reusable tools.
1. The Protocol: __enter__ and __exit__
A context manager is any object with two methods:
class CM: def __enter__(self): # set up; return whatever `as` should bind to return self def __exit__(self, exc_type, exc_val, exc_tb): # tear down; runs on success, exception, or return return False
with cm as x: runs x = cm.__enter__(). When the block exits, for any reason, cm.__exit__(exc_type, exc_val, exc_tb) runs. If the block finished without an error, all three arguments are None. On an exception they describe what was raised.
The full desugaring of with cm as x: body is roughly:
x = cm.__enter__() try: body except BaseException: if not cm.__exit__(*sys.exc_info()): raise else: cm.__exit__(None, None, None)
setup added so this can run · defines body, cm, sys
# Lightweight mock for objects whose attributes/methods aren't critical class _AutoMock: def __init__(self, name='mock'): self._name = name def __getattr__(self, k): return _AutoMock(self._name + '.' + k) def __call__(self, *a, **kw): print('-> ' + self._name + '() called') return _AutoMock(self._name + '()') def __repr__(self): return '<mock ' + self._name + '>' def __str__(self): return '<mock ' + self._name + '>' def __bool__(self): return True def __iter__(self): return iter([]) def __len__(self): return 0 def __getitem__(self, k): return _AutoMock(self._name + '[...]') def __setitem__(self, k, v): pass def __enter__(self): return self def __exit__(self, *a): return False async def __aenter__(self): return self async def __aexit__(self, *a): return False def __add__(self, o): return self def __radd__(self, o): return self def __sub__(self, o): return self def __mul__(self, o): return self def __rmul__(self, o): return self def __truediv__(self, o): return self def __eq__(self, o): return isinstance(o, _AutoMock) def __hash__(self): return hash(self._name) def __lt__(self, o): return True def __le__(self, o): return True def __gt__(self, o): return False def __ge__(self, o): return False def __mro_entries__(self, bases): return (object,) body = _AutoMock('body') cm = _AutoMock('cm') sys = _AutoMock('sys')
This long form shows two important things:
__exit__always runs.- If
__exit__returns a truthy value, the exception is suppressed (Python does not raise it).
2. Suppressing Exceptions From __exit__
Returning True from __exit__ tells Python "I handled it. Don't pass it on." This is sometimes useful, but most of the time it is not what you want.
class SwallowKeyError: def __enter__(self): return self def __exit__(self, exc_type, exc_val, exc_tb): return exc_type is KeyError # True only for KeyError; let everything else through d = {"a": 1} with SwallowKeyError(): print(d["missing"]) # KeyError raised... print("survived") # ...and swallowed. We get here. with SwallowKeyError(): 1 / 0 # ZeroDivisionError is NOT swallowed, so it propagates
The default return False (or no return at all) is what you want almost every time. Use return True only when the context manager exists specifically to ignore a known exception type. Even then, prefer contextlib.suppress (section 7) for the common cases.
3. A Class-Based Context Manager: Timer
This prints how long the with body took. It is one of the smallest useful context managers (CMs) you can write.
import time class Timer: def __init__(self, label="elapsed"): self.label = label def __enter__(self): self.start = time.perf_counter() return self # bind via `as t:` if the caller wants the object def __exit__(self, exc_type, exc_val, exc_tb): self.elapsed = time.perf_counter() - self.start print(f"{self.label}: {self.elapsed*1000:.2f} ms") return False # do not suppress anything with Timer("sum 1M squares"): sum(i*i for i in range(1_000_000)) # sum 1M squares: 58.43 ms with Timer("with binding") as t: sum(range(100_000)) print(f"caller can read t.elapsed = {t.elapsed:.4f}s")
Notice __enter__ returns self, so as t binds the timer object. If you don't need as, you can return None, and the with form with Timer("x"): still works.
time.perf_counter() is the right clock for measuring durations: it never goes backwards, and it has high resolution. Don't use time.time() here.
4. @contextmanager: Generator-Based CMs
For most single-use context managers, a class is more code than you need. contextlib.contextmanager lets you write the setup, the yield, and the teardown in one function.
import time from contextlib import contextmanager @contextmanager def timer(label="elapsed"): start = time.perf_counter() try: yield # the body of the `with` block runs here finally: elapsed = time.perf_counter() - start print(f"{label}: {elapsed*1000:.2f} ms") with timer("expensive op"): sum(i*i for i in range(1_000_000))
The shape is fixed: everything before yield is __enter__, and everything after is __exit__. The try/finally is required if you want the cleanup to run when there is an exception too. Without it, your cleanup is skipped when the body raises.
The value passed to yield is what as binds to:
@contextmanager def opened(path, mode="r"): f = open(path, mode) try: yield f # bind the file to `as` finally: f.close() with opened("/tmp/x.txt", "w") as f: f.write("hi")
setup added so this can run · re-runs 1 earlier block on this page first
# Lightweight mock for objects whose attributes/methods aren't critical class _AutoMock: def __init__(self, name='mock'): self._name = name def __getattr__(self, k): return _AutoMock(self._name + '.' + k) def __call__(self, *a, **kw): print('-> ' + self._name + '() called') return _AutoMock(self._name + '()') def __repr__(self): return '<mock ' + self._name + '>' def __str__(self): return '<mock ' + self._name + '>' def __bool__(self): return True def __iter__(self): return iter([]) def __len__(self): return 0 def __getitem__(self, k): return _AutoMock(self._name + '[...]') def __setitem__(self, k, v): pass def __enter__(self): return self def __exit__(self, *a): return False async def __aenter__(self): return self async def __aexit__(self, *a): return False def __add__(self, o): return self def __radd__(self, o): return self def __sub__(self, o): return self def __mul__(self, o): return self def __rmul__(self, o): return self def __truediv__(self, o): return self def __eq__(self, o): return isinstance(o, _AutoMock) def __hash__(self): return hash(self._name) def __lt__(self, o): return True def __le__(self, o): return True def __gt__(self, o): return False def __ge__(self, o): return False def __mro_entries__(self, bases): return (object,) contextmanager = _AutoMock('contextmanager') # Earlier blocks on this page that this one builds on, re-run quietly first. try: import time from contextlib import contextmanager except Exception: pass _pm_earlier = [ 'import time\nfrom contextlib import contextmanager\n\n@contextmanager\ndef timer(label="elapsed"):\n start = time.perf_counter()\n try:\n yield # the body of the `with` block runs here\n finally:\n elapsed = time.perf_counter() - start\n print(f"{label}: {elapsed*1000:.2f} ms")\n\nwith timer("expensive op"):\n sum(i*i for i in range(1_000_000))', ] def _pm_run_earlier(_srcs, _g): import sys as _s, io as _io, warnings as _w for _src in _srcs: _o, _e = _s.stdout, _s.stderr _s.stdout = _s.stderr = _io.StringIO() try: with _w.catch_warnings(): _w.simplefilter('ignore') exec(compile(_src, '<earlier block>', 'exec'), _g) except (Exception, SystemExit): pass finally: _s.stdout, _s.stderr = _o, _e if 'matplotlib.pyplot' in _s.modules: _s.modules['matplotlib.pyplot'].close('all') # their plots belong to their own blocks _pm_run_earlier(_pm_earlier, globals()) del _pm_earlier, _pm_run_earlier
An important rule: use exactly one yield. Two yields raise RuntimeError: generator didn't stop. The decorator advances the generator exactly twice, once to start and once to finish, and the second time it must reach the end of the function.
5. Handling Exceptions Inside @contextmanager
When the with body raises, the exception is re-raised at the yield inside your generator. You can catch it, log it, suppress it, or let it propagate:
@contextmanager def safe_step(label): print(f"start {label}") try: yield except ValueError as e: print(f"swallowed ValueError in {label}: {e}") # not re-raising suppresses the error (like __exit__ returning True) finally: print(f"end {label}") with safe_step("parse"): raise ValueError("bad input") # start parse # swallowed ValueError in parse: bad input # end parse print("survived")
setup added so this can run · re-runs 1 earlier block on this page first
# Lightweight mock for objects whose attributes/methods aren't critical class _AutoMock: def __init__(self, name='mock'): self._name = name def __getattr__(self, k): return _AutoMock(self._name + '.' + k) def __call__(self, *a, **kw): print('-> ' + self._name + '() called') return _AutoMock(self._name + '()') def __repr__(self): return '<mock ' + self._name + '>' def __str__(self): return '<mock ' + self._name + '>' def __bool__(self): return True def __iter__(self): return iter([]) def __len__(self): return 0 def __getitem__(self, k): return _AutoMock(self._name + '[...]') def __setitem__(self, k, v): pass def __enter__(self): return self def __exit__(self, *a): return False async def __aenter__(self): return self async def __aexit__(self, *a): return False def __add__(self, o): return self def __radd__(self, o): return self def __sub__(self, o): return self def __mul__(self, o): return self def __rmul__(self, o): return self def __truediv__(self, o): return self def __eq__(self, o): return isinstance(o, _AutoMock) def __hash__(self): return hash(self._name) def __lt__(self, o): return True def __le__(self, o): return True def __gt__(self, o): return False def __ge__(self, o): return False def __mro_entries__(self, bases): return (object,) contextmanager = _AutoMock('contextmanager') # Earlier blocks on this page that this one builds on, re-run quietly first. try: import time from contextlib import contextmanager except Exception: pass _pm_earlier = [ 'import time\nfrom contextlib import contextmanager\n\n@contextmanager\ndef timer(label="elapsed"):\n start = time.perf_counter()\n try:\n yield # the body of the `with` block runs here\n finally:\n elapsed = time.perf_counter() - start\n print(f"{label}: {elapsed*1000:.2f} ms")\n\nwith timer("expensive op"):\n sum(i*i for i in range(1_000_000))', ] def _pm_run_earlier(_srcs, _g): import sys as _s, io as _io, warnings as _w for _src in _srcs: _o, _e = _s.stdout, _s.stderr _s.stdout = _s.stderr = _io.StringIO() try: with _w.catch_warnings(): _w.simplefilter('ignore') exec(compile(_src, '<earlier block>', 'exec'), _g) except (Exception, SystemExit): pass finally: _s.stdout, _s.stderr = _o, _e if 'matplotlib.pyplot' in _s.modules: _s.modules['matplotlib.pyplot'].close('all') # their plots belong to their own blocks _pm_run_earlier(_pm_earlier, globals()) del _pm_earlier, _pm_run_earlier
If you want to let the exception through but still run cleanup, drop the except and use finally alone. The try/except/finally shape around yield covers every case you need:
@contextmanager def transaction(db): db.begin() try: yield db except Exception: db.rollback() raise # re-raise after rollback else: db.commit()
setup added so this can run · re-runs 1 earlier block on this page first
# Lightweight mock for objects whose attributes/methods aren't critical class _AutoMock: def __init__(self, name='mock'): self._name = name def __getattr__(self, k): return _AutoMock(self._name + '.' + k) def __call__(self, *a, **kw): print('-> ' + self._name + '() called') return _AutoMock(self._name + '()') def __repr__(self): return '<mock ' + self._name + '>' def __str__(self): return '<mock ' + self._name + '>' def __bool__(self): return True def __iter__(self): return iter([]) def __len__(self): return 0 def __getitem__(self, k): return _AutoMock(self._name + '[...]') def __setitem__(self, k, v): pass def __enter__(self): return self def __exit__(self, *a): return False async def __aenter__(self): return self async def __aexit__(self, *a): return False def __add__(self, o): return self def __radd__(self, o): return self def __sub__(self, o): return self def __mul__(self, o): return self def __rmul__(self, o): return self def __truediv__(self, o): return self def __eq__(self, o): return isinstance(o, _AutoMock) def __hash__(self): return hash(self._name) def __lt__(self, o): return True def __le__(self, o): return True def __gt__(self, o): return False def __ge__(self, o): return False def __mro_entries__(self, bases): return (object,) contextmanager = _AutoMock('contextmanager') # Earlier blocks on this page that this one builds on, re-run quietly first. try: import time from contextlib import contextmanager except Exception: pass _pm_earlier = [ 'import time\nfrom contextlib import contextmanager\n\n@contextmanager\ndef timer(label="elapsed"):\n start = time.perf_counter()\n try:\n yield # the body of the `with` block runs here\n finally:\n elapsed = time.perf_counter() - start\n print(f"{label}: {elapsed*1000:.2f} ms")\n\nwith timer("expensive op"):\n sum(i*i for i in range(1_000_000))', ] def _pm_run_earlier(_srcs, _g): import sys as _s, io as _io, warnings as _w for _src in _srcs: _o, _e = _s.stdout, _s.stderr _s.stdout = _s.stderr = _io.StringIO() try: with _w.catch_warnings(): _w.simplefilter('ignore') exec(compile(_src, '<earlier block>', 'exec'), _g) except (Exception, SystemExit): pass finally: _s.stdout, _s.stderr = _o, _e if 'matplotlib.pyplot' in _s.modules: _s.modules['matplotlib.pyplot'].close('all') # their plots belong to their own blocks _pm_run_earlier(_pm_earlier, globals()) del _pm_earlier, _pm_run_earlier
else (after a try that has an except) runs only when the body completed without exception. This fits the commit-on-success / rollback-on-failure pattern.
6. ExitStack: A Changing Number of Context Managers
Sometimes you don't know, when you write the code, how many resources you will need. Opening N files in one block, or building a configurable pipeline of CMs, doesn't fit the with a, b, c: syntax. ExitStack solves this.
from contextlib import ExitStack def cat_all(paths): """Open every path under one stack; close everything in reverse on exit.""" with ExitStack() as stack: files = [stack.enter_context(open(p)) for p in paths] for f in files: for line in f: print(line, end="") # All files are guaranteed closed when the `with` block ends, even if one raises. cat_all(["a.txt", "b.txt", "c.txt"])
stack.enter_context(cm) runs cm.__enter__() and registers the matching __exit__ to run when the stack is closed, in LIFO order (last opened, first closed).
You can also register a callback that isn't itself a CM:
with ExitStack() as stack: conn = open_connection() stack.callback(conn.close) # plain cleanup function stack.callback(print, "all done") # any callable, plus its arguments do_work(conn)
setup added so this can run · defines ExitStack, open_connection, do_work
# Lightweight mock for objects whose attributes/methods aren't critical class _AutoMock: def __init__(self, name='mock'): self._name = name def __getattr__(self, k): return _AutoMock(self._name + '.' + k) def __call__(self, *a, **kw): print('-> ' + self._name + '() called') return _AutoMock(self._name + '()') def __repr__(self): return '<mock ' + self._name + '>' def __str__(self): return '<mock ' + self._name + '>' def __bool__(self): return True def __iter__(self): return iter([]) def __len__(self): return 0 def __getitem__(self, k): return _AutoMock(self._name + '[...]') def __setitem__(self, k, v): pass def __enter__(self): return self def __exit__(self, *a): return False async def __aenter__(self): return self async def __aexit__(self, *a): return False def __add__(self, o): return self def __radd__(self, o): return self def __sub__(self, o): return self def __mul__(self, o): return self def __rmul__(self, o): return self def __truediv__(self, o): return self def __eq__(self, o): return isinstance(o, _AutoMock) def __hash__(self): return hash(self._name) def __lt__(self, o): return True def __le__(self, o): return True def __gt__(self, o): return False def __ge__(self, o): return False def __mro_entries__(self, bases): return (object,) def ExitStack(*_a, **_kw): print('-> ExitStack() called') return _AutoMock('ExitStack()') def open_connection(*_a, **_kw): print('-> open_connection() called') return _AutoMock('open_connection()') def do_work(*_a, **_kw): print('-> do_work() called') return _AutoMock('do_work()')
You can also use ExitStack to hold a resource only in some cases: put it on the stack only when you need it.
7. contextlib.suppress: Recap
Covered in exceptions for the "ignore exactly this error" pattern:
from contextlib import suppress import os with suppress(FileNotFoundError): os.remove("scratch.tmp") # does nothing if the file doesn't exist
This one line replaces try/except SpecificError: pass. Don't use suppress(Exception), because it is the same anti-pattern as except Exception: pass.
8. contextlib.closing: Wrap a .close()-Only Object
Some objects have .close() but no __enter__/__exit__. closing creates a CM around them:
from contextlib import closing from urllib.request import urlopen with closing(urlopen("https://example.com")) as page: html = page.read() # page.close() is called automatically
Modern stdlib objects (sockets, files, urlopen on 3.x) already implement the protocol directly, so closing mostly shows up for third-party libraries or DB cursors that predate the protocol.
9. contextlib.nullcontext: Conditional CMs
Sometimes you want the with form whether or not there's a real resource to hold. nullcontext(x) is a no-op CM whose __enter__ returns x.
from contextlib import nullcontext def process(path, verbose=False): log_cm = open("log.txt", "a") if verbose else nullcontext() with log_cm as log: # `log` is a file object or None; one code path covers both do_work(log)
setup added so this can run · defines do_work
# Lightweight mock for objects whose attributes/methods aren't critical class _AutoMock: def __init__(self, name='mock'): self._name = name def __getattr__(self, k): return _AutoMock(self._name + '.' + k) def __call__(self, *a, **kw): print('-> ' + self._name + '() called') return _AutoMock(self._name + '()') def __repr__(self): return '<mock ' + self._name + '>' def __str__(self): return '<mock ' + self._name + '>' def __bool__(self): return True def __iter__(self): return iter([]) def __len__(self): return 0 def __getitem__(self, k): return _AutoMock(self._name + '[...]') def __setitem__(self, k, v): pass def __enter__(self): return self def __exit__(self, *a): return False async def __aenter__(self): return self async def __aexit__(self, *a): return False def __add__(self, o): return self def __radd__(self, o): return self def __sub__(self, o): return self def __mul__(self, o): return self def __rmul__(self, o): return self def __truediv__(self, o): return self def __eq__(self, o): return isinstance(o, _AutoMock) def __hash__(self): return hash(self._name) def __lt__(self, o): return True def __le__(self, o): return True def __gt__(self, o): return False def __ge__(self, o): return False def __mro_entries__(self, bases): return (object,) def do_work(*_a, **_kw): print('-> do_work() called') return _AutoMock('do_work()')
This replaces the repeated code in if cm: with cm: ... else: .... It is also useful for a lock that you only need with concurrency, or a profiler that is sometimes disabled.
10. redirect_stdout and redirect_stderr: Capturing Prints
This sends print output into a file (or io.StringIO) while a block runs:
import io from contextlib import redirect_stdout buf = io.StringIO() with redirect_stdout(buf): print("captured") help(int) # prints a lot of text, but into buf, not the terminal captured = buf.getvalue() print(f"caught {len(captured)} chars")
This is useful for testing CLI output, for capturing output from libraries that print a lot and don't accept a file argument, and for showing REPL-style help inside a TUI. redirect_stderr does the same for sys.stderr.
11. Async Context Managers: async with
Async resources (HTTP sessions, async DB connections, async locks) implement __aenter__ and __aexit__ and are entered with async with:
import aiohttp async def fetch(url): async with aiohttp.ClientSession() as session: async with session.get(url) as resp: return await resp.text()
setup added so this can run · defines session, resp
# Lightweight mock for objects whose attributes/methods aren't critical class _AutoMock: def __init__(self, name='mock'): self._name = name def __getattr__(self, k): return _AutoMock(self._name + '.' + k) def __call__(self, *a, **kw): print('-> ' + self._name + '() called') return _AutoMock(self._name + '()') def __repr__(self): return '<mock ' + self._name + '>' def __str__(self): return '<mock ' + self._name + '>' def __bool__(self): return True def __iter__(self): return iter([]) def __len__(self): return 0 def __getitem__(self, k): return _AutoMock(self._name + '[...]') def __setitem__(self, k, v): pass def __enter__(self): return self def __exit__(self, *a): return False async def __aenter__(self): return self async def __aexit__(self, *a): return False def __add__(self, o): return self def __radd__(self, o): return self def __sub__(self, o): return self def __mul__(self, o): return self def __rmul__(self, o): return self def __truediv__(self, o): return self def __eq__(self, o): return isinstance(o, _AutoMock) def __hash__(self): return hash(self._name) def __lt__(self, o): return True def __le__(self, o): return True def __gt__(self, o): return False def __ge__(self, o): return False def __mro_entries__(self, bases): return (object,) session = _AutoMock('session') resp = _AutoMock('resp')
The rules are the same as in the sync protocol (the same arguments to __aexit__, the same suppression rule), except that the methods are coroutines and you await them. The @contextmanager equivalent is @asynccontextmanager from contextlib. Full coverage in async.
12. Real-World Examples
Three patterns you will see often:
# Database transactions: commit on success, roll back on an exception with db.transaction(): db.execute("INSERT INTO users (...) VALUES (...)") db.execute("INSERT INTO audit (...) VALUES (...)") # Either both inserts happened, or neither. # File locks: the lock is always released on exit with lockfile.lock("/tmp/job.lock"): do_exclusive_work() # pytest's `raises`: check that the block raised the expected exception import pytest with pytest.raises(ValueError): int("not-a-number")
setup added so this can run · defines do_exclusive_work, db, lockfile
# Lightweight mock for objects whose attributes/methods aren't critical class _AutoMock: def __init__(self, name='mock'): self._name = name def __getattr__(self, k): return _AutoMock(self._name + '.' + k) def __call__(self, *a, **kw): print('-> ' + self._name + '() called') return _AutoMock(self._name + '()') def __repr__(self): return '<mock ' + self._name + '>' def __str__(self): return '<mock ' + self._name + '>' def __bool__(self): return True def __iter__(self): return iter([]) def __len__(self): return 0 def __getitem__(self, k): return _AutoMock(self._name + '[...]') def __setitem__(self, k, v): pass def __enter__(self): return self def __exit__(self, *a): return False async def __aenter__(self): return self async def __aexit__(self, *a): return False def __add__(self, o): return self def __radd__(self, o): return self def __sub__(self, o): return self def __mul__(self, o): return self def __rmul__(self, o): return self def __truediv__(self, o): return self def __eq__(self, o): return isinstance(o, _AutoMock) def __hash__(self): return hash(self._name) def __lt__(self, o): return True def __le__(self, o): return True def __gt__(self, o): return False def __ge__(self, o): return False def __mro_entries__(self, bases): return (object,) def do_exclusive_work(*_a, **_kw): print('-> do_exclusive_work() called') return _AutoMock('do_exclusive_work()') db = _AutoMock('db') lockfile = _AutoMock('lockfile')
The transaction and lock patterns are the original motivation for with: pairs of operations where forgetting the second one corrupts state. A context manager runs the second operation for you, so you can't forget it.
Common Mistakes
1. Returning True from __exit__ by accident
def __exit__(self, exc_type, exc_val, exc_tb): self.cleanup() return self.cleanup() # if cleanup returns a truthy value, all exceptions are swallowed
Don't return anything from __exit__ unless you really want to suppress the exception. Explicit return False or no return at all is the safe default.
2. Forgetting that yield's value is what as binds
@contextmanager def db_conn(): conn = connect() try: yield # yields nothing, so the caller's `as c` gets None! finally: conn.close() with db_conn() as c: c.execute(...) # AttributeError: NoneType has no 'execute'
setup added so this can run · defines contextmanager, connect
# Lightweight mock for objects whose attributes/methods aren't critical class _AutoMock: def __init__(self, name='mock'): self._name = name def __getattr__(self, k): return _AutoMock(self._name + '.' + k) def __call__(self, *a, **kw): print('-> ' + self._name + '() called') return _AutoMock(self._name + '()') def __repr__(self): return '<mock ' + self._name + '>' def __str__(self): return '<mock ' + self._name + '>' def __bool__(self): return True def __iter__(self): return iter([]) def __len__(self): return 0 def __getitem__(self, k): return _AutoMock(self._name + '[...]') def __setitem__(self, k, v): pass def __enter__(self): return self def __exit__(self, *a): return False async def __aenter__(self): return self async def __aexit__(self, *a): return False def __add__(self, o): return self def __radd__(self, o): return self def __sub__(self, o): return self def __mul__(self, o): return self def __rmul__(self, o): return self def __truediv__(self, o): return self def __eq__(self, o): return isinstance(o, _AutoMock) def __hash__(self): return hash(self._name) def __lt__(self, o): return True def __le__(self, o): return True def __gt__(self, o): return False def __ge__(self, o): return False def __mro_entries__(self, bases): return (object,) contextmanager = _AutoMock('contextmanager') def connect(*_a, **_kw): print('-> connect() called') return _AutoMock('connect()')
yield conn if you want the caller to use it.
3. Two yields in a @contextmanager generator
@contextmanager def broken(): yield "first" yield "second" # RuntimeError: generator didn't stop
setup added so this can run · re-runs 1 earlier block on this page first
# Lightweight mock for objects whose attributes/methods aren't critical class _AutoMock: def __init__(self, name='mock'): self._name = name def __getattr__(self, k): return _AutoMock(self._name + '.' + k) def __call__(self, *a, **kw): print('-> ' + self._name + '() called') return _AutoMock(self._name + '()') def __repr__(self): return '<mock ' + self._name + '>' def __str__(self): return '<mock ' + self._name + '>' def __bool__(self): return True def __iter__(self): return iter([]) def __len__(self): return 0 def __getitem__(self, k): return _AutoMock(self._name + '[...]') def __setitem__(self, k, v): pass def __enter__(self): return self def __exit__(self, *a): return False async def __aenter__(self): return self async def __aexit__(self, *a): return False def __add__(self, o): return self def __radd__(self, o): return self def __sub__(self, o): return self def __mul__(self, o): return self def __rmul__(self, o): return self def __truediv__(self, o): return self def __eq__(self, o): return isinstance(o, _AutoMock) def __hash__(self): return hash(self._name) def __lt__(self, o): return True def __le__(self, o): return True def __gt__(self, o): return False def __ge__(self, o): return False def __mro_entries__(self, bases): return (object,) contextmanager = _AutoMock('contextmanager') # Earlier blocks on this page that this one builds on, re-run quietly first. try: import time from contextlib import contextmanager except Exception: pass _pm_earlier = [ 'import time\nfrom contextlib import contextmanager\n\n@contextmanager\ndef timer(label="elapsed"):\n start = time.perf_counter()\n try:\n yield # the body of the `with` block runs here\n finally:\n elapsed = time.perf_counter() - start\n print(f"{label}: {elapsed*1000:.2f} ms")\n\nwith timer("expensive op"):\n sum(i*i for i in range(1_000_000))', ] def _pm_run_earlier(_srcs, _g): import sys as _s, io as _io, warnings as _w for _src in _srcs: _o, _e = _s.stdout, _s.stderr _s.stdout = _s.stderr = _io.StringIO() try: with _w.catch_warnings(): _w.simplefilter('ignore') exec(compile(_src, '<earlier block>', 'exec'), _g) except (Exception, SystemExit): pass finally: _s.stdout, _s.stderr = _o, _e if 'matplotlib.pyplot' in _s.modules: _s.modules['matplotlib.pyplot'].close('all') # their plots belong to their own blocks _pm_run_earlier(_pm_earlier, globals()) del _pm_earlier, _pm_run_earlier
The decorator drives the generator exactly twice. Anything other than "one yield, then end" is a bug.
4. Not handling exceptions inside __exit__
If __exit__ itself raises, the new exception replaces the original, and the original is reported only as its __context__. Wrap teardown that could fail (closing a flaky socket) in its own try/except:
def __exit__(self, exc_type, exc_val, exc_tb): try: self.conn.close() except Exception: logger.exception("close failed; preserving original error") return False # always let the original through
setup added so this can run · defines logger
# Lightweight mock for objects whose attributes/methods aren't critical class _AutoMock: def __init__(self, name='mock'): self._name = name def __getattr__(self, k): return _AutoMock(self._name + '.' + k) def __call__(self, *a, **kw): print('-> ' + self._name + '() called') return _AutoMock(self._name + '()') def __repr__(self): return '<mock ' + self._name + '>' def __str__(self): return '<mock ' + self._name + '>' def __bool__(self): return True def __iter__(self): return iter([]) def __len__(self): return 0 def __getitem__(self, k): return _AutoMock(self._name + '[...]') def __setitem__(self, k, v): pass def __enter__(self): return self def __exit__(self, *a): return False async def __aenter__(self): return self async def __aexit__(self, *a): return False def __add__(self, o): return self def __radd__(self, o): return self def __sub__(self, o): return self def __mul__(self, o): return self def __rmul__(self, o): return self def __truediv__(self, o): return self def __eq__(self, o): return isinstance(o, _AutoMock) def __hash__(self): return hash(self._name) def __lt__(self, o): return True def __le__(self, o): return True def __gt__(self, o): return False def __ge__(self, o): return False def __mro_entries__(self, bases): return (object,) logger = _AutoMock('logger')
5. Reusing a generator-based CM
cm = timer("once") with cm: pass # works with cm: pass # AttributeError: '_GeneratorContextManager' object has no attribute 'args'
setup added so this can run · re-runs 1 earlier block on this page first
# Lightweight mock for objects whose attributes/methods aren't critical class _AutoMock: def __init__(self, name='mock'): self._name = name def __getattr__(self, k): return _AutoMock(self._name + '.' + k) def __call__(self, *a, **kw): print('-> ' + self._name + '() called') return _AutoMock(self._name + '()') def __repr__(self): return '<mock ' + self._name + '>' def __str__(self): return '<mock ' + self._name + '>' def __bool__(self): return True def __iter__(self): return iter([]) def __len__(self): return 0 def __getitem__(self, k): return _AutoMock(self._name + '[...]') def __setitem__(self, k, v): pass def __enter__(self): return self def __exit__(self, *a): return False async def __aenter__(self): return self async def __aexit__(self, *a): return False def __add__(self, o): return self def __radd__(self, o): return self def __sub__(self, o): return self def __mul__(self, o): return self def __rmul__(self, o): return self def __truediv__(self, o): return self def __eq__(self, o): return isinstance(o, _AutoMock) def __hash__(self): return hash(self._name) def __lt__(self, o): return True def __le__(self, o): return True def __gt__(self, o): return False def __ge__(self, o): return False def __mro_entries__(self, bases): return (object,) def timer(*_a, **_kw): print('-> timer() called') return _AutoMock('timer()') # Earlier blocks on this page that this one builds on, re-run quietly first. try: import time from contextlib import contextmanager except Exception: pass _pm_earlier = [ 'import time\nfrom contextlib import contextmanager\n\n@contextmanager\ndef timer(label="elapsed"):\n start = time.perf_counter()\n try:\n yield # the body of the `with` block runs here\n finally:\n elapsed = time.perf_counter() - start\n print(f"{label}: {elapsed*1000:.2f} ms")\n\nwith timer("expensive op"):\n sum(i*i for i in range(1_000_000))', ] def _pm_run_earlier(_srcs, _g): import sys as _s, io as _io, warnings as _w for _src in _srcs: _o, _e = _s.stdout, _s.stderr _s.stdout = _s.stderr = _io.StringIO() try: with _w.catch_warnings(): _w.simplefilter('ignore') exec(compile(_src, '<earlier block>', 'exec'), _g) except (Exception, SystemExit): pass finally: _s.stdout, _s.stderr = _o, _e if 'matplotlib.pyplot' in _s.modules: _s.modules['matplotlib.pyplot'].close('all') # their plots belong to their own blocks _pm_run_earlier(_pm_earlier, globals()) del _pm_earlier, _pm_run_earlier
@contextmanager produces a single-use CM. Call the function again each time: with timer("twice"): .... Class-based CMs can be reusable if you write them that way, but the generator form can't.
🎯 Your Turn: Build chdir(path)
Write a context manager that changes the current working directory on entry and restores it on exit, even if the block raises.
Requirements:
- Capture
os.getcwd()on entry. - Call
os.chdir(path). - On exit (success or exception),
os.chdirback to the captured directory. - Do not suppress exceptions.
- Implement it twice: once as a class, once with
@contextmanager.
import os from contextlib import contextmanager class chdir: def __init__(self, path): # TODO 1 ... def __enter__(self): # TODO 2: save the cwd, change to self.path, return self (or the new cwd) ... def __exit__(self, exc_type, exc_val, exc_tb): # TODO 3: restore the saved cwd; do not suppress ... @contextmanager def chdir_gen(path): # TODO 4: save cwd, chdir, yield, restore in finally ... with chdir("/tmp"): print(os.getcwd()) # /tmp print(os.getcwd()) # back to the original directory with chdir_gen("/tmp"): raise RuntimeError("boom") # the cwd is still restored, and the exception still propagates
Hint 1: Save before you change
Saveos.getcwd() into self._saved (class) or a local (generator) before the os.chdir(path) call. If chdir raises (bad path), you haven't moved, and nothing has been changed.
Hint 2: try/finally for the generator form
The generator must use try: yield; finally: os.chdir(saved). Without finally, an exception in the with body skips the restore.
Show full solution
import os from contextlib import contextmanager class chdir: """Class-based: change CWD on enter, restore on exit.""" def __init__(self, path): self.path = path def __enter__(self): self._saved = os.getcwd() os.chdir(self.path) return self.path # bind the new directory via `as` def __exit__(self, exc_type, exc_val, exc_tb): os.chdir(self._saved) # no return gives None, which is falsy, so any exception propagates @contextmanager def chdir_gen(path): """Generator-based: same behaviour, fewer lines.""" saved = os.getcwd() os.chdir(path) try: yield path finally: os.chdir(saved) # Demo: both forms behave the same way original = os.getcwd() with chdir("/tmp") as p: assert os.getcwd() == "/tmp" print(f"inside class CM: {p}") assert os.getcwd() == original try: with chdir_gen("/tmp"): assert os.getcwd() == "/tmp" raise RuntimeError("boom") except RuntimeError as e: print(f"caught: {e}") assert os.getcwd() == original print("cwd restored after exception")
Both versions guarantee restoration on success, on exception, and on early return. The class version is more explicit and reusable across multiple with blocks. The generator version is half the code and has the same shape as most single-use CMs.
Python 3.11 added this exact context manager to the standard library as contextlib.chdir. Before that, you wrote it yourself, usually as a class like the one above. Now you can from contextlib import chdir instead of writing either version, and you know what it does inside.
What You Learned
- The protocol:
__enter__(self)runs on entry, and__exit__(self, exc_type, exc_val, exc_tb)runs on exit, always. - Returning
Truefrom__exit__suppresses the exception. This is almost never what you want. Usereturn Falseor no return. - Class-based CMs for stateful or reusable resources.
__enter__returns whateverasshould bind to. @contextmanagerturns a generator into a CM: setup beforeyield, teardown after.try/finallyaround theyieldis required for safe cleanup. Exactly oneyield.ExitStackfor a changing number of CMs: open N things in one block, register cleanup callbacks, or hold a resource only when needed.contextlib.suppressfor "ignore this specific error",closingfor.close()-only objects, andnullcontextfor conditional CMs.redirect_stdout/redirect_stderrfor capturing print output.async withuses__aenter__/__aexit__: the same protocol, with async methods.- Real-world uses: transactions, locks,
pytest.raises, timing blocks, directory changes. Use one anywhere a pair of operations must always happen together.
Next: Async: async/await, event loops, and the async versions of context managers and iterators.
Found this useful? Share it with a friend who’s learning Python.