Write a Context Manager with a Generator

A test needs to switch settings.mode to "test" for a few lines, then put the old value back, even if the test fails halfway. You could write a class with __enter__ and __exit__ methods. The contextmanager decorator from contextlib lets you write the same context manager as one generator function:

from collections.abc import Iterator
from contextlib import contextmanager


class Settings:
    def __init__(self) -> None:
        self.mode = "production"


@contextmanager
def temp_mode(settings: Settings, mode: str) -> Iterator[None]:
    original = settings.mode
    settings.mode = mode
    try:
        yield
    finally:
        settings.mode = original


settings = Settings()
with temp_mode(settings, "test"):
    print(settings.mode)  # test
print(settings.mode)      # production

How a generator becomes a context manager

You met yield in "Generators and One-Shot Iterators". Here the decorator advances the generator for you:

1. Calling temp_mode(...) returns a context manager. As with any generator function, none of the body runs yet.
2. When the with block starts, the decorator calls next() on the generator. The body runs up to yield and pauses there. The value you yield goes to the as target. This one yields nothing, so the annotation is Iterator[None].
3. When the block ends, the generator resumes after yield and runs the cleanup. If the block raised an exception, Python raises that exception inside the generator, at the yield line.

Annotate the generator function as Iterator[T], where T is the type you yield. mypy then knows that the decorated function returns a context manager that gives you a T.

Put the cleanup in finally

A common mistake is to write the cleanup with no try:

@contextmanager
def broken_temp_mode(settings: Settings, mode: str) -> Iterator[None]:
    original = settings.mode
    settings.mode = mode
    yield
    settings.mode = original  # skipped when the block raises


try:
    with broken_temp_mode(settings, "test"):
        raise ValueError("test failed")
except ValueError:
    pass
print(settings.mode)  # test, so the old value is lost

If the with block raises, the exception comes out of yield, and the line after it never runs. Wrap the yield in try and put the cleanup in finally, as temp_mode does. The finally block runs whether the block finishes or raises, and the exception still reaches the caller afterwards.

Handle any attribute, even a missing one

The built-in functions getattr, setattr and delattr read, set and delete an attribute whose name is a string. With them, one context manager works for any attribute.

If the attribute did not exist before, putting it back means deleting it. getattr(obj, name, None) cannot tell you whether the attribute was missing or held None. Use a sentinel instead, a unique object that no caller can ever pass in:

_MISSING = object()


@contextmanager
def temp_attr(obj: object, name: str, value: object) -> Iterator[None]:
    original = getattr(obj, name, _MISSING)
    setattr(obj, name, value)
    try:
        yield
    finally:
        if original is _MISSING:
            delattr(obj, name)
        else:
            setattr(obj, name, original)


config = Settings()
with temp_attr(config, "debug", True):
    print(getattr(config, "debug"))  # True
print(hasattr(config, "debug"))      # False

Nest calls, but enter each one once

Nested calls work with no extra code. The inner call saves the value the outer call set, and restores it when the inner block ends.

with temp_attr(config, "mode", "outer"):
    with temp_attr(config, "mode", "inner"):
        print(config.mode)  # inner
    print(config.mode)      # outer
print(config.mode)          # production

Each call creates a new context manager, and each one works only once, because a generator runs only once. Entering the same object a second time raises an exception. Python does not promise which exception class, so do not write code that depends on it.

Going further

• Write a class with __enter__ and __exit__ instead when the object has other methods too, or when code must enter it more than once.
• If the cleanup itself raises, its exception replaces the one from the block. Python keeps the first one in the new exception's __context__.
• A context manager made with @contextmanager also works as a decorator. From Python 3.15, decorating a generator function with it keeps the context open while the caller loops over the generator. On 3.12 to 3.14 it covers only the call that creates the generator.
• See the contextlib documentation.