contextlib.contextmanager decorator
Overview
The `@contextmanager` decorator from `contextlib` lets you build a context manager using a generator function instead of a full class with `__enter__` and `__exit__`. It breaks the generator down into three phases: - The code before the `yield` statement acts as `__enter__` (setup) - The value yielded is bound to the `as` variable - The code after the `yield` acts as `__exit__` (teardown) from contextlib import contextmanager @contextmanager def no_yield(): return # Forgot to yield! with no_yield(): pass # RuntimeError: generator didn't yield A `@contextmanager` must `yield` exactly once. Failing to `yield` raises a `RuntimeError`.
It drastically reduces boilerplate. Instead of writing a full class with dunder methods, you write a single generator function, keeping setup and teardown logic localized in one place.
Where used: `FastAPI` database sessions, Patching in unit tests, Timing execution blocks
Why learn this
- Write cleaner, more readable context managers without class overhead.
- Understand the primary way `FastAPI` manages dependencies like database sessions.
Code walkthrough
from contextlib import contextmanager
@contextmanager
def log_error():
try:
yield
except Exception as e:
print(f'Logged: {e}')
with log_error():
raise ValueError('oops')Focus: try:\n yield
Aha moment
from contextlib import contextmanager
@contextmanager
def bad_cleanup():
yield
print('Cleaning')
try:
with bad_cleanup():
raise ValueError('Fail')
except ValueError:
print('Caught')Prediction: Is `'Cleaning'` printed before `'Caught'`?
Common guess: `Yes`
Because the `yield` is not wrapped in `try`/`finally`, the `ValueError` raised inside the block crashes the generator at the `yield` line. `'Cleaning'` is never reached.
Common mistakes
- Forgetting the try/finally block: If an exception happens inside the `with` block, it is re-raised at the `yield` statement. If you don't wrap the `yield` in a `try`/`finally` block, your teardown code will never run.
- Yielding multiple values: The generator must `yield` exactly once. Yielding multiple times, or not yielding at all, will raise a `RuntimeError`.
Glossary
- generator function
- A special function that can pause its execution and `yield` multiple values over time, e.g. `def f(): yield 1`.
- teardown
- The process of cleaning up resources, like closing a file or database connection, after they are no longer needed, e.g. `file.close()`.
Recall questions
- Which part of the generator function acts as the `__exit__` method?
- What happens if an exception occurs inside the `with` block?
- How do you ensure teardown code runs even if an exception occurs?
- What happens if a `@contextmanager` generator returns without yielding?
Understanding checks
What is printed to the console?
<p> text </p>
The code before `yield` prints `<p>`, then the `with` block runs printing `text`, and finally the code after `yield` prints `</p>`.
What is printed to the console?
Start Value End
The generator yields the string `'Value'`, which gets bound to `s`. It prints `'Start'`, then `'Value'`, then `'End'`.
Practice tasks
Add Safe Teardown
Given this context manager, modify it to guarantee that `'Disconnecting'` is printed even if the `with` block raises an exception.
Challenge
A Directory Switcher
Write a context manager using `@contextmanager` called `temp_dir` that takes a dictionary and a key. It should temporarily set the key to `'temp_value'`, `yield` the new dictionary, and then restore the original value of the key afterwards.
Continue learning
Previous: __enter__ and __exit__
Previous: Function Decorators
Previous: yield keyword and generator functions