Chapter 15

Context Managers

Whatever you acquire must be released, on every path out, including the one an exception takes. A context manager pairs acquisition with release so cleanup cannot be skipped.

The with statement, introduced in Control Flow, marks out a span of execution. It runs setup before a block and cleanup after it, even if the block raises an exception. That guarantee is far more reliable than the __del__() approach. Writing your own context managers shows how with works.

The payoff is a borrower’s contract two lines long:

with pool.lease() as conn:
    conn.query("SELECT name FROM users")

The connection returns to the pool on every path out of that block, including the exception path, and the borrower writes nothing to arrange it.

A Basic Context Manager

The simplest way to write a context manager is a generator function with a single yield. The contextlib.contextmanager decorator turns that function into a context manager. The yield here works the way it does in a pytest fixture that yields its value. Everything before the yield is setup, and everything after it is teardown. Iterators covers generators in full. A context manager needs nothing beyond that shape.

# trace_gen.py
from collections.abc import Iterator
from contextlib import contextmanager

@contextmanager
def trace(name: str) -> Iterator[str]:
    print(f"enter {name}")  # Setup
    try:
        yield name  # The block runs here
    finally:
        print(f"exit {name}")  # Cleanup

if __name__ == "__main__":
    with trace("A") as t:
        print(f"inside {t}")
#: enter A
#: inside A
#: exit A

with trace("A") as t: runs the body of trace() up to the yield, printing enter A. as binds the yielded value, so t is "A". The block under the with then runs. When it finishes, trace() resumes just after the yield and prints exit A.

The manager object trace("A") returns is single-use. Its generator runs once, so reusing the same object in a second with fails with a message that names nothing useful: AttributeError: '_GeneratorContextManager' object has no attribute 'args'. Construct a fresh manager for each with statement. Nor does a loop around the yield make the manager reusable. @contextmanager allows one yield, and a generator that reaches a second one makes the manager raise RuntimeError: generator didn't stop when the block ends.

The finally makes the cleanup dependable. An exception raised in the block appears at the yield, and finally still runs the cleanup before the exception propagates.

Leaving the try/finally out is the common mistake:

# no_finally.py
from collections.abc import Iterator
from contextlib import contextmanager

@contextmanager
def careless(name: str) -> Iterator[str]:
    print(f"enter {name}")
    yield name
    print(f"exit {name}")

try:
    with careless("A"):
        raise ValueError("boom")
except ValueError as error:
    print("caught:", error)
#: enter A
#: caught: boom

Without the try/finally, @contextmanager’s __exit__() calls the generator’s throw() method, which resumes the generator by raising the block’s exception at the yield, so the code after the yield, exit A included, never runs. Nothing warns you. The generator silently skips the cleanup on the one path where it matters most. Wrap the yield in try/finally in every @contextmanager generator.

The Protocol

How does with know what to run? It knows nothing about generators or @contextmanager. A context manager is any object that implements two methods: __enter__(), which runs at the start of the block, and __exit__(), which runs at the end. @contextmanager manufactures such an object from a generator function.

Every hand-written context manager class in this chapter keeps __init__() in longhand rather than becoming a @dataclass, so every line between the class statement and the two protocol methods reads as ordinary Python:

# trace_cm.py
from typing import Self

class Trace:
    def __init__(self, name: str) -> None:
        self.name = name

    def __enter__(self) -> Self:
        print(f"enter {self.name}")
        return self

    def __exit__(self, exc_type: type[BaseException] | None,
                 exc: object, tb: object) -> None:
        print(f"exit {self.name}")

if __name__ == "__main__":
    with Trace("A") as t:
        print(f"inside {t.name}")
#: enter A
#: inside A
#: exit A

Python runs with Trace("A") as t: in five steps:

Evaluate Trace("A") to produce a manager object.
Call the manager’s __enter__().
Bind __enter__()’s return value to t.
Run the block.
Call the manager’s __exit__(), no matter how the block finished.

__enter__() returns the object that as binds, often self. The return annotation Self declares an instance of the class on which you call the method, so it adapts to subclasses.

The three parameters of __exit__() describe the exception that ended the block. The block in trace_cm.py ends normally, so all three are None.

In generator terms, __enter__() is the portion before the yield. __exit__() is the portion after it.

Trace is also reusable. The same instance can appear in a second with, because __enter__() just runs again. A generator manager cannot appear in a second with. That is the single-use caution from trace_gen.py. A class manager that stores per-with state stays reusable only if __enter__() resets that state.

The generator form is usually the clearest choice. Use a class when the manager needs to hold methods or state beyond a single setup and teardown.

Guaranteed Cleanup

Here, __exit__() runs when the block raises an exception, but before the exception propagates:

# exit_on_error.py
from trace_cm import Trace

try:
    with Trace("A"):
        raise ValueError("boom")
except ValueError as error:
    print("caught:", error)
#: enter A
#: exit A
#: caught: boom

exit A prints before caught, so the cleanup runs on the exception path. That ordering is the same guarantee a try/finally gives, packaged as a reusable object.

The guarantee covers the block, not the setup:

# enter_fails.py

class Fragile:
    def __enter__(self) -> None:
        print("enter fails")
        raise RuntimeError("no resource")

    def __exit__(self, *exc: object) -> None:
        print("exit runs")

try:
    with Fragile():
        print("body")
except RuntimeError as error:
    print("caught:", error)
#: enter fails
#: caught: no resource

exit runs is missing from the output, because Python registers the cleanup when __enter__() returns, and this __enter__() raises an exception instead. An __enter__() that acquires several things must clean up its own partial work before it raises an exception. ExitStack, later in this chapter, is the standard tool for that. When a later entry fails, ExitStack unwinds the managers it entered earlier. A with naming several managers applies the same rule per manager. The ones that entered still exit, and the failing one alone gets no __exit__() call.

Fragile.__exit__(self, *exc: object) collects the three arguments into a tuple the method ignores, the shorter form for a cleanup that ignores why the block ended.

The guarantee has a matching gap on the other side. Cleanup can fail. When __exit__() raises an exception, that new exception replaces the block’s original one, and the original survives only as the new exception’s __context__:

# exit_masks.py

class Careless:
    def __enter__(self) -> None:
        return None

    def __exit__(self, *exc: object) -> None:
        raise ValueError("cleanup error")

try:
    with Careless():
        raise KeyError("original")
except ValueError as error:
    print("caught:", repr(error))
    print("context:", repr(error.__context__))
#: caught: ValueError('cleanup error')
#: context: KeyError('original')

KeyError('original') never reaches the except. Careless.__exit__() raises a ValueError before with can propagate the KeyError, so the caller sees the cleanup’s failure instead of the block’s. Python keeps original. It becomes error.__context__, the same chaining a nested except produces. A broken cleanup path that hides the real failure is one of the most common context-manager bugs in practice, so write __exit__() methods that fail only for reasons worse than the original exception.

The `exit()` Arguments

__exit__(self, exc_type, exc, tb) receives the details of an exception raised in the block. When the block finishes normally, all three are None. When it raises an exception, they hold the exception’s class, its instance, and its traceback object. Trace.__exit__() in trace_cm.py types exc and tb as object, the most general type, since it ignores both. A manager that reads them types exc as BaseException | None, as expected_one below does, and tb as TracebackType | None, with TracebackType imported from types.

The return value decides that exception’s fate. A falsy value lets it propagate. Falsy includes the implicit None of a method with no return, so the exception propagates by default. A truthy value suppresses it: the with statement swallows the exception, and execution continues after the block.

Any truthy value suppresses, not True alone. An __exit__() that returns the result of its last cleanup call, such as a count or a status string, swallows every exception its block raises. Annotate __exit__() as -> None, the way Trace does, unless suppressing is the manager’s job. The type checker then rejects any value the method returns.

A generator manager suppresses by catching:

# suppress_in_generator.py
from collections.abc import Iterator
from contextlib import contextmanager

@contextmanager
def ignoring(kind: type[BaseException]) -> Iterator[None]:
    try:
        yield
    except kind as error:
        print(f"swallowed {error!r}")

with ignoring(ZeroDivisionError):
    print("before")
    1 / 0
    print("after")
#: before
#: swallowed ZeroDivisionError('division by zero')
print("survived")
#: survived

A generator manager has no return value to set. The exception arrives at the yield, so catching it there and swallowing it is the equivalent of __exit__() returning True. Re-raising it, or omitting the except clause, is the equivalent of returning a falsy value.

The standard library provides this behavior ready-made, as contextlib.suppress:

# suppress_exceptions.py
from contextlib import suppress

with suppress(ZeroDivisionError):
    print("before")
    1 / 0
    print("after")  # Never runs
#: before
print("survived")
#: survived

suppress is a class named like a function because you use it like one.

A class manager suppresses through the return value of __exit__(). expected_one returns True for one exception type, after printing the exception it caught:

# expected_one.py

class expected_one:
    def __init__(self, kind: type[BaseException]) -> None:
        self.kind = kind

    def __enter__(self) -> None:
        return None

    def __exit__(self, exc_type: type[BaseException] | None,
                 exc: BaseException | None,
                 tb: object) -> bool:
        if (exc_type is not None
            and issubclass(exc_type, self.kind)):
            print(f"{exc!r}")
            return True
        return False

with expected_one(ZeroDivisionError):
    print("before")
    1 / 0
    # Never runs: the error jumps to __exit__
    print("after")
#: before
#: ZeroDivisionError('division by zero')
print("survived")
#: survived

__exit__() receives exc_type: type[BaseException] | None because Python passes it the raised exception’s class, or None when the block finishes cleanly. type[...] means the class, such as ZeroDivisionError, not an instance of it. issubclass(cls, classinfo) returns True if cls is classinfo or a subclass of it, so a ZeroDivisionError still matches expected_one(ArithmeticError). exc!r prints the exception’s repr(), which includes both its type and its arguments, not just exc_type.__name__. __enter__() returns None because this manager has nothing to hand to as. An as clause still works, but it binds None.

The `expected` Manager

expected is a fuller version of expected_one. It takes several types at once, and with no argument it catches everything. Its name says how the book uses it. The block is expected to raise an exception, and the manager shows what the block raised. It is useful enough to reuse elsewhere in the book, so it lives in utils/, where any chapter can import it:

# utils/exceptions.py
import textwrap
from collections.abc import Awaitable, Callable
from typing import Final

WIDTH: Final[int] = 57
ALL = sentinel("ALL")
type Types = (type[BaseException]
              | tuple[type[BaseException], ...])

def report(e: BaseException) -> None:
    line = f"[{type(e).__name__}] {e}"
    print(textwrap.fill(line, WIDTH))

class expected:
    def __init__(self, types: Types | ALL = ALL) -> None:
        self.types = types

    def __enter__(self) -> None:
        return None

    def __exit__(self, exc_type: type[BaseException] | None,
                 exc: BaseException | None,
                 tb: object) -> bool:
        if exc_type is None or exc is None:
            return False
        if self.types is not ALL:
            if not issubclass(exc_type, self.types):
                return False
        report(exc)
        return True

def expect[**P](
    types: Types, fn: Callable[P, object],
    /, *args: P.args, **kwargs: P.kwargs
) -> None:
    try:
        fn(*args, **kwargs)
    except types as e:
        report(e)
        return
    raise AssertionError("no exception raised")

async def aexpect[**P](
    types: Types, fn: Callable[P, Awaitable[object]],
    /, *args: P.args, **kwargs: P.kwargs
) -> None:
    try:
        await fn(*args, **kwargs)
    except types as e:
        report(e)
        return
    raise AssertionError("no exception raised")

expected adds two things to expected_one. The first is the tuple form. issubclass() accepts a tuple of classes as its second argument, matching if cls is a subclass of any one of them, so expected((ZeroDivisionError, TypeError)) covers several types in one manager. The Types alias names that “one class or a tuple of classes” shape once instead of writing it out at every use.

The second is the default. The constructor’s types parameter defaults to the ALL sentinel, which makes expected() with no argument catch everything. self.types is not ALL narrows self.types from Types | ALL down to Types, and the earlier if exc_type is None or exc is None: return False narrows exc_type to a class and exc to an exception, so issubclass(exc_type, self.types) type-checks.

__exit__() also hands the exception to report(), which prints it as [Type] message through textwrap.fill() at WIDTH, so a long message wraps instead of overrunning the listing width. expect(), in the same file, prints through report() too, so every exception the book demonstrates through expected or expect() has one form.

suppress treats the no-argument call the opposite way. suppress() suppresses nothing, because a raised exception has no listed type to match. An expected() that catches everything also catches KeyboardInterrupt and SystemExit, so name the types you expect unless you want a block that nothing escapes.

# demo_exceptions.py
from exceptions import expected

with expected(ZeroDivisionError):
    print("before")
    1 / 0
    # Never runs: the error jumps to __exit__
    print("after")
#: before
#: [ZeroDivisionError] division by zero
print("survived")
#: survived

with expected():  # No argument means ALL
    print("before")
    raise KeyError("anything")
#: before
#: [KeyError] 'anything'
print("survived")
#: survived

with expected() as x:  # [1]
    print(f"{x = }")
#: x = None

The 1 / 0 raises an exception, __exit__() prints the exception it caught, then returns True, and the with statement suppresses the error so survived still prints.

In the with at [1], x receives the return value of __enter__(), which for expected() is None.

The `expect()` Function

Many listings in this book call something to show the exception it raises. expect() is the function form of that demonstration. It names the types it expects, calls fn with the remaining arguments, and prints the exception as [Type] message, wrapped by textwrap.fill() when that line is wider than WIDTH. WIDTH is 57 because this book’s listings are 60 columns wide and a #: output line spends three of them on its prefix. An exception of another type propagates, and a call that raises nothing fails with an AssertionError, so expect() flags a demo that stops failing instead of quietly printing nothing.

The / makes types and fn positional-only, so every keyword argument goes to fn. **P ties *args and **kwargs to fn’s own signature, so the type checker checks the forwarded arguments as if you had called fn:

# demo_expect.py
import json
from exceptions import expect

def parse(text: str, *, base: int = 10) -> int:
    return int(text, base)

expect(ValueError, parse, "ff")
#: [ValueError] invalid literal for int() with base 10: 'ff'
expect((ValueError, TypeError), parse, "ff", base=1)  # [1]
#: [ValueError] int() base must be >= 2 and <= 36, or 0
expect(json.JSONDecodeError, json.loads, "{bad")  # [2]
#: [JSONDecodeError] Expecting property name enclosed in
#: double quotes: line 1 column 2 (char 1)

The call at [1] names two types in a tuple and forwards base=1 as a keyword, which parse() passes on to int(). The message from [2] is too long for one line, so it wraps.

aexpect() is the async form. It awaits the call instead of making it, for a coroutine function whose failure is the demonstration.

Where a demonstration needs several statements or an assignment in the guarded block, expected remains the right tool. expect() covers the common case of one call.

Context Manager as Decorator

A context manager brackets a block of statements, with setup before and cleanup after. A typical decorator brackets a function call the same way. contextlib.ContextDecorator connects the two. A subclass works both as a context manager and as a decorator. Every manager @contextmanager produces inherits from ContextDecorator, so banner works as a decorator, although ContextDecorator is absent from context_decorator.py:

# context_decorator.py
from collections.abc import Iterator
from contextlib import contextmanager

@contextmanager
def banner(title: str) -> Iterator[None]:
    print(f"=== {title} ===")
    try:
        yield
    finally:
        print(f"=== {title} ends ===")

@banner("report")
def report() -> None:
    print("quarterly numbers")

if __name__ == "__main__":
    report()
    with banner("meeting"):
        print("agenda")
#: === report ===
#: quarterly numbers
#: === report ends ===
#: === meeting ===
#: agenda
#: === meeting ends ===

banner works both as a decorator for report() and in a with in __main__. The parentheses in @banner("report") matter. The call constructs the manager, which then decorates the function. Each call of the decorated function builds a fresh manager, so report() can run any number of times, each with its own enter and exit. The single-use caution from trace_gen.py still holds for the manager object you name in a with. ContextDecorator applies functools.wraps, so report() keeps its name and docstring.

# banner_cm.py
from contextlib import ContextDecorator

class banner(ContextDecorator):
    def __init__(self, title: str) -> None:
        self.title = title

    def __enter__(self) -> None:
        print(f"=== {self.title} ===")

    def __exit__(self, *exc: object) -> None:
        print(f"=== {self.title} ends ===")

@banner("report")
def report() -> None:
    print("quarterly numbers")

if __name__ == "__main__":
    report()
    with banner("meeting"):
        print("agenda")
#: === report ===
#: quarterly numbers
#: === report ends ===
#: === meeting ===
#: agenda
#: === meeting ends ===

Like suppress and expected, the class version of banner uses a lowercase name because you use it like a function. Its __exit__() ignores why the block ended, so it takes the *exc: object form and returns None, which lets an exception from report() propagate. Unlike the generator form, the class form re-enters the same instance on every call to report(), so every call shares any state the instance holds.

banner’s wrapper comes from ContextDecorator (as the base class in banner_cm.py, or by way of @contextmanager in context_decorator.py), and that wrapper always calls the function once, unchanged, with setup before it and cleanup after. So banner sees neither the arguments nor the return value of report(), and it cannot skip the call. A hand-written decorator can do all three, because it defines its own wrapper function. repeat forwards *args and **kwargs and returns the wrapped function’s result, while hijack returns a replacement that runs instead of the original function, skipping its call. banner offers one definition instead, usable both as a with block and as a @ decorator. Use ContextDecorator when setup and cleanup should be identical on every call.

The wrapper treats a generator function differently. Before Python 3.15, it wrapped only the call that creates the generator object, so the context entered and exited before the generator’s body ran, and the manager closed any resource it opened before the body could use it. On 3.15 the context opens at the first next() and closes after the generator’s last yield:

# decorated_generator.py
from collections.abc import Iterator
from contextlib import contextmanager

@contextmanager
def banner(title: str) -> Iterator[None]:
    print(f"=== {title} ===")
    try:
        yield
    finally:
        print(f"=== {title} ends ===")

@banner("rows")
def rows() -> Iterator[int]:
    for n in range(3):
        print(f"yield {n}")
        yield n

gen = rows()
print("created")
#: created
for row in gen:
    print(f"got {row}")
#: === rows ===
#: yield 0
#: got 0
#: yield 1
#: got 1
#: yield 2
#: got 2
#: === rows ends ===
print("done")
#: done

created prints before the opening banner, because calling rows() builds the generator object. The banner opens when the for loop asks for the first value, and it stays open each time the body resumes after a yield. Before 3.15, both banner lines printed before created. Code written for those versions leaves the decorator off and puts a with inside the generator around its body, which works on every version.

Combining Context Managers

A single with can include several managers, separated by commas. They enter left to right and exit in reverse:

# multiple.py
from collections.abc import Iterator
from contextlib import contextmanager

@contextmanager
def tag(name: str) -> Iterator[str]:
    print(f"<{name}>")
    try:
        yield name
    finally:
        print(f"</{name}>")

with tag("ul") as outer, tag("li") as inner:
    print(f"  {outer} then {inner}")
#: <ul>
#: <li>
#:   ul then li
#: </li>
#: </ul>

When the managers do not fit on one line, parentheses group them without changing the behavior:

with (tag("ul") as outer,
      tag("li") as inner):
    ...

When you do not know the number of managers until runtime, contextlib.ExitStack holds a dynamic set of managers and unwinds them in reverse when the block ends:

# exit_stack.py
from collections.abc import Iterator
from contextlib import ExitStack, contextmanager

@contextmanager
def tag(name: str) -> Iterator[str]:
    print(f"open {name}")
    try:
        yield name
    finally:
        print(f"close {name}")

def wrap(names: list[str]) -> None:
    with ExitStack() as stack:
        open_tags = [
            stack.enter_context(tag(n)) for n in names]
        print("using", open_tags)

wrap(["a", "b"])
#: open a
#: open b
#: using ['a', 'b']
#: close b
#: close a
wrap(["a", "b", "c"])
#: open a
#: open b
#: open c
#: using ['a', 'b', 'c']
#: close c
#: close b
#: close a

wrap() finds out how many managers to enter when it runs, and a comma-separated with cannot express that.

Every entry in wrap() succeeds, so wrap() leaves ExitStack’s other guarantee untested. When a later entry fails, the stack unwinds the managers that entered earlier. In exit_stack_fails.py the third manager fails to enter. The first two unwind, and the third’s cleanup does not run:

# exit_stack_fails.py
from collections.abc import Iterator
from contextlib import ExitStack, contextmanager
from exceptions import expected

@contextmanager
def tag(name: str, fail: bool = False) -> Iterator[str]:
    print(f"open {name}")
    if fail:
        raise RuntimeError(f"{name} failed to open")
    try:
        yield name
    finally:
        print(f"close {name}")

with expected(RuntimeError):
    with ExitStack() as stack:
        stack.enter_context(tag("a"))
        stack.enter_context(tag("b"))
        stack.enter_context(tag("c", fail=True))
#: open a
#: open b
#: open c
#: close b
#: close a
#: [RuntimeError] c failed to open

c gets no close c line, because its __enter__() raises a RuntimeError before ExitStack can register it. a and b entered first, so both unwind in reverse, the same rule a comma-separated with follows.

The `contextlib` Toolkit

The contextlib module provides ready-made managers. Choose these before writing __enter__() and __exit__() by hand.

suppress(*exceptions) ignores the listed exceptions, covering the case the expected class handles, without expected’s printing or its catch-everything default.
closing(obj) calls obj.close() on exit, for objects that have close() but are not context managers.
ExitStack manages a dynamic or conditional set of managers.
ContextDecorator lets a context manager double as a decorator, as banner shows.
nullcontext(value) is a do-nothing manager whose __enter__() returns value, useful when a with is optional and one code path should cover it.

A function might write to a path it opens, to a stream the caller hands it, or to standard output by default. When it finishes, the function should close the file it opened and leave the caller’s stream and standard output open. nullcontext lets a single with block serve all three cases:

# nullcontext_demo.py
import sys
import tempfile
from contextlib import AbstractContextManager, nullcontext
from io import StringIO
from pathlib import Path
from typing import IO

def emit(lines: list[str],
         out: IO[str] | Path | None = None) -> None:
    manager: AbstractContextManager[IO[str]]
    match out:
        case Path():
            manager = out.open("w")
        case None:
            manager = nullcontext(sys.stdout)
        case _:
            manager = nullcontext(out)
    with manager as stream:
        for line in lines:
            print(line, file=stream)

emit(["alpha", "beta"])  # Default: stdout, left open
#: alpha
#: beta
buffer = StringIO()
emit(["gamma"], buffer)  # Caller's stream, left open
print(buffer.getvalue().strip(), buffer.closed)
#: gamma False
path = Path(tempfile.gettempdir()) / "emit.txt"
# emit() opened it, so emit() closes it
emit(["delta"], path)
print(path.read_text().strip())
#: delta
path.unlink()

AbstractContextManager[IO[str]] is the type of any context manager whose __enter__() returns an IO[str], so one variable can hold the open file in one branch and a nullcontext in the others.

emit() closes only the file it opened. Exiting a nullcontext does nothing, so the same with block closes the file in the Path branch and leaves the caller’s stream, and sys.stdout, open, as the caller expects.

The Async Protocol

with calls __enter__() and __exit__(). async with calls __aenter__() and __aexit__(), which are coroutines, so the setup and the cleanup can both await. contextlib.asynccontextmanager builds such a manager from an async generator, the same way @contextmanager builds the synchronous form, and AsyncExitStack is the ExitStack equivalent:

# async_manager.py
import asyncio
from collections.abc import AsyncIterator
from contextlib import asynccontextmanager

@asynccontextmanager
async def session(name: str) -> AsyncIterator[str]:
    print(f"open {name}")
    await asyncio.sleep(0.01)  # Setup that waits
    try:
        yield name
    finally:
        await asyncio.sleep(0.01)  # Cleanup that waits
        print(f"close {name}")

async def main() -> None:
    async with session("db") as s:
        print(f"using {s}")

asyncio.run(main())
#: open db
#: using db
#: close db

session() is the generator form with async in front of it. asyncio.run() starts the event loop those awaits need, and Concurrency covers that loop. Everything this chapter says about ordering, the three exception arguments, and suppression through a truthy return applies unchanged. The Concurrency chapter uses async with throughout, for asyncio.TaskGroup, locks, and semaphores. Each of those is an object with the two a-prefixed methods.

An Object Pool

Some objects are expensive to create or rationed by the outside world: database connections, worker processes, licensed sessions. The Object Pool pattern creates a set number of these expensive objects and lends them out. Lending is the dangerous half. Every borrower must return the object on every path out of their code, including the exception path, or the pool slowly drains until the program starves. A context manager can guarantee that every borrowed object comes back. In Python, a pool is a queue plus a @contextmanager method:

# object_pool.py
from collections.abc import Iterator
from contextlib import contextmanager
from dataclasses import dataclass
from queue import Queue

@dataclass(frozen=True)
class Connection:
    number: int

    def query(self, sql: str) -> str:
        return f"connection {self.number}: {sql}"

class Pool[R]:
    def __init__(self, *items: R) -> None:
        self._available: Queue[R] = Queue()
        for item in items:
            self._available.put(item)

    @contextmanager
    def lease(self) -> Iterator[R]:
        item = self._available.get()
        try:
            yield item
        finally:
            self._available.put(item)

    def available(self) -> int:
        return self._available.qsize()

lease() takes an item out of the queue and yields it to the with block. The finally puts it back.

Pool is generic over the pooled type, and it neither creates nor destroys anything. Its one job is tracking custody.

pool_demo.py leases a connection, then crashes inside a second lease:

# pool_demo.py
from contextlib import suppress
from object_pool import Connection, Pool

pool = Pool(Connection(1), Connection(2))
with pool.lease() as conn:
    print(conn.query("SELECT name FROM users"))
    print("available during lease:", pool.available())
#: connection 1: SELECT name FROM users
#: available during lease: 1
print("available after lease:", pool.available())
#: available after lease: 2
with suppress(RuntimeError), pool.lease():
    raise RuntimeError("crash during query")
print("available after crash:", pool.available())
#: available after crash: 2

The second with block crashes, and the finally in lease() still returns the connection, so the count is back to two.

An object pool differs from Flyweight, its nearest neighbor. A flyweight is immutable and shared by everyone at once. A pooled object is usually mutable or stateful, so the pool lends it to one borrower at a time, and the lease exists to take it back.

An Empty Pool Blocks the Borrower

The queue does more than store the idle items. Queue is thread-safe, and get() blocks while the pool is empty, so a borrower waits until another borrower’s with block ends and returns an item to the queue. When several threads share one pool, the queue limits concurrent use, the way a real database connection pool does. That blocking is for threads. In an asyncio program a blocking get() stalls the event loop, so an async pool holds an asyncio.Queue, awaits its get(), and builds lease() with asynccontextmanager.

pool_contention.py puts the pool under real contention. Eight threads share a pool of two connections and lease and release two hundred times each.

# pool_contention.py
import threading
from typing import Final
from object_pool import Connection, Pool

WORKERS: Final[int] = 8
ROUNDS: Final[int] = 200
pool = Pool(Connection(1), Connection(2))
lock = threading.Lock()
held = 0
over_capacity = False

def borrow() -> None:
    global held, over_capacity
    for _ in range(ROUNDS):
        with pool.lease():
            with lock:
                held += 1
                if held > 2:
                    over_capacity = True
            with lock:
                held -= 1

workers = [threading.Thread(target=borrow)
           for _ in range(WORKERS)]
for worker in workers:
    worker.start()
for worker in workers:
    worker.join()
print("over capacity:", over_capacity)
#: over capacity: False
print("pool size after:", pool.available())
#: pool size after: 2

held counts the threads inside the counted window of a lease, incremented once the lease begins and decremented right after the check, so over_capacity turns True when three of those windows overlap. Across sixteen hundred lease-and-release cycles, spread over eight threads competing for two connections, held stays at or below two. A thread that arrives while the pool is empty blocks in get() instead of racing past it.

available() is a snapshot for the demo, not a synchronization primitive. Queue.qsize() becomes approximate once more than one thread is borrowing, because another thread can lease or return between the count and its use.

Testing the Lease

Three tests pin down what the lease guarantees: the item leaves the pool and comes back, it comes back even when the block raises an exception, and the pool hands out the same object rather than a new one:

# test_object_pool.py
import pytest
from object_pool import Connection, Pool

def test_lease_takes_then_returns() -> None:
    pool = Pool(Connection(1), Connection(2))
    with pool.lease():
        assert pool.available() == 1
    assert pool.available() == 2

def test_returned_on_exception() -> None:
    pool = Pool(Connection(1))
    with pytest.raises(RuntimeError):
        with pool.lease():
            raise RuntimeError("boom")
    assert pool.available() == 1

def test_lease_hands_out_same_object() -> None:
    pool = Pool(Connection(1))
    with pool.lease() as first:
        pass
    with pool.lease() as second:
        assert second is first

What the Skeleton Leaves Out

A production pool adds refinements to the Pool skeleton, such as lazily creating items on first demand, validating an item before lending it out, and giving get() a timeout so a starved borrower gets a queue.Empty exception instead of waiting forever.

None of those refinements guard against the opposite mistake. Nothing in Pool stops a borrower from keeping a reference after the with block ends, then using it once the lease has moved on to someone else. The pool hands out the same object again, not a copy, so a stale reference and the new borrower’s object are one and the same:

# leaked_lease.py
from object_pool import Connection, Pool

pool = Pool(Connection(1))
with pool.lease() as first:
    stale = first  # Escapes the block

with pool.lease() as second:
    print(stale is second)
    print(stale.query("late"))
#: True
#: connection 1: late

The assignment to stale stands for any place a borrower stores the reference, such as an attribute or a list. The name first escapes without that assignment. with creates no scope, so first still names the connection after its block ends, and test_lease_hands_out_same_object() relies on that.

stale still points at the Connection that second now legitimately holds. Calling stale.query() after the first with block ended works as if second had called it, because they are the same object. For a mutable pooled resource, that is the source of corruption. Two borrowers each believe they have exclusive use of one connection. Guarding against a stale reference takes a wrapper that invalidates the borrower’s handle on exit, one more refinement the Pool skeleton leaves out.

Each of those refinements is a change inside lease(), invisible to every with pool.lease(). That is the protocol’s payoff. The borrower’s contract is two lines long, the return is impossible to forget, and everything hard about custody lives on the other side of the yield.

Choosing a Form

Not every setup and teardown pair needs a context manager. A pair used once, in one place, is often clearest as a try/finally written inline. Write a manager once you want the with syntax at the call site, or once several places need the same setup and teardown.

Four forms give you a context manager, once you have decided you want one. Try them in this order. Use a contextlib manager when one fits, since suppress, closing, nullcontext, and ExitStack cover most of what people write by hand. Otherwise write a generator with @contextmanager, which is the shortest thing that can express setup, teardown, and a try/finally between them. Write a class with __enter__() and __exit__() when the manager needs state beyond one setup and teardown, methods beyond the two protocol ones, or reuse across several with statements. A generator provides none of those. Add ContextDecorator only when the same bracket should also wrap whole functions.

Whichever form you choose, the borrower’s side contains two lines, and every change you make later goes inside the manager.

Exercises

This chapter’s solutions give a hint, usually the shape of the code, and a full answer for each exercise.

In trace_cm.py, nest a second with Trace("B") as u: block inside the body of the first with Trace("A") as t: block, with its own print(f"inside {u.name}"). Before running it, predict the order in which the six “enter”/“inside”/“exit” lines appear.
In demo_exceptions.py, change expected(ZeroDivisionError) to expected((ZeroDivisionError, TypeError)), then raise a TypeError instead of dividing by zero, and confirm that expected catches and prints it too.
Add a third manager to the with statement in multiple.py, tag("li") again for a second item, and confirm the exit order still reverses the entry order.
In test_object_pool.py, add a test that leases both connections at once, entering a second with pool.lease() block inside the first, and confirms pool.available() reaches 0.
Stack @banner("outer") and @banner("inner") from context_decorator.py on a single function and predict the order of the four bracketing lines before running it.
Write a context manager ignore_missing whose __exit__() suppresses KeyError and lets everything else through, without using contextlib.suppress. Test it with a block that raises a KeyError and a block that raises a ValueError.
Rewrite exit_stack.py to take its names from sys.argv[1:], run it with no arguments and with three, and confirm the close order reverses the open order in both cases.
Wrap the yield in careless() from no_finally.py in try/finally, with the exit line in the finally. Before running it, predict where exit A appears relative to caught: boom.