Python’s contextlib module lets you build custom context managers with minimal boilerplate, but its real power is in how it exposes the underlying generator protocol for control flow.
Let’s watch a simple contextlib.contextmanager in action. Imagine we need to temporarily change the current working directory, perform some operations, and then restore the original directory.
import os
from contextlib import contextmanager
@contextmanager
def change_dir(new_dir):
original_dir = os.getcwd()
os.chdir(new_dir)
print(f"Entered directory: {os.getcwd()}")
try:
yield # This is where the 'with' block's code runs
finally:
os.chdir(original_dir)
print(f"Restored to directory: {os.getcwd()}")
# Create a dummy directory for demonstration
os.makedirs("temp_subdir", exist_ok=True)
print(f"Current directory before: {os.getcwd()}")
with change_dir("temp_subdir"):
print("Inside the 'with' block.")
# Imagine doing some file operations here
with open("test_file.txt", "w") as f:
f.write("Hello from temp_subdir!")
print(f"Current directory after: {os.getcwd()}")
When you run this, you’ll see:
Current directory before: /path/to/your/script
Entered directory: /path/to/your/script/temp_subdir
Inside the 'with' block.
Restored to directory: /path/to/your/script
Current directory after: /path/to/your/script
The @contextmanager decorator transforms a generator function into a context manager. The code before the yield statement runs when the with block is entered. The yield statement itself is the point where control is transferred to the code inside the with block. Crucially, anything after the yield (within a finally block, as is best practice for cleanup) runs when the with block is exited, regardless of whether an exception occurred.
This pattern is incredibly useful for managing resources that need setup and teardown. Think about database connections, file locks, or temporary configuration changes. The contextlib module provides several helpers to build more complex scenarios. closing wraps an object with a close() method, suppress lets you ignore specific exceptions within a with block, and redirect_stdout/redirect_stderr are handy for capturing output.
The core mechanism powering contextlib is the generator’s ability to pause execution. When yield is encountered, the generator’s state is saved, and control returns to the caller. The with statement’s __enter__ method effectively calls next() on the generator, running it up to the yield. When the with block finishes, the __exit__ method is called, which resumes the generator by calling next() again, allowing the code after yield to execute. If an exception occurred in the with block, it’s passed to the generator’s throw() method, giving your cleanup code a chance to handle it.
A common pitfall is forgetting the finally block. If an exception occurs before the yield (during setup), the code after yield in the generator will not run, and your resource might be left in an inconsistent state. Always wrap your yield statement in a try...finally block within your @contextmanager function to guarantee cleanup.
The contextlib.ExitStack is a powerful tool for managing multiple nested context managers that might be determined dynamically, allowing you to enter and exit them in a single with statement, ensuring all are properly cleaned up even if exceptions occur.