PythonMastery
advanced 18 min read · lesson 3 of 9 in Python Advanced

Context Managers: Beyond with-open

1 · The lesson

read

You'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:

python
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:

python
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.

python
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.

python
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.

python
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:

python
@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:

python
@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:

python
@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.

python
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:

python
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:

python
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:

python
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.

python
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:

python
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:

python
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:

python
# 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

python
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

python
@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

python
@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:

python
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

python
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.chdir back to the captured directory.

  • Do not suppress exceptions.

  • Implement it twice: once as a class, once with @contextmanager.

python
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 Stash os.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
python
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 True from __exit__ suppresses the exception. Almost never what you want; default to return False or no return.
  • Class-based CMs for stateful or reusable resources; __enter__ returns whatever as should bind to.
  • @contextmanager turns a generator into a CM: setup before yield, teardown after. try/finally around the yield is mandatory for safe cleanup. Exactly one yield.
  • ExitStack for a dynamic number of CMs — open N things in one block, register cleanup callbacks, conditional acquisition.
  • contextlib.suppress for "ignore this specific error"; closing for .close()-only objects; nullcontext for conditional CMs.
  • redirect_stdout / redirect_stderr for capturing print output.
  • async with uses __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.