Advanced context manager patterns and techniques
Advanced context manager patterns solve complex scenarios: dynamic cleanup with an unknown number of resources, managing context-dependent state, resource pooling, and fault tolerance. ExitStack is the Swiss Army knife of context managers, allowing you to register cleanup functions at runtime in LIFO order. Context manager composition enables layering concerns: timing, logging, error handling, and resource management in a single, readable with statement. These patterns are used extensively in production frameworks and libraries.
Mastering these patterns transforms you from a user of context managers into a designer of them — capable of building robust, composable, fault-tolerant systems where resource management is declarative and auditable.
Pattern 1: contextlib.ExitStack for Dynamic Cleanup
ExitStack allows you to register cleanup functions that will be called in reverse order, enabling cleanup of resources that are unknown at write time.
# Example 1: Dynamic file management with ExitStack
from contextlib import ExitStack
import tempfile
import os
def process_many_files(filenames, operation):
"""Process multiple files, opening each with ExitStack."""
with ExitStack() as stack:
files = {}
try:
# Open all files and register cleanup
for filename in filenames:
f = open(filename, 'r')
stack.enter_context(f) # Register f.__exit__ for cleanup
files[filename] = f
# Process all files
for filename, f in files.items():
operation(filename, f.read())
except Exception as e:
print(f"Error during processing: {e}")
raise
# All files are automatically closed here, in reverse order
# Test with temporary files
with tempfile.TemporaryDirectory() as tmpdir:
filenames = []
for i in range(3):
fname = os.path.join(tmpdir, f"file_{i}.txt")
with open(fname, 'w') as f:
f.write(f"Content {i}")
filenames.append(fname)
def print_file(fname, content):
print(f"{os.path.basename(fname)}: {content}")
process_many_files(filenames, print_file)
# Output:
# file_0.txt: Content 0
# file_1.txt: Content 1
# file_2.txt: Content 2
Pattern 2: ExitStack with Custom Cleanup Functions
Register arbitrary cleanup functions for non-context-manager resources.
# Example 2: Custom cleanup with ExitStack
from contextlib import ExitStack
import tempfile
import os
def setup_complex_resource():
"""Setup a resource that requires custom cleanup."""
print("Acquiring complex resource")
return {"status": "active", "data": []}
def cleanup_complex_resource(resource):
"""Custom cleanup function."""
print(f"Cleaning up resource with status: {resource['status']}")
resource['status'] = 'cleaned'
with ExitStack() as stack:
# Acquire resource and register custom cleanup
resource = setup_complex_resource()
stack.callback(cleanup_complex_resource, resource)
# Use the resource
resource['data'].append("item1")
resource['data'].append("item2")
print(f"Resource data: {resource['data']}")
# Output:
# Acquiring complex resource
# Resource data: ['item1', 'item2']
# Cleaning up resource with status: active
Pattern 3: Chaining Context Managers at Runtime
Dynamically decide which context managers to enter.
# Example 3: Conditional context management with ExitStack
from contextlib import ExitStack, contextmanager
@contextmanager
def operation_logger(name, verbose=False):
"""Log operation start and end."""
if verbose:
print(f"Starting {name}")
try:
yield
finally:
if verbose:
print(f"Finished {name}")
@contextmanager
def profiling_timer(name, profile=False):
"""Optionally time an operation."""
import time
if profile:
start = time.time()
print(f"Starting profiling: {name}")
try:
yield
finally:
if profile:
elapsed = time.time() - start
print(f"Profiling result: {name} took {elapsed:.3f}s")
def run_operation(operation, name, verbose=False, profile=False):
"""Run operation with conditionally applied context managers."""
with ExitStack() as stack:
# Conditionally enter context managers
if verbose:
stack.enter_context(operation_logger(name, verbose=True))
if profile:
stack.enter_context(profiling_timer(name, profile=True))
# Run the operation
operation()
def my_operation():
import time
print("Doing work...")
time.sleep(0.1)
print("--- No logging or profiling ---")
run_operation(my_operation, "Operation A", verbose=False, profile=False)
print("\n--- With verbose logging ---")
run_operation(my_operation, "Operation B", verbose=True, profile=False)
print("\n--- With profiling ---")
run_operation(my_operation, "Operation C", verbose=False, profile=True)
# Output:
# --- No logging or profiling ---
# Doing work...
#
# --- With verbose logging ---
# Starting Operation B
# Doing work...
# Finished Operation B
#
# --- With profiling ---
# Starting profiling: Operation C
# Doing work...
# Profiling result: Operation C took 0.101s
Pattern 4: Resource Pool Context Manager
Manage a pool of reusable resources (connections, threads, etc.).
# Example 4: Simple connection pool with context manager
from contextlib import contextmanager
import threading
from queue import Queue
class ConnectionPool:
"""A simple connection pool implemented with ExitStack."""
def __init__(self, create_connection, max_size=5):
self.create_connection = create_connection
self.max_size = max_size
self.available = Queue(maxsize=max_size)
self.active = set()
self.lock = threading.Lock()
# Pre-populate the pool
for _ in range(max_size):
conn = create_connection()
self.available.put(conn)
@contextmanager
def get_connection(self):
"""Acquire a connection from the pool, release when done."""
conn = self.available.get()
with self.lock:
self.active.add(id(conn))
try:
yield conn
finally:
with self.lock:
self.active.discard(id(conn))
self.available.put(conn)
def create_fake_connection():
"""Create a fake database connection."""
return f"Connection_{id(id)}"
# Test the pool
pool = ConnectionPool(create_fake_connection, max_size=2)
def worker(worker_id):
for i in range(2):
with pool.get_connection() as conn:
print(f"Worker {worker_id}: using {conn}")
threads = [threading.Thread(target=worker, args=(i,)) for i in range(3)]
for t in threads:
t.start()
for t in threads:
t.join()
print("All workers completed")
# Output (order varies):
# Worker 0: using Connection_...
# Worker 1: using Connection_...
# Worker 0: using Connection_...
# Worker 2: using Connection_...
# Worker 1: using Connection_...
# Worker 2: using Connection_...
# All workers completed
Pattern 5: Multi-Context Composition
Stack multiple context managers for orthogonal concerns.
# Example 5: Composing context managers for multiple concerns
from contextlib import contextmanager, ExitStack
import time
@contextmanager
def profiling(name):
"""Profile execution time."""
start = time.time()
print(f"[PROFILING] Starting {name}")
try:
yield
finally:
elapsed = time.time() - start
print(f"[PROFILING] {name} took {elapsed:.3f}s")
@contextmanager
def error_recovery(operation_name):
"""Handle and log errors."""
try:
yield
except Exception as e:
print(f"[ERROR] {operation_name} failed: {e}")
raise
@contextmanager
def state_machine(operation_name):
"""Track state transitions."""
print(f"[STATE] {operation_name} transitioning to RUNNING")
try:
yield
finally:
print(f"[STATE] {operation_name} transitioning to IDLE")
def complex_operation():
"""An operation that needs multiple concerns."""
with ExitStack() as stack:
stack.enter_context(profiling("ComplexOp"))
stack.enter_context(error_recovery("ComplexOp"))
stack.enter_context(state_machine("ComplexOp"))
print("Executing complex operation...")
time.sleep(0.1)
print("Complex operation done")
complex_operation()
# Output:
# [PROFILING] Starting ComplexOp
# [ERROR] ComplexOp failed: ... (if error)
# [STATE] ComplexOp transitioning to RUNNING
# Executing complex operation...
# Complex operation done
# [STATE] ComplexOp transitioning to IDLE
# [PROFILING] ComplexOp took 0.10s
Pattern 6: Fault-Tolerant Cleanup with ExitStack
Ensure cleanup succeeds even if individual steps fail.
# Example 6: Resilient cleanup
from contextlib import ExitStack
class ResilientCleanup:
"""Cleanup multiple resources, continuing even if some fail."""
@staticmethod
def cleanup_with_fallback():
with ExitStack() as stack:
resources = []
# Register multiple cleanup steps
def cleanup_resource(name):
try:
print(f"Cleaning {name}")
if name == "problematic":
raise RuntimeError(f"Failed to clean {name}")
resources.append(f"{name} cleaned")
except Exception as e:
print(f"Warning: {e}")
stack.callback(cleanup_resource, "A")
stack.callback(cleanup_resource, "problematic")
stack.callback(cleanup_resource, "C")
print("All resources acquired")
print(f"Resources cleaned: {resources}")
ResilientCleanup.cleanup_with_fallback()
# Output:
# All resources acquired
# Cleaning C
# Cleaning problematic
# Warning: Failed to clean problematic
# Cleaning A
# Resources cleaned: ['C cleaned', 'A cleaned']
Pattern 7: Context Manager Decorator Pattern
Create a decorator that applies context manager logic to functions.
# Example 7: Decorator using context managers
from contextlib import contextmanager
import functools
import time
@contextmanager
def timing_context(func_name):
"""Context manager for timing function execution."""
start = time.time()
print(f"Timing {func_name}...")
try:
yield
finally:
elapsed = time.time() - start
print(f"{func_name} took {elapsed:.3f}s")
def timed(func):
"""Decorator that times function execution."""
@functools.wraps(func)
def wrapper(*args, **kwargs):
with timing_context(func.__name__):
return func(*args, **kwargs)
return wrapper
@timed
def slow_function():
"""A function that is timed by the decorator."""
time.sleep(0.1)
return "result"
result = slow_function()
print(f"Result: {result}")
# Output:
# Timing slow_function...
# slow_function took 0.102s
# Result: result
Pattern 8: State-Preserving Context Managers
Context managers that preserve and restore state across function calls.
# Example 8: Stack-based state preservation
from contextlib import contextmanager
class ConfigContext:
"""Manage configuration state with rollback."""
def __init__(self):
self.config = {}
self.state_stack = []
@contextmanager
def temporary_config(self, key, value):
"""Temporarily set config value, restore on exit."""
old_value = self.config.get(key)
self.config[key] = value
try:
yield self.config
finally:
if old_value is None:
self.config.pop(key, None)
else:
self.config[key] = old_value
config = ConfigContext()
config.config['debug'] = False
print(f"Initial: debug={config.config.get('debug')}")
with config.temporary_config('debug', True):
print(f"Inside context: debug={config.config.get('debug')}")
print(f"After context: debug={config.config.get('debug')}")
# Output:
# Initial: debug=False
# Inside context: debug=True
# After context: debug=False
Key Takeaways
ExitStackenables dynamic cleanup registration for resources unknown at write time.- Register context managers with
enter_context()or custom cleanup withcallback(). - Compose multiple context managers for orthogonal concerns: profiling, logging, error handling.
- Resource pools use context managers to safely acquire and release from a shared pool.
- Fault-tolerant cleanup continues even if individual cleanup steps fail.
- Context manager decorators apply resource management transparently to functions.
- State-preserving context managers implement rollback and restoration patterns.
Frequently Asked Questions
What is the difference between ExitStack.callback and enter_context?
enter_context() calls __enter__ on an existing context manager. callback() registers an arbitrary function to call on exit. Use enter_context() for context managers, callback() for custom cleanup.
Can I dynamically choose cleanup order with ExitStack?
No, ExitStack always exits in LIFO order (reverse of entry). If you need custom order, use multiple callback() calls or manually track cleanup order in a data structure.
How do I handle exceptions in ExitStack cleanup?
Exceptions in cleanup callbacks propagate after all cleanups run. To handle them, wrap your with block in try-except or use a callback that catches exceptions internally.
Can I nest ExitStack with other context managers?
Yes. You can use with ExitStack() as stack: inside another with block, or use enter_context() to register an ExitStack inside another ExitStack.
Is ExitStack reentrant?
Yes, but each entry creates a separate context. Nested uses are safe because each level has its own exit stack.