Python contextlib — Deep Dive

The context manager protocol

Before diving into contextlib’s advanced features, it helps to understand the raw protocol. A context manager implements two methods:

class ManagedResource:
    def __enter__(self):
        # Setup: acquire resource
        # Return value becomes the "as" variable
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        # Cleanup: release resource
        # Return True to suppress the exception, False/None to propagate
        self.close()
        return False

The __exit__ signature receives exception info if one occurred inside the with block. This is the mechanism behind contextlib’s exception handling capabilities.

@contextmanager exception handling in depth

The @contextmanager decorator wraps your generator in a class that handles the protocol. When an exception occurs in the with block, it’s thrown into the generator at the yield point:

from contextlib import contextmanager

@contextmanager
def transaction(conn):
    conn.begin()
    try:
        yield conn
    except Exception:
        conn.rollback()
        raise           # re-raise to propagate
    else:
        conn.commit()   # only if no exception

If you catch the exception and don’t re-raise it, the exception is suppressed:

@contextmanager
def suppress_value_errors():
    try:
        yield
    except ValueError:
        pass  # exception suppressed — with block appears to succeed

with suppress_value_errors():
    int("not a number")  # ValueError suppressed
print("continues normally")

This is powerful but dangerous — suppressing exceptions silently can hide bugs. Be explicit about what you’re catching.

Async context managers

contextlib provides async equivalents for all major features:

@asynccontextmanager

from contextlib import asynccontextmanager

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

async with async_db_session(pool) as conn:
    await conn.execute("SELECT 1")

AsyncExitStack

from contextlib import AsyncExitStack

async def process_connections(urls):
    async with AsyncExitStack() as stack:
        connections = []
        for url in urls:
            conn = await stack.enter_async_context(connect(url))
            connections.append(conn)
        # all connections managed, all cleaned up on exit

ExitStack advanced patterns

Transferring ownership

ExitStack.pop_all() transfers all registered cleanups to a new stack, letting you prevent cleanup when you want to return a resource to the caller:

from contextlib import ExitStack

def open_managed_files(paths):
    """Open files; caller takes ownership of cleanup."""
    with ExitStack() as stack:
        files = [stack.enter_context(open(p)) for p in paths]

        # Validate all files before returning
        for f in files:
            if f.read(1) == "":
                raise ValueError(f"Empty file: {f.name}")
            f.seek(0)

        # Transfer cleanup responsibility to caller
        return stack.pop_all(), files

# Caller receives cleanup_stack and must close it
cleanup, files = open_managed_files(["a.txt", "b.txt"])
try:
    process(files)
finally:
    cleanup.close()  # closes all files

Conditional context managers

from contextlib import ExitStack

def process_data(path, use_profiling=False, use_lock=False):
    with ExitStack() as stack:
        if use_profiling:
            stack.enter_context(profiler())
        if use_lock:
            stack.enter_context(file_lock(path))

        data = open(path).read()
        return transform(data)

Cleanup callbacks with arguments

from contextlib import ExitStack

with ExitStack() as stack:
    tmpdir = create_temp_dir()
    stack.callback(shutil.rmtree, tmpdir)

    # push with specific exception info
    stack.callback(print, "cleanup complete", flush=True)

Callbacks are called in LIFO order, same as atexit.

AbstractContextManager and AbstractAsyncContextManager

These ABCs (from contextlib) provide default implementations and can be used for type checking:

from contextlib import AbstractContextManager

class Timer(AbstractContextManager):
    def __enter__(self):
        self.start = time.perf_counter()
        return self

    def __exit__(self, *exc):
        self.elapsed = time.perf_counter() - self.start
        return False  # don't suppress exceptions

Using AbstractContextManager gives you a __enter__ that returns self by default, so you only need to implement __exit__.

chdir context manager (Python 3.11+)

Python 3.11 added contextlib.chdir():

from contextlib import chdir
import os

print(os.getcwd())  # /home/user

with chdir("/tmp"):
    print(os.getcwd())  # /tmp
    # do work in /tmp

print(os.getcwd())  # /home/user — restored automatically

Thread safety warning: os.chdir() affects the entire process. In multi-threaded applications, chdir() is not safe — use absolute paths instead.

Exception handling patterns

Selective exception suppression with logging

from contextlib import contextmanager
import logging

logger = logging.getLogger(__name__)

@contextmanager
def log_and_suppress(*exception_types):
    try:
        yield
    except exception_types as e:
        logger.warning("Suppressed exception: %s: %s", type(e).__name__, e)

Exception transformation

@contextmanager
def convert_exceptions(from_type, to_type):
    try:
        yield
    except from_type as e:
        raise to_type(str(e)) from e

with convert_exceptions(KeyError, ValueError):
    data = {"a": 1}
    _ = data["missing"]  # KeyError → ValueError

Retry context manager

@contextmanager
def retry_on_failure(max_retries=3, delay=1):
    """Not a true context manager pattern — shows yield-once limitation."""
    for attempt in range(max_retries):
        try:
            yield attempt
            return  # success
        except Exception as e:
            if attempt == max_retries - 1:
                raise
            time.sleep(delay)

Note: this doesn’t work as written because @contextmanager requires exactly one yield. For retry logic, use a regular function or loop — context managers aren’t the right abstraction.

Production patterns

Database transaction manager

from contextlib import contextmanager

@contextmanager
def transaction(engine, isolation_level="READ COMMITTED"):
    conn = engine.connect()
    conn.execute(f"SET TRANSACTION ISOLATION LEVEL {isolation_level}")
    conn.begin()
    try:
        yield conn
    except Exception:
        conn.rollback()
        raise
    else:
        conn.commit()
    finally:
        conn.close()

Temporary environment variables

from contextlib import contextmanager
import os

@contextmanager
def env_vars(**kwargs):
    old_values = {}
    for key, value in kwargs.items():
        old_values[key] = os.environ.get(key)
        if value is None:
            os.environ.pop(key, None)
        else:
            os.environ[key] = value
    try:
        yield
    finally:
        for key, old in old_values.items():
            if old is None:
                os.environ.pop(key, None)
            else:
                os.environ[key] = old

with env_vars(DATABASE_URL="sqlite:///test.db", DEBUG="1"):
    run_tests()
# Original env vars restored

Timing and profiling

from contextlib import contextmanager
import time

@contextmanager
def timed(label, threshold_ms=None):
    start = time.perf_counter()
    yield
    elapsed_ms = (time.perf_counter() - start) * 1000
    msg = f"{label}: {elapsed_ms:.1f}ms"
    if threshold_ms and elapsed_ms > threshold_ms:
        logger.warning(f"SLOW: {msg}")
    else:
        logger.debug(msg)

Performance notes

The @contextmanager decorator adds minimal overhead — about 2-3μs per with block entry/exit compared to a raw class-based context manager (~0.5μs). For code that enters context managers millions of times per second, the class-based approach is faster. For everything else, the readability of @contextmanager wins.

ExitStack adds more overhead (~5μs per entry) because it maintains a list of callbacks. This is negligible unless you’re managing hundreds of resources in a tight loop.

One thing to remember

contextlib’s power comes from composability — @contextmanager for simple setup/cleanup, ExitStack for dynamic resource counts, suppress for clean exception handling, and async variants for the async world. Master these four and you can express any resource management pattern cleanly.