contextlib.contextmanager decorator simplified

The contextlib.contextmanager decorator transforms a generator function into a context manager, eliminating the need to define __enter__ and __exit__ methods. Instead, you write a simple generator: code before the yield runs in __enter__, and code after the yield runs in __exit__. This approach is cleaner, more concise, and ideal for lightweight, single-purpose context managers.

The decorator is part of the Python standard library and has been available since Python 2.5. It is one of the most elegant features of the contextlib module and is used extensively in real code for temporary state changes, resource management, and error handling patterns that would otherwise require boilerplate class definitions.

How contextlib.contextmanager Works

When you decorate a generator function with @contextlib.contextmanager, the decorator returns a context manager factory. Each call to the decorated function returns a new context manager instance. When the context manager is entered, the generator runs until it yields; the yielded value is returned to the as clause. When the context manager exits, the generator resumes after the yield and runs to completion.

from contextlib import contextmanager

# Simplest example: generator-based context manager
@contextmanager
def simple_context():
    print("Entering (before yield)")
    yield "resource"
    print("Exiting (after yield)")

# Using the context manager
with simple_context() as res:
    print(f"Inside with block: {res}")

# Output:
# Entering (before yield)
# Inside with block: resource
# Exiting (after yield)

The decorator handles the protocol for you: the code before yield becomes __enter__, and the code after yield becomes __exit__. This reduces boilerplate from 5-10 lines to 2-3, making simple context managers trivial to write.

Example 1: File Writing with State Cleanup

A common use case is temporarily changing state, using a resource, and restoring the original state. The generator pattern makes this crystal clear.

from contextlib import contextmanager
import sys

@contextmanager
def redirected_output(target_file):
    """Redirect stdout to a file, then restore the original stdout."""
    original_stdout = sys.stdout
    try:
        sys.stdout = open(target_file, 'w')
        yield  # Code in with block runs here
    finally:
        # Always restore stdout, even if an exception occurs
        sys.stdout.close()
        sys.stdout = original_stdout

# Using the context manager
import tempfile
import os

with tempfile.NamedTemporaryFile(mode='w', delete=False, suffix='.txt') as tmp:
    tmp_name = tmp.name

try:
    with redirected_output(tmp_name):
        print("This goes to the file")
        print("So does this")
    
    print("This goes to the console")
    
    with open(tmp_name, 'r') as f:
        print(f"File contents: {f.read()}")
finally:
    os.unlink(tmp_name)

# Output:
# This goes to the console
# File contents: This goes to the file
# So does this

Example 2: Database Transaction Context Manager

Transaction management is a perfect use case for generator-based context managers: begin the transaction before yield, commit on success, rollback on failure.

from contextlib import contextmanager

class FakeDatabase:
    def __init__(self):
        self.data = {}
        self.in_transaction = False
    
    def begin(self):
        print("BEGIN TRANSACTION")
        self.in_transaction = True
    
    def commit(self):
        print("COMMIT")
        self.in_transaction = False
    
    def rollback(self):
        print("ROLLBACK")
        self.in_transaction = False
    
    def set(self, key, value):
        if not self.in_transaction:
            raise RuntimeError("Not in transaction")
        self.data[key] = value
        print(f"SET {key} = {value}")
    
    def get(self, key):
        return self.data.get(key)

db = FakeDatabase()

@contextmanager
def transaction(database):
    """Context manager for database transactions."""
    database.begin()
    try:
        yield database
        database.commit()
    except Exception as e:
        database.rollback()
        print(f"Exception: {e}")
        raise

# Successful transaction
print("--- Successful transaction ---")
with transaction(db) as db_conn:
    db_conn.set("name", "Alice")
    db_conn.set("age", 30)

print(f"Data: {db.data}\n")

# Failed transaction with rollback
print("--- Transaction with error ---")
db.data = {}  # Reset
try:
    with transaction(db) as db_conn:
        db_conn.set("name", "Bob")
        raise ValueError("Invalid operation")
except ValueError:
    pass

print(f"Data after rollback: {db.data}")

# Output:
# --- Successful transaction ---
# BEGIN TRANSACTION
# SET name = Alice
# SET age = 30
# COMMIT
# Data: {'name': 'Alice', 'age': 30}
#
# --- Transaction with error ---
# BEGIN TRANSACTION
# SET name = Bob
# Exception: Invalid operation
# ROLLBACK
# Data after rollback: {}

Example 3: Temporary File or Directory

The contextlib approach makes it easy to create temporary resources that are automatically cleaned up.

from contextlib import contextmanager
import tempfile
import os
import shutil

@contextmanager
def temporary_directory(prefix="tmp_"):
    """Create and clean up a temporary directory."""
    tmpdir = tempfile.mkdtemp(prefix=prefix)
    try:
        print(f"Created temporary directory: {tmpdir}")
        yield tmpdir
    finally:
        print(f"Removing temporary directory: {tmpdir}")
        shutil.rmtree(tmpdir, ignore_errors=True)

# Using the temporary directory
with temporary_directory(prefix="myapp_") as tmpdir:
    test_file = os.path.join(tmpdir, "test.txt")
    with open(test_file, 'w') as f:
        f.write("Temporary content")
    print(f"Files in {tmpdir}: {os.listdir(tmpdir)}")

# Directory is automatically cleaned up
print("Temporary directory cleaned up")

# Output:
# Created temporary directory: /tmp/myapp_abcd1234
# Files in /tmp/myapp_abcd1234: ['test.txt']
# Removing temporary directory: /tmp/myapp_abcd1234
# Temporary directory cleaned up

Exception Handling with contextmanager

By default, exceptions in the with block are re-raised after the finally block. You can inspect exceptions using sys.exc_info() or by catching them around the yield.

from contextlib import contextmanager
import sys

@contextmanager
def exception_aware():
    """Context manager that logs exceptions."""
    try:
        print("Acquiring resource")
        yield
        print("Normal exit")
    except Exception as e:
        print(f"Caught exception: {type(e).__name__}: {e}")
        # Re-raise to propagate
        raise
    finally:
        print("Cleanup")

# Test with exception
print("--- With exception ---")
try:
    with exception_aware():
        raise ValueError("Test error")
except ValueError:
    print("Exception propagated to caller")

print()

# Test without exception
print("--- Without exception ---")
with exception_aware():
    print("Normal operation")

# Output:
# --- With exception ---
# Acquiring resource
# Caught exception: ValueError: Test error
# Cleanup
# Exception propagated to caller
#
# --- Without exception ---
# Acquiring resource
# Normal operation
# Cleanup

Parameterized Context Managers

Context managers created with the decorator accept parameters just like regular functions, making them flexible and reusable.

from contextlib import contextmanager

@contextmanager
def logging_context(operation, verbose=False):
    """Log the start and end of an operation."""
    if verbose:
        print(f"Starting operation: {operation}")
    try:
        yield
    finally:
        if verbose:
            print(f"Finished operation: {operation}")

# Using with different parameters
print("--- Verbose mode ---")
with logging_context("Database query", verbose=True):
    print("Running query...")

print("\n--- Silent mode ---")
with logging_context("Background task", verbose=False):
    print("Running task...")

# Output:
# --- Verbose mode ---
# Starting operation: Database query
# Running query...
# Finished operation: Database query
#
# --- Silent mode ---
# Running task...

contextlib.ExitStack: Dynamic Context Management

For complex scenarios with an unknown number of resources, contextlib.ExitStack provides a powerful pattern: dynamically register cleanup functions that will be called in reverse order.

from contextlib import ExitStack
import tempfile
import os

@contextmanager
def many_files(filenames):
    """Open multiple files dynamically and clean them all up."""
    with ExitStack() as stack:
        files = {}
        for filename in filenames:
            # Register each file for cleanup
            files[filename] = stack.enter_context(open(filename, 'w'))
        yield files

import tempfile

# Create temporary files for testing
files_to_create = []
for i in range(3):
    tmp = tempfile.NamedTemporaryFile(mode='w', delete=False, suffix=f'_{i}.txt')
    files_to_create.append(tmp.name)
    tmp.close()

try:
    with many_files(files_to_create) as files:
        for fname, fobj in files.items():
            fobj.write(f"Content of {fname}")
        print(f"Opened {len(files)} files")
    print("All files cleaned up")
finally:
    for fname in files_to_create:
        os.unlink(fname)

# Output:
# Opened 3 files
# All files cleaned up

contextmanager vs. Class-Based Context Managers

Aspect	@contextmanager	Class-Based
Lines of code	2-5	6-15
Readability	Very clear; sequential	Must read both methods
State tracking	Via local variables	Via instance variables
Exception handling	Simple try-finally	Explicit __exit__ logic
Reusability	Each call is new	Can reuse instance
Debugging	Stack trace clear	One extra frame

Choose @contextmanager for simple, single-use context managers. Use class-based for complex state, multiple cleanup steps, or if you need to reuse a single instance.

Key Takeaways

• The @contextlib.contextmanager decorator converts a generator function into a context manager, eliminating boilerplate.
• Code before yield runs on entry; code after yield runs on exit, with exception safety via try-finally.
• This pattern is ideal for state cleanup, temporary resources, transactions, and error handling.
• Parameters work just like regular functions, enabling flexible, reusable context managers.
• For dynamic cleanup with an unknown number of resources, use contextlib.ExitStack.

Frequently Asked Questions

What if my generator yields None?

You can yield None or nothing (implicit None). The yield value is what gets bound to the as variable. If you don't need to expose a resource, yielding nothing is fine.

Can I catch and suppress exceptions in a @contextmanager?

Yes. Wrap the yield in a try-except block. Returning True from __exit__ (implicitly via the decorator) would suppress the exception. However, the standard pattern is to let exceptions propagate.

What happens if code after yield raises an exception?

Any exception after yield propagates after the context manager exits, replacing the original exception (if any). Use exception chaining (raise ... from ...) to preserve the original.

Can I make a reusable @contextmanager-based context manager?

Yes, but be careful. The decorator returns a new context manager each call, so reusing the same instance requires calling it again in a new with statement. For reuse, use a class-based context manager instead.

How is ExitStack different from nesting with statements?

ExitStack allows dynamic cleanup registration (you don't know the number of resources in advance). With nesting, you must know the resources at write time. For most cases, nesting or multiple context managers is clearer.

Further Reading

• contextlib module documentation
• PEP 343: contextmanager decorator design
• ExitStack for advanced patterns
• Real Python: contextlib recipes