Python Context Managers with contextlib

The contextlib module is the standard library's toolkit for working with context managers, and learning Python context managers with contextlib opens up patterns that are difficult or verbose to implement with classes alone. The module provides the contextmanager decorator, which turns a generator function into a context manager with a single yield statement splitting the setup from the teardown. It provides pre-built context managers for common tasks like suppressing exceptions and redirecting output. And it provides ExitStack, a context manager that manages a dynamic collection of other context managers, solving the problem of not knowing at compile time how many resources a block will need.

The previous article showed how to create custom context managers by writing classes with enter and exit methods. That approach gives you full control and is the right choice for complex resource management. But many context managers are simple wrappers around a setup-teardown pair, and writing a full class for each one adds boilerplate that obscures the actual logic. The contextmanager decorator reduces a context manager to a generator function with a single yield, where the code before the yield is the enter logic and the code after the yield is the exit logic. The pattern is concise, readable, and covers the majority of custom context managers you will write.

Understanding contextlib also means understanding ExitStack, which solves a problem that the comma-separated with statement cannot handle: managing a variable number of context managers. When you know at compile time that you need exactly two files open, you write them in a single with statement separated by a comma, as shown in using the with statement in Python. But when you need to open a list of files whose length is determined at runtime, you need ExitStack. It collects context managers as you push them, enters each one when pushed, and exits all of them in reverse order when the stack itself is exited.

The contextmanager decorator

The contextmanager decorator transforms a generator function into a context manager. The generator must yield exactly once. The code before the yield becomes the enter method, running when the with block begins. The yielded value becomes the return value of enter, available through the as clause. The code after the yield becomes the exit method, running when the with block ends. If the block raises an exception, the generator's exit logic still runs because the with statement calls the generator's close or throw method to inject the exception at the yield point.

Here is a simple context manager written with the decorator that temporarily changes the current working directory and restores it afterwards:

python

import os
from contextlib import contextmanager
 
@contextmanager
def change_dir(path):
    old_path = os.getcwd()
    os.chdir(path)
    try:
        yield
    finally:
        os.chdir(old_path)

The code before the yield saves the current directory and changes to the new one. The yield itself produces no value, so the as clause is omitted. The code after the yield restores the original directory. The try-finally around the yield ensures that the directory is restored even if the block raises an exception. This try-finally pattern is essential in generator-based context managers because the yield point is where the block runs, and any exception from the block is raised at that point. Wrapping the yield in try-finally guarantees that the teardown code runs regardless.

Using this context manager is identical to using a class-based one. The decorator makes the generator function return a context manager object when called, and the with statement uses that object's enter and exit methods, which the decorator synthesized from the generator:

python

with change_dir("/tmp"):
    print(os.getcwd())

The block runs with the current directory set to the temporary path. When the block ends, the original directory is restored. If the block raised an exception, the directory would still be restored because the finally clause around the yield runs before the exception propagates. Since Python 3.11, the standard library ships this exact tool as contextlib.chdir, so in real code you would reach for that, but writing it yourself is the clearest possible demonstration of the decorator pattern.

What the decorator cannot do

The contextmanager decorator is concise, but exception handling works differently than in a class. In a class-based context manager, the exit method receives the exception type, value, and traceback directly as parameters, so branching on them is natural. In a generator-based context manager, the exception is instead raised at the yield point, so any inspection or branching has to happen in an except clause wrapped around the yield.

If you need to react differently to success and failure, the generator-based approach still works, but it requires an explicit try-except-else structure around the yield. Here is a transaction context manager that commits on success and rolls back on failure, written with the decorator:

python

from contextlib import contextmanager
 
@contextmanager
def transaction(connection):
    connection.begin()
    try:
        yield connection
    except Exception:
        connection.rollback()
        raise
    else:
        connection.commit()

The yield produces the connection so the block can use it. The except clause catches any exception, rolls back the transaction, and re-raises the exception so the caller knows the operation failed. The else clause runs only if no exception occurred, and it commits the transaction. This pattern works correctly but is slightly less readable than the class-based version where commit and rollback are symmetric branches of an if-statement in exit.

For context managers that need to distinguish between normal and exceptional exit, a class with enter and exit methods is often clearer than a generator with try-except-else around the yield. The class version puts the two paths side by side in a single method, while the generator version scatters them across except and else clauses. Choose the generator approach when the teardown logic is the same regardless of exceptions (like restoring a directory or closing a file). Choose the class approach when the teardown logic branches based on the exception.

Pre-built context managers in contextlib

The contextlib module includes several ready-made context managers for common patterns. The suppress context manager catches and ignores specified exceptions, replacing the try-except-pass pattern with a named context. It is useful when an operation might fail for expected reasons and you do not need to handle the failure beyond ignoring it:

python

import os
from contextlib import suppress
 
with suppress(FileNotFoundError):
    os.remove("optional_file.txt")

Use suppress only when the failure is genuinely acceptable to ignore, such as deleting a file that may not exist. If the caller needs to know the operation failed, an explicit try-except that logs or reports the error is the better tool. The module also provides closing, which wraps any object that has a close method in a context manager, useful for older libraries that predate the context manager protocol.

The redirect_stdout and redirect_stderr context managers temporarily replace the standard output and error streams. They are primarily used in tests to capture output that would normally go to the console. The stream is restored when the with block ends:

python

import io
from contextlib import redirect_stdout
 
buffer = io.StringIO()
with redirect_stdout(buffer):
    print("This goes to the buffer")
output = buffer.getvalue()

ExitStack for dynamic resource management

ExitStack solves the problem of managing a variable number of context managers. In a standard with statement, you list all context managers separated by commas, and all are known at compile time. But when you are iterating over a list of filenames and opening each one, you cannot use a comma-separated with statement because the number of files depends on the list length.

ExitStack lets you push context managers onto a stack inside a with block. Each context manager's enter method runs when you push it. When the ExitStack's own with block ends, it calls exit on every context manager in reverse order, guaranteeing cleanup. Here is an example that opens every file in a list and collects their contents:

python

from contextlib import ExitStack
 
def read_all(file_paths):
    with ExitStack() as stack:
        files = [stack.enter_context(open(path)) for path in file_paths]
        return [f.read() for f in files]

The enter_context method on the ExitStack enters the context manager (opening the file) and returns the value from its enter method (the file object). The files list collects the open file objects. When the ExitStack block ends, all files are closed in reverse order. If any file operation raises an exception, the files that were already opened are still closed.

ExitStack is also useful for combining context managers that depend on each other. When one context manager's creation depends on a value from a previously entered context manager, the comma-separated with statement cannot express the dependency. ExitStack handles this naturally because you push context managers one at a time with full access to values from earlier ones.

The next article covers nested context managers in Python, exploring how context managers compose and how the with statement handles nesting, both explicit and implicit, along with the parenthesized multi-line syntax introduced in Python 3.10.