Context managers (with, contextlib)

In Chapter 6 errors and exception handling we briefly noted that with is cleaner than try/finally. This chapter covers that topic. We cover every pattern of with, user-defined context managers, and the tools in contextlib.

This chapter’s @contextmanager builds on the generators in Chapter 11 iterables, generators, yield from, so the two chapters are paired. The async with we preview at the end shows up in earnest in Chapter 14 asyncio intro.

with — resources collapsed to one line #

The most basic example — file handling.

f = open("data.txt")
try:
    data = f.read()
finally:
    f.close()

with open("data.txt") as f:
    data = f.read()
# f.close() is called automatically here

with automatically runs cleanup code on block exit. Whether you exit normally, via an exception, or via return, cleanup happens.

Resources that suit with #

What they share: “step into some state briefly, then revert when leaving.” That’s the problem context managers solve.

Multiple resources at once #

with open("input.txt") as src, open("output.txt", "w") as dst:
    dst.write(src.read())

From 3.10 onward you can also wrap in parentheses and split across lines.

with (
    open("input.txt") as src,
    open("output.txt", "w") as dst,
    Lock() as lock,
):
    dst.write(src.read())

Reads well — same shape as a function with many parameters.

Building your own — __enter__ / __exit__ #

Objects that work with with have two methods.

• __enter__(self) — called on entering with, returns the value bound after as
• __exit__(self, exc_type, exc_value, traceback) — called on exit

Simplest example — temporary chdir #

When you want to briefly change the working directory.

import os

class chdir:
    def __init__(self, path: str):
        self.path = path

    def __enter__(self):
        self.old = os.getcwd()
        os.chdir(self.path)
        return self.path

    def __exit__(self, exc_type, exc_value, tb):
        os.chdir(self.old)

with chdir("/tmp"):
    print(os.getcwd())     # /tmp
print(os.getcwd())          # original directory

__exit__’s return value — whether to swallow exceptions #

__exit__ swallows exceptions when it returns a truthy value. If you don’t know this and return True, every exception disappears and debugging gets hard.

class Bad:
    def __enter__(self): return self
    def __exit__(self, *args):
        return True   # ✗ swallows every exception

In most cases, return nothing, or return False / None.

To intentionally swallow only specific exceptions, look at exc_type and decide.

class IgnoreFileNotFound:
    def __enter__(self): return self
    def __exit__(self, exc_type, exc_value, tb):
        return exc_type is not None and issubclass(exc_type, FileNotFoundError)

with IgnoreFileNotFound():
    with open("missing.txt") as f:
        ...
# passes without exception

This common pattern already exists in the standard library (see suppress below).

@contextmanager — one function, short #

Building a class with both __enter__ and __exit__ every time is tedious. There’s a way to make a context manager from a single function with yield in the middle.

from contextlib import contextmanager

@contextmanager
def chdir(path: str):
    old = os.getcwd()
    os.chdir(path)
    try:
        yield path        # the with block runs at this point
    finally:
        os.chdir(old)

This one function does exactly what the class version above does.

How to read it:

• Before yield = what __enter__ does
• The yielded value = what as binds
• After yield = what __exit__ does
• Wrap with try/finally to guarantee cleanup even on exceptions

How @contextmanager is built on top of generators is revisited in Chapter 11 iterables, generators, yield from.

If you want to catch exceptions #

@contextmanager
def transactional(conn):
    tx = conn.begin()
    try:
        yield tx
    except Exception:
        tx.rollback()
        raise
    else:
        tx.commit()

If the with block throws an exception at the yield point, that exception is re-thrown at the generator’s yield point. So you can catch it with a regular try/except.

Other helpers in contextlib #

suppress — ignore specific exceptions #

from contextlib import suppress
import os

with suppress(FileNotFoundError):
    os.remove("maybe-missing.txt")
# passes through even if the file is missing

The IgnoreFileNotFound class we built above is right there in the standard library.

closing — wrap an object that only has a close() method #

from contextlib import closing
from urllib.request import urlopen

with closing(urlopen("https://example.com")) as resp:
    data = resp.read()
# resp.close() runs automatically

Often used in older versions where urlopen’s response object didn’t have __exit__. These days many response objects support with directly.

redirect_stdout — intercept stdout #

import io
from contextlib import redirect_stdout

buf = io.StringIO()
with redirect_stdout(buf):
    print("hello")
print("captured:", buf.getvalue())
# captured: hello

Useful in tests, captures, and log conversion. redirect_stderr works the same way.

nullcontext — conditional with #

from contextlib import nullcontext

def process(path: str | None):
    ctx = open(path) if path else nullcontext()
    with ctx as f:
        if f is None:
            print("no file")
        else:
            print(f.read())

“Open the file if there is one, otherwise just pass through.” with always needs a context manager after it, so you need a harmless no-op context manager. nullcontext is that.

ExitStack — manage many resources dynamically #

When the number of resources is decided at runtime.

from contextlib import ExitStack

def merge(paths: list[str]) -> str:
    with ExitStack() as stack:
        files = [stack.enter_context(open(p)) for p in paths]
        return "\n".join(f.read() for f in files)
# all files are closed automatically

When you can’t write N copies of with open(p) as f (because the paths list length is dynamic), ExitStack solves it. All registered resources are cleaned up in reverse order.

ExitStack can also register plain cleanup functions that aren’t __enter__-based.

def setup_temp_resources():
    stack = ExitStack()
    tmp_dir = create_tmp_dir()
    stack.callback(remove_dir, tmp_dir)
    # ... register more resources
    return stack

with setup_temp_resources() as stack:
    ...

async context managers — a quick preview #

For async resources, use async with.

async with open_async_db() as conn:
    await conn.execute(...)

The methods here are __aenter__ / __aexit__. There’s also @asynccontextmanager, and usage is the same as the sync version. Details in Chapter 14 asyncio intro.

User-defined — class vs @contextmanager #

When in doubt, start with @contextmanager and move to a class as things get more complex.

Common patterns — in one place #

Temporary environment variables #

import os
from contextlib import contextmanager

@contextmanager
def set_env(**kwargs: str):
    old = {}
    for k, v in kwargs.items():
        old[k] = os.environ.get(k)
        os.environ[k] = v
    try:
        yield
    finally:
        for k, v in old.items():
            if v is None:
                del os.environ[k]
            else:
                os.environ[k] = v

Often used when verifying env-var-dependent code in tests.

Timing #

import time
from contextlib import contextmanager

@contextmanager
def timer(label: str):
    start = time.perf_counter()
    yield
    elapsed = time.perf_counter() - start
    print(f"{label}: {elapsed:.3f}s")

with timer("query"):
    do_expensive_thing()
# query: 0.123s

This pattern isn’t a precise measurement tool. Precise profiling is covered in Chapter 21 performance — cProfile, py-spy, memory profiling.

Locks #

import threading

lock = threading.Lock()

with lock:
    # critical section
    update_shared_state()

Lock, RLock, and Semaphore all support with. Code that calls acquire() / release() directly is almost never written.

Exercises #

1. Write set_env(**kwargs: str) with @contextmanager. On entering the with block, set the environment variables; on exit, restore the original values. The case where the original value was missing (None) must also be restored correctly.
2. Write a class TempFile:. __enter__ creates a temp file and returns the path; __exit__ deletes the file. Verify the file is deleted even if an exception occurs inside the with block.
3. Write def merge(paths: list[str]) -> str: with ExitStack. Open every path passed in, concatenate the contents, and return them. Confirm that already-opened files are properly closed even if one of the paths mid-list raises FileNotFoundError (trace open file handles with a debugger).

In one line: with compresses the enter / cleanup boundary into one line. User-defined versions are either a class (__enter__/__exit__) or @contextmanager (function + yield). __exit__’s return value is normally False — True swallows exceptions. contextlib’s suppress / closing / redirect_stdout / nullcontext / ExitStack are the practical helpers. Use ExitStack for a dynamic number of resources. Async uses async with.

Next chapter #

In Chapter 11 iterables, generators, yield from we cover the tools beyond comprehensions — user-defined iterables, generator functions, and yield from. This chapter’s @contextmanager is in fact a tool built on top of generators.