15 · 上下文管理器

with 语句在控制流中已经介绍过，它会划出一段执行区间：进入代码块之前执行 setup，离开代码块之后执行 cleanup，即使代码块抛出异常也一样。这比10 · 清理中的 __del__() 方案可靠得多。本章会展示如何编写自己的 context manager，以及 with 到底如何工作。

最终得到的好处，是借用方只需要写两行契约：

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

无论以什么路径离开这个 block，connection 都会被归还到 pool，包括异常路径；借用方无需写任何归还逻辑。对象池会构建它。

基础上下文管理器

编写 context manager 最简单的方式，是写一个只包含一次 yield 的 generator function，再由 contextlib.contextmanager 装饰器把它转换成 context manager。这里的 yield 与 pytest fixture 中yield 出 fixture 值的方式相同： yield 之前的一切都是 setup，之后的一切都是 teardown。迭代器会完整介绍 generator。本章只需要理解这种形状。

# 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: 会运行 trace() 函数体直到 yield，打印 enter A。 as 会绑定 yield 出来的值，所以 t 是 "A"。接着运行 with 下方的 block。当 block 结束时，trace() 从 yield 之后恢复执行，并打印 exit A。

finally 让 cleanup 可靠：如果 block 抛出异常，这个异常会出现在 yield 所在位置，而 finally 仍会先执行清理，然后异常才继续向外传播。

遗漏 try/finally 是最常见的错误：

# 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

没有 try/finally 时， Python 会通过“在 yield 处抛出 block 的异常”的方式恢复 generator，所以 yield 后面的普通代码根本不会执行，exit A 也不会打印。没有任何地方会警告你：generator 会静默跳过 cleanup，偏偏是在最需要 cleanup 的异常路径上。因此，每个 @contextmanager generator 都应该用 try/finally 包住 yield。

还有一点需要小心：trace("A") 返回的 manager object 只能使用一次。它内部的 generator 只运行一轮，所以在第二个 with 中重复使用同一个对象会失败，而且错误消息并不直观： AttributeError: '_GeneratorContextManager' object has no attribute 'args'。每条 with 语句都应该构造一个新的 manager。在 yield 外面套循环也无法绕过这一点： @contextmanager 强制 generator 只能有一次有效 yield，如果 block 结束后 generator 又运行到第二个 yield，manager 会抛出 RuntimeError （generator didn't stop）。

协议

with 怎么知道应该运行什么？它其实并不了解 generator，也不了解 @contextmanager。一个*上下文管理器（context manager）*就是任何实现了两个方法的对象： __enter__() 在 block 开始时运行，__exit__() 在 block 结束时运行。 @contextmanager 只是根据 generator function 制造了这样一个对象。手写一个类，可以直接看到这套机制。本章所有手写 context manager 类都会把 __init__() 明确写出来，而不改成 @dataclass，这样从 class statement 到两个 protocol method 之间不需要再额外解读：

# 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

with Trace("A") as t: 会依次执行：

计算 Trace("A")，得到 manager object。