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

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

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')

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.

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

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

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

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

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

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

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 is closed, in LIFO order (last opened, first closed).

You can also register 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, 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:

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

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

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()')

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:

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

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

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

python
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

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

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

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                             # 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.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: 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 Save 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, 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
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 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 True from __exit__ suppresses the exception. This is almost never what you want. Use 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 required for safe cleanup. Exactly one yield.
  • ExitStack for a changing number of CMs: open N things in one block, register cleanup callbacks, or hold a resource only when needed.
  • contextlib.suppress for "ignore this specific error", closing for .close()-only objects, and nullcontext for conditional CMs.
  • redirect_stdout / redirect_stderr for capturing print output.
  • async with uses __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.