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 whole point of with — and open isn't the only object that offers it. Locks, database transactions, timing blocks, redirected stdout, temporary working directories, dynamically opened batches of files: anything with a "set up, do work, tear down" shape is a context manager waiting to be written.
This lesson covers the underlying protocol, the two ways to build one, and the bits of contextlib that turn small try/finally patterns into reusable primitives.
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. On the happy path 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')
Two things drop out of that desugaring that matter:
__exit__always runs.- If
__exit__returns a truthy value, the exception is swallowed.
2. Suppressing Exceptions From __exit__
Returning True from __exit__ tells Python "I handled it; don't propagate." Useful occasionally — most of the time, 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 — NOT swallowed; propagates
The default return False (or simply no return) is what you want 95% of the time. Reach for return True only when the context manager exists specifically to absorb a known exception type — and even then, prefer contextlib.suppress (section 7) for the common cases.
3. A Class-Based Context Manager: Timer
Print how long the with body took. One of the smallest useful CMs you'll 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; the with form with Timer("x"): still works.
time.perf_counter() is the right clock for durations — monotonic, high-resolution. Don't use time.time() here.
4. @contextmanager — Generator-Based CMs
For most one-shot context managers, the class is overkill. 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 # ← `with` body 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__. Everything after is __exit__. The try/finally is mandatory if you want teardown to run on exception too — otherwise 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 · defines contextmanager
# 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')
Critical rule: exactly one yield. Two yields raise RuntimeError: generator didn't stop. The decorator pumps the generator exactly twice — once to start, once to finish — and the second pump must hit 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, swallow 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 = suppression (same as __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 · defines contextmanager
# 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')
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 is the only state machine 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 · defines contextmanager
# 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')
else (after a try that has an except) runs only when the body completed without exception. Perfect for the commit-on-success / rollback-on-failure pattern.
6. ExitStack — Dynamic Number of Context Managers
You sometimes don't know at write-time how many resources you need. Opening N files in one block, or wiring up a configurable pipeline of CMs, doesn't fit with a, b, c: syntax. ExitStack is the answer.
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 unwinds (LIFO order — last opened, first closed).
You can also stash 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 + args/kwargs 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()')
ExitStack is also how you build CMs that conditionally hold a resource — push 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") # no-op if it doesn't exist
One-liner replacement for try/except SpecificError: pass. Don't use suppress(Exception) — same anti-pattern as except Exception: pass.
8. contextlib.closing — Wrap a .close()-Only Object
Some objects have .close() but no __enter__/__exit__. closing synthesises 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() 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()')
Replaces the awkward if cm: with cm: ... else: ... duplication. Same value when threading a lock that's None if no concurrency, or a profiler that's None if disabled.
10. redirect_stdout and redirect_stderr — Capturing Prints
Routing print output into a file (or io.StringIO) for the duration of a block:
import io from contextlib import redirect_stdout buf = io.StringIO() with redirect_stdout(buf): print("captured") help(int) # prints a wall of text — into buf, not the terminal captured = buf.getvalue() print(f"caught {len(captured)} chars")
Handy for testing CLI output, taming chatty libraries that don't accept a file argument, and embedding REPL-style help in a TUI. redirect_stderr is the matching pair 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 mechanics are identical to the sync protocol — same arguments to __aexit__, same suppression rule — except 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'll see constantly:
# Database transactions — commit on success, rollback on exception with db.transaction(): db.execute("INSERT INTO users (...) VALUES (...)") db.execute("INSERT INTO audit (...) VALUES (...)") # Either both inserts happened, or neither. # File locks — release on exit no matter what with lockfile.lock("/tmp/job.lock"): do_exclusive_work() # pytest's `raises` — assert a 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. Context managers make forgetting impossible.
Common Mistakes
1. Returning True from __exit__ by accident
def __exit__(self, exc_type, exc_val, exc_tb): self.cleanup() return self.cleanup() # cleanup returns truthy → all exceptions swallowed
Don't return anything from __exit__ unless you genuinely 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 # ← caller's `as c` binds to 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 · defines contextmanager
# 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')
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 the __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 # RuntimeError: generator already executed
setup added so this can run · defines timer
# 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()')
@contextmanager produces a one-shot CM. Call the factory fresh each time: with timer("twice"): .... Class-based CMs can be reusable if you write them to be — but the generator form isn'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: stash 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 original with chdir_gen("/tmp"): raise RuntimeError("boom") # cwd still restored, exception still propagates
Hint 1 — Save before you change
Stashos.getcwd() into self._saved (class) or a local (generator) before the os.chdir(path) call. If chdir raises (bad path), you haven't moved — but you also haven't corrupted state.
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 → falsy → exception (if any) 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 identically 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 matches the shape of most one-shot CMs.
Python 3.11 added contextlib.chdir to the standard library — this exact context manager. Before that, you wrote it yourself, and the class form above was the canonical recipe. Good news: now you can from contextlib import chdir and skip both implementations. Better news: you understand exactly what it's doing under the hood.
What You Learned
- The protocol:
__enter__(self)runs on entry;__exit__(self, exc_type, exc_val, exc_tb)runs on exit, always. - Returning
Truefrom__exit__suppresses the exception. Almost never what you want; default toreturn 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 mandatory for safe cleanup. Exactly oneyield.ExitStackfor a dynamic number of CMs — open N things in one block, register cleanup callbacks, conditional acquisition.contextlib.suppressfor "ignore this specific error";closingfor.close()-only objects;nullcontextfor conditional CMs.redirect_stdout/redirect_stderrfor capturing print output.async withuses__aenter__/__aexit__— same protocol, async methods.- Real-world wins: transactions, locks,
pytest.raises, timing blocks, directory changes — anywhere a pair of operations must always happen together.
Next: Async — async/await, event loops, and the parallel universe of async context managers and iterators.