Custom Context Managers

Writing your own context manager — as a class or a generator function — lets any setup/teardown pair in your code get the same guaranteed-cleanup treatment as open() and threading.Lock.

The context management protocol is open for any object to implement, which means you can write custom context managers tailored to your own resources and patterns: database transactions, temporary state changes, timing blocks, mock patches. This topic focuses on the practical craft of writing them well — both the class-based approach and the generator-based @contextmanager approach — and the details that separate a correct implementation from a subtly broken one.

Class-Based Custom Context Managers

A class-based context manager implements __enter__ and __exit__ directly. This approach is well suited when the context manager needs to hold meaningful state, expose other methods, or be reused as a normal object outside of a with block too.

__exit__ inspects exc_type to decide commit vs rollback — this is exactly the kind of exception-aware cleanup logic that's easy to express in __exit__ but awkward with a plain finally block.

class DatabaseTransaction:
    def __init__(self, connection):
        self.connection = connection

    def __enter__(self):
        self.connection.begin()
        return self.connection

    def __exit__(self, exc_type, exc_value, tb):
        if exc_type is None:
            self.connection.commit()
        else:
            self.connection.rollback()
        return False  # never suppress the exception

with DatabaseTransaction(conn) as c:
    c.execute("INSERT INTO users VALUES (?)", ("alice",))
# committed if no exception, rolled back otherwise

Generator-Based Custom Context Managers

The same transaction manager can be written more concisely with @contextmanager, when the logic is linear enough not to need a full class. This is often the more Pythonic choice for straightforward setup/teardown pairs.

try/except/else around yield mirrors the class version: except handles the failure path (rollback, then re-raise so the caller still sees the error), else handles the success path (commit only if no exception occurred).

from contextlib import contextmanager

@contextmanager
def db_transaction(connection):
    connection.begin()
    try:
        yield connection
    except Exception:
        connection.rollback()
        raise
    else:
        connection.commit()

with db_transaction(conn) as c:
    c.execute("INSERT INTO users VALUES (?)", ("alice",))

Class vs Generator: When to Choose Which

Prefer @contextmanager for simple, linear setup/teardown logic with no extra methods needed — it's shorter and reads top-to-bottom like a procedure. Prefer a class when the context manager needs to expose additional methods beyond enter/exit, hold state accessed by other parts of the program, support being subclassed, or when the object already exists as a class for other reasons and the protocol is just being added to it.

Both approaches produce an object satisfying the exact same protocol — code using with cannot tell which style was used to build it, so the choice is purely about what's more maintainable to write.

Return Value from __enter__

__enter__'s return value becomes whatever as name binds to — it does not have to be self. Returning something other than self is useful when the with block should interact with a different object than the context manager itself (e.g. the manager sets something up, but the block should work with the resource it produced).

The caller only ever interacts with the plain directory path string; the TemporaryDirectory instance itself, which tracks internal state for cleanup, stays hidden.

If a with block doesn't use as name at all, __enter__'s return value is simply discarded — this is fine and common when a context manager exists purely for its side effects (e.g. acquiring a lock).

class TemporaryDirectory:
    def __enter__(self):
        import tempfile
        self._path = tempfile.mkdtemp()
        return self._path  # returns a plain string, not self

    def __exit__(self, exc_type, exc_value, tb):
        import shutil
        shutil.rmtree(self._path)
        return False

with TemporaryDirectory() as path:
    print(path)  # path is a str, not a TemporaryDirectory instance

Exception Handling Inside __exit__: The exc_type Check

__exit__ receives (exc_type, exc_value, traceback). All three are None when the block completed without error — this is the standard way to branch on success vs failure inside __exit__, as shown in the transaction example above. Never assume an exception occurred just because __exit__ was called; it's called unconditionally.

Checking exc_type is None (not truthiness of some other flag) is the correct, idiomatic way to detect whether the with block raised.

class Logger:
    def __enter__(self):
        print("starting operation")
        return self

    def __exit__(self, exc_type, exc_value, tb):
        if exc_type is None:
            print("operation succeeded")
        else:
            print(f"operation failed: {exc_type.__name__}: {exc_value}")
        return False

Suppressing Only Specific Exceptions

Returning True unconditionally from __exit__ is a common bug — it suppresses every exception, including ones you never intended to swallow. Always check exc_type (or use issubclass(exc_type, SomeError)) before deciding to suppress, and return False for anything you don't specifically intend to handle.

Only TimeoutError on the first attempt is suppressed; any other exception type, or a TimeoutError on a later attempt, correctly propagates instead of being silently discarded.

class RetryOnce:
    def __init__(self):
        self.attempts = 0

    def __enter__(self):
        self.attempts += 1
        return self

    def __exit__(self, exc_type, exc_value, tb):
        if exc_type is TimeoutError and self.attempts == 1:
            print("suppressing first timeout, caller should retry")
            return True
        return False  # everything else propagates normally

Re-entrancy and Reuse

A class-based context manager instance is not automatically safe to use in more than one with block at a time (or nested with itself) unless it's specifically designed to be — instance state set in __enter__ (like a start time or a file handle) can be overwritten by a second entry before the first has exited. If reuse across multiple with blocks is needed, either make the object a factory that creates fresh state each time, or explicitly design for re-entrancy.

Because Timer stores its start time as self._start, the inner __enter__ overwrites the outer's start time, breaking the outer block's elapsed-time calculation — a fresh Timer() per with block avoids this entirely.

timer = Timer()
with timer:
    with timer:  # re-enters the SAME instance -- likely to misbehave
        pass

Generator-Based Pitfall: Multiple yields or No yield

A @contextmanager generator function must yield exactly once. Zero yields, or more than one reachable yield on different code paths, raises a RuntimeError when the context manager is used, because the underlying machinery expects precisely one 'enter' point and one 'resume after block' point.

The single-yield rule exists because @contextmanager's machinery resumes the generator exactly once after the with block ends — a second yield has nowhere defined to go.

from contextlib import contextmanager

@contextmanager
def broken(flag: bool):
    if flag:
        yield "a"
    else:
        yield "b"
    yield "c"  # BUG: unreachable in practice, but illustrates the single-yield rule

# A correct version has exactly one yield on every path:
@contextmanager
def fixed(flag: bool):
    value = "a" if flag else "b"
    yield value

Key Takeaways

Class-based context managers implement __enter__/__exit__ directly; generator-based ones use @contextmanager with exactly one yield.

__enter__'s return value becomes whatever as name binds to — it can be self, a different object, or nothing meaningful if as is unused.

__exit__ receives (exc_type, exc_value, traceback), all None on success — check exc_type is None to branch on success vs failure.

Only return True from __exit__ (or omit re-raising after a caught exception in a generator) for exception types you deliberately intend to suppress.

A @contextmanager generator must yield exactly once; class-based managers are not automatically safe to reuse or nest with themselves unless designed for it.

def __exit__(self, exc_type, exc_value, tb): cleanup() return True # suppresses everything, even bugs def __exit__(self, exc_type, exc_value, tb): cleanup() return exc_type is SpecificExpectedError

Unconditionally returning True from __exit__ suppresses every exception, including unrelated bugs raised inside the with block; only suppress the specific exception type(s) you actually intend to handle.

@contextmanager def cm(): setup() yield yield # second yield -- RuntimeError at use time

@contextmanager def cm(): setup() try: yield finally: teardown()

A @contextmanager generator must yield exactly once; a reachable second yield causes a RuntimeError when the context manager is actually used in a with statement.

timer = Timer() with timer: with timer: pass # reuses the same instance, state gets overwritten with Timer() as t1: with Timer() as t2: pass # each with block gets its own fresh instance

Reusing one context manager instance across nested or sequential with blocks can corrupt its internal state unless the class was specifically designed to be re-entrant; prefer a fresh instance per with block.

What determines the value bound by as name in a with statement?

Whatever __enter__ returns (for a class-based context manager) or whatever is yielded (for a @contextmanager generator) — it does not have to be self, and can be any object, or nothing meaningful if as isn't used.

How do you correctly detect, inside __exit__, whether the with block raised an exception?

Check whether exc_type is None: it's None when the block completed successfully, and set to the exception's class when the block raised. Never assume an exception occurred just because __exit__ was called — it always runs.

Why is return True unconditionally from __exit__ usually a bug?

It suppresses every exception raised in the with block, including ones that were never intended to be swallowed (like unrelated bugs); __exit__ should check exc_type and return True only for the specific exception type(s) it deliberately wants to handle.

What happens if a @contextmanager-decorated generator function yields twice on a reachable path?

It raises a RuntimeError when the context manager is actually used, because the underlying machinery only expects to resume the generator once after the with block ends — exactly one yield is required per use.

In the generator-based transaction example, why does the except clause re-raise after calling rollback()?

So the caller of the with block still sees the original exception (the failure isn't silently hidden); rollback() undoes the transaction's effects, but the exception itself should still propagate unless suppressing it is deliberately intended.

Why might reusing a single class-based context manager instance across two separate with blocks cause bugs?

Instance attributes set in __enter__ (like a start time or file handle) are shared state; entering the same instance a second time before or after the first can overwrite that state, corrupting whichever with block relies on it — unless the class is specifically designed to be re-entrant.

How would you write a context manager whose __enter__ returns something other than self, and why might you want that?

Have __enter__ return the specific resource the caller needs (e.g. return self._path instead of return self), which is useful when the with block should interact directly with the produced resource rather than the manager object that tracks internal setup/teardown state.

Write a class-based context manager Suppress(*exception_types) equivalent to contextlib.suppress, that ignores only the given exception types.

class Suppress: def __init__(self, *exception_types: type[BaseException]) -> None: self.exception_types = exception_types def __enter__(self): return self def __exit__(self, exc_type, exc_value, tb) -> bool: return exc_type is not None and issubclass(exc_type, self.exception_types)

__exit__ suppresses the exception (returns True) only when exc_type is set and is a subclass of one of the stored exception_types; any other exception, or no exception, correctly returns False/propagates.

What is printed, and does 'D' get reached?

Prints A, B, C, D. The ValueError is thrown into the generator at yield, caught by except ValueError (prints B), then execution continues to print C; since the exception was not re-raised, it's suppressed, so the with block completes normally and D is reached.

Write a generator-based context manager atomic_write(path) that writes to a temporary file and only replaces the target file if no exception occurred, using try/except/else.

from contextlib import contextmanager from pathlib import Path import os import tempfile @contextmanager def atomic_write(path: str): dir_ = Path(path).parent fd, tmp_path = tempfile.mkstemp(dir=dir_) try: with os.fdopen(fd, "w") as f: yield f except Exception: os.remove(tmp_path) raise else: os.replace(tmp_path, path)

Writes go to a temporary file; on success (else branch) the temp file atomically replaces the target, on failure the temp file is removed and the exception re-raised, so the target file is only ever updated by a fully successful write.

from contextlib import contextmanager

@contextmanager
def cm():
    print("A")
    try:
        yield
    except ValueError:
        print("B")
    print("C")

with cm():
    raise ValueError("x")
print("D")