Context managers and contextlib

Why this matters

with is how Python guarantees cleanup — closing files, releasing locks, rolling back transactions — without verbose try/finally. Knowing how to write context managers (both class-based and generator-based) is the difference between "I can use open(...)" and "I can build robust resource APIs."

Learning objectives

Use with for files, locks, network resources.
Write context managers two ways: class with __enter__/__exit__ and @contextmanager.
Use ExitStack for dynamic / variable-number resources.
Use the with (a, b, c): parenthesised form (3.10+).
Write async context managers (__aenter__/__aexit__ / @asynccontextmanager).

1. The protocol

python

with EXPR as var:
    BODY

Desugars to:

python

cm = EXPR
var = cm.__enter__()
try:
    BODY
except:
    if not cm.__exit__(*sys.exc_info()):
        raise
else:
    cm.__exit__(None, None, None)

__enter__ runs at entry; its return value binds to var.
__exit__(exc_type, exc, tb) always runs. If it returns truthy, the exception is suppressed.

2. Class-based context manager

python

class Timer:
    def __enter__(self):
        self.start = time.perf_counter()
        return self                    # accessible as `t` in `with Timer() as t`

    def __exit__(self, exc_type, exc, tb):
        self.elapsed = time.perf_counter() - self.start
        # return None (falsy) -> don't suppress exceptions

with Timer() as t:
    do_stuff()
print(f"{t.elapsed:.3f}s")

Class-based is the right choice when:

You need access to the instance after with.
You manage multi-step state.
You inherit cleanup from a base.

3. Generator-based with `@contextmanager`

python

from contextlib import contextmanager

@contextmanager
def timer():
    start = time.perf_counter()
    try:
        yield                          # `with timer():` binds nothing; `as t` would bind this value
    finally:
        print(f"{time.perf_counter() - start:.3f}s")

with timer():
    do_stuff()

The pattern:

Setup code before yield.
The yielded value is what as ... binds.
Teardown in finally.

try/except/finally lets you handle exceptions cleanly:

python

@contextmanager
def db_transaction(conn):
    cur = conn.cursor()
    try:
        yield cur
        conn.commit()
    except Exception:
        conn.rollback()
        raise
    finally:
        cur.close()

with db_transaction(conn) as cur:
    cur.execute("INSERT ...")

Generator-based is preferred for simple context managers — fewer lines than class-based.

4. `with EXPR1, EXPR2, ...` and the parenthesised form

Multiple resources without nesting:

python

with open("in.txt") as f_in, open("out.txt", "w") as f_out:
    f_out.write(f_in.read())

3.10+ allows parens for multi-line readability:

python

with (
    open("in.txt") as f_in,
    open("out.txt", "w") as f_out,
    Timer() as t,
):
    f_out.write(f_in.read())

5. `contextlib` essentials

python

from contextlib import (
    contextmanager, asynccontextmanager,
    suppress, ExitStack, AsyncExitStack,
    redirect_stdout, redirect_stderr,
    closing, nullcontext,
    chdir,                                # 3.11+
)

`suppress(*exceptions)`

"Ignore this exception if it happens":

python

with suppress(FileNotFoundError):
    Path("tmp").unlink()

`closing(thing)`

Wrap an object with a close() method as a context manager:

python

from contextlib import closing
with closing(get_legacy_connection()) as conn:
    use(conn)

`nullcontext`

A no-op context manager. Useful as a default:

python

def maybe_lock(lock):
    return lock if lock is not None else nullcontext()

with maybe_lock(lock):
    do_thing()

`redirect_stdout` / `redirect_stderr`

Capture print output:

python

import io
buf = io.StringIO()
with redirect_stdout(buf):
    noisy_function()
print(buf.getvalue())                  # what it would have printed

Great for testing legacy code that prints.

`chdir(path)` (3.11+)

Temporarily change directory:

python

with chdir(Path("/tmp")):
    do_work_in_tmp()

6. `ExitStack` — variable-number context managers

When the number of resources is dynamic:

python

from contextlib import ExitStack

paths = ["a.txt", "b.txt", "c.txt"]
with ExitStack() as stack:
    files = [stack.enter_context(open(p)) for p in paths]
    process(files)
# All files closed here, in reverse order

Or build up complex setup that can partially fail and unwind:

python

with ExitStack() as stack:
    conn = stack.enter_context(db.connect(url))
    cur = stack.enter_context(closing(conn.cursor()))
    lock = stack.enter_context(lock_for(resource_id))
    do_work(conn, cur)

If any setup fails midway, already-entered contexts are exited correctly.

Detaching / preserving

python

with ExitStack() as stack:
    f = stack.enter_context(open("x.txt"))
    stack.pop_all().close()       # caller now owns cleanup
    return f                      # leak-free transfer

7. Async context managers

python

class AsyncTimer:
    async def __aenter__(self):
        self.start = time.perf_counter()
        return self
    async def __aexit__(self, exc_type, exc, tb):
        self.elapsed = time.perf_counter() - self.start

async def main():
    async with AsyncTimer() as t:
        await do_async_work()
    print(t.elapsed)

async with calls __aenter__ / __aexit__ as coroutines.

`@asynccontextmanager`

python

from contextlib import asynccontextmanager

@asynccontextmanager
async def db_tx(pool):
    conn = await pool.acquire()
    try:
        async with conn.transaction():
            yield conn
    finally:
        await pool.release(conn)

async with db_tx(pool) as conn:
    await conn.execute("INSERT ...")

Phase 4.4 (asyncio) leans on these heavily — every httpx.AsyncClient, asyncpg.Pool, FastAPI lifespan, etc. is built on __aenter__/__aexit__.

8. Suppressing vs propagating exceptions

Return truthy from __exit__ to swallow:

python

class IgnoreError:
    def __init__(self, *exc_types):
        self.exc_types = exc_types
    def __enter__(self): return self
    def __exit__(self, exc_type, exc, tb):
        return exc_type is not None and issubclass(exc_type, self.exc_types)

with IgnoreError(FileNotFoundError):
    Path("never").unlink()

This is basically contextlib.suppress. Roll your own only if you need extra logic.

9. Common patterns

Lock

python

with lock:
    critical_section()

Temp file / temp dir

python

from tempfile import TemporaryDirectory, NamedTemporaryFile
with TemporaryDirectory() as d:
    work_in(Path(d))                   # cleaned up

Stop a server / stream

python

@contextmanager
def run_server(port):
    server = start_server(port)
    try:
        yield server
    finally:
        server.shutdown()

`with` in tests (pytest)

python

import pytest
def test_raises():
    with pytest.raises(ValueError, match="bad"):
        my_function(bad_input)

10. Worked example: a chunked uploader

python

from contextlib import contextmanager
import httpx

@contextmanager
def chunked_upload(url: str, chunk_size: int = 1024 * 1024):
    buffer = bytearray()
    client = httpx.Client(timeout=30)

    def write(data: bytes) -> None:
        buffer.extend(data)
        while len(buffer) >= chunk_size:
            chunk, _ = buffer[:chunk_size], buffer.__setitem__(slice(0, chunk_size), b"")
            client.post(url, content=chunk)

    try:
        yield write
        if buffer:
            client.post(url, content=bytes(buffer))     # flush remainder
    finally:
        client.close()

with chunked_upload("https://upload.example.com") as write:
    for block in read_blocks(huge_file):
        write(block)

The caller just calls write(data). Setup, flush, and cleanup are guaranteed.

Hands-on lab (1.5 hours)

Write @contextmanager def timer(): that prints elapsed time.
Re-implement contextlib.suppress as a class.
Build chdir(path) (yourself) without using the 3.11 helper; verify cwd restores even on exception.
Use ExitStack to open N files (from a list) and pass them to a function.
Build db_transaction(conn) as an @contextmanager; commit on success, rollback on error.
Write @asynccontextmanager for an httpx.AsyncClient with auth headers wired in.
Bonus: build a context manager that captures stdout, then asserts on its contents (for testing).

Common pitfalls

Forgetting try/finally in a generator-based context manager.
Returning a truthy value from __exit__ accidentally suppresses exceptions.
Using with on something that doesn't implement the protocol → TypeError.
Class-based CM that doesn't handle exceptions in __exit__.
Using sync with on an async resource.
Inside ExitStack, forgetting to call enter_context (just constructing the CM doesn't enter it).

Self-check

What does __exit__ return to suppress an exception?
When use @contextmanager vs class-based?
What does ExitStack solve?
Difference between with and async with?
What is nullcontext for?

References

PEP 343 — The "with" Statement.
PEP 492 — Coroutines with async/await syntax (async with).
PEP 617 — New PEG parser.
Fluent Python, Ramalho — Chapter 18.
contextlib docs.

Why this matters

Learning objectives

1. The protocol

2. Class-based context manager

3. Generator-based with @contextmanager

4. with EXPR1, EXPR2, ... and the parenthesised form

5. contextlib essentials

suppress(*exceptions)

closing(thing)

nullcontext

redirect_stdout / redirect_stderr

chdir(path) (3.11+)

6. ExitStack — variable-number context managers

Detaching / preserving

7. Async context managers

@asynccontextmanager

8. Suppressing vs propagating exceptions

9. Common patterns

Lock

Temp file / temp dir

Stop a server / stream

with in tests (pytest)

10. Worked example: a chunked uploader

Hands-on lab (1.5 hours)

Common pitfalls

Self-check

References

3. Generator-based with `@contextmanager`

4. `with EXPR1, EXPR2, ...` and the parenthesised form

5. `contextlib` essentials

`suppress(*exceptions)`

`closing(thing)`

`nullcontext`

`redirect_stdout` / `redirect_stderr`

`chdir(path)` (3.11+)

6. `ExitStack` — variable-number context managers

`@asynccontextmanager`

`with` in tests (pytest)