Python中的上下文管理器和with语句

1. 总览

打开文件、数据库连接或获取锁都需要在最后进行清理。即使中间的代码抛出异常，这个步骤也必须运行。Python通过上下文管理器和with语句来解决这个问题。

本文介绍了with如何在底层工作，如何用类或contextlib编写上下文管理器。你还将学到如何在__exit__中处理异常、如何为上下文管理器编写测试，以及使用它们时常见的错误。

2. 问题：用try/finally清理资源

让我们看看不使用with的文件读取方式：

f = open("orders.csv", encoding="utf-8")
try:
    header = f.readline()
    print(header)
finally:
    f.close()

这种方式是对的，但很冗长。也容易忘记。如果你把open()放在try里，而open()又失败了，那么f还没有被赋值，finally块会抛出NameError。with语句完美地封装了上述模式：

with open("orders.csv", encoding="utf-8") as f:
    header = f.readline()
    print(header)

print(f.closed)

order_id,customer,amount

True

当退出with块时，文件已被关闭。无论是正常退出还是有异常，这都是成立的。

3. 上下文管理器协议

任何拥有__enter__和__exit__两个方法的对象都可以与with一起使用。在概念上，该语句：

with EXPR as target:
    BODY

等价于以下代码：

manager = EXPR
target = manager.__enter__()
try:
    BODY
except BaseException as e:
    if not manager.__exit__(type(e), e, e.__traceback__):
        raise
else:
    manager.__exit__(None, None, None)

有三点需要记住：

• 赋给target的值是**__enter__()的结果**，不是EXPR本身。
• __exit__总是被调用。当有异常时，它接收(exc_type, exc_value, traceback)。当没有错误时，所有三个参数都是None。
• 如果__exit__返回真值，异常会被抑制。否则，异常会被重新抛出。

3.1. 用类编写上下文管理器

第一个例子是一个Timer，用来测量代码块的运行时间：

import time


class Timer:
    def __init__(self, label: str):
        self.label = label
        self.elapsed: float | None = None

    def __enter__(self):
        self._start = time.perf_counter()
        return self

    def __exit__(self, exc_type, exc, tb):
        self.elapsed = time.perf_counter() - self._start
        print(f"[{self.label}] {self.elapsed:.3f}s")
        return False


with Timer("tinh_tong_don_hang") as t:
    total = sum(range(5_000_000))

print(t.elapsed > 0)

[tinh_tong_don_hang] 0.081s
True

逐部分解释：

• __enter__记录开始时间并返回self。因此t就是Timer对象本身。
• __exit__计算经过的时间，打印到屏幕，然后返回False。任何异常都不会被抑制。
• 变量t在with块之后仍然可用。原因是with不会创建新的作用域。

3.2. 在__exit__中处理异常

当你需要根据代码块是否成功进行不同的反应时，上下文管理器的作用最为明显。以下例子是一个简单的账户交易（ledger）：

class Transaction:
    def __init__(self, ledger: list):
        self.ledger = ledger
        self.pending: list = []

    def __enter__(self):
        return self.pending

    def __exit__(self, exc_type, exc, tb):
        if exc_type is None:
            self.ledger.extend(self.pending)
            print(f"commit {len(self.pending)} 条记录")
        else:
            print(f"rollback 原因是 {exc_type.__name__}: {exc}")
        return False


ledger = []

with Transaction(ledger) as tx:
    tx.append({"order_id": 1, "amount": 250_000})
    tx.append({"order_id": 2, "amount": 120_000})

try:
    with Transaction(ledger) as tx:
        tx.append({"order_id": 3, "amount": 90_000})
        raise ValueError("余额不足")
except ValueError:
    pass

print(len(ledger))

commit 2 条记录
rollback 原因是 ValueError: 余额不足
2

在第二个with块中，订单3没有被写入ledger。因为__exit__返回False，ValueError仍然会被抛出，调用者可以决定如何处理。

注意： 只有当你故意想要抑制异常时，才从__exit__返回True。一个__exit__方法不小心return self或返回其他真值对象会导致错误无声地消失，没有任何迹象。

4. 用@contextmanager简化写法

为了一个简单的setup和teardown步骤而写一个完整的类太冗长了。模块contextlib提供了一个@contextmanager装饰器，可以将生成器转换为上下文管理器：

from contextlib import contextmanager


class PaymentService:
    def connect(self):
        print("打开支付网关连接")

    def charge(self, customer: str, amount: int):
        print(f"从{customer}扣除{amount}")

    def disconnect(self):
        print("关闭连接")


@contextmanager
def payment_session(service: PaymentService):
    service.connect()
    try:
        yield service
    finally:
        service.disconnect()


with payment_session(PaymentService()) as payment:
    payment.charge("Lan", 300_000)

打开支付网关连接
从Lan扣除300000
关闭连接

装饰器如何将生成器映射到上下文管理器协议：

• yield之前的代码充当__enter__。
• 被yield的值是分配给as后面的变量的值。
• yield之后的代码充当__exit__。如果with块抛出异常，该异常会在**yield处**被重新抛出。

因为异常在yield处出现，yield周围的try/finally块是必需的。没有它，disconnect()在出错时不会运行。

要单独处理不同类型的错误，可以使用额外的except：

@contextmanager
def safe_charge(order_id: int):
    try:
        yield
    except ConnectionError as e:
        print(f"订单{order_id}：网络错误，稍后重试({e})")


with safe_charge(42):
    raise ConnectionError("timeout")

print("继续处理")

订单42：网络错误，稍后重试(timeout)
继续处理

生成器捕获错误但不重新抛出，所以异常被抑制。这等价于__exit__返回True。

5. contextlib中的实用工具

5.1. suppress: 忽略特定的异常

import os
from contextlib import suppress

with suppress(FileNotFoundError):
    os.remove("cache_orders.tmp")

print("完成")

完成

上面的代码比try: ... except FileNotFoundError: pass更短，意图也更清晰。

5.2. ExitStack: 资源数量未知的情况

假设你需要合并多个发票文件，但文件数量只在运行时才知道：

from contextlib import ExitStack
from pathlib import Path


def merge_invoices(paths: list[Path], output: Path) -> None:
    with ExitStack() as stack:
        files = [stack.enter_context(open(p, encoding="utf-8")) for p in paths]
        with open(output, "w", encoding="utf-8") as out:
            for f in files:
                out.write(f.read())


for i in range(3):
    Path(f"invoice_{i}.txt").write_text(f"发票{i}\n", encoding="utf-8")

merge_invoices([Path(f"invoice_{i}.txt") for i in range(3)], Path("all.txt"))
print(Path("all.txt").read_text(encoding="utf-8"))

发票0
发票1
发票2

enter_context调用__enter__并记住相应的__exit__。当ExitStack结束时，文件按打开时的相反顺序被关闭。如果第三个文件打开失败，前两个文件仍然会被正确关闭。

5.3. closing和nullcontext

• closing(obj)将只有close()方法的对象转换为上下文管理器。
• nullcontext(value)是一个什么都不做的上下文管理器。当一个参数可能是也可能不是上下文管理器时很有用：

from contextlib import nullcontext
import threading

def update_stock(product: dict, qty: int, lock: threading.Lock | None = None):
    with lock if lock is not None else nullcontext():
        product["stock"] -= qty

product = {"sku": "AO-01", "stock": 10}
update_stock(product, 3)
update_stock(product, 2, threading.Lock())
print(product["stock"])

5

6. 一个with语句中的多个上下文管理器

你可以在同一个with语句中声明多个上下文管理器。从Python 3.10开始，使用括号的语法允许优雅地换行：

with (
    open("orders.csv", encoding="utf-8") as src,
    open("orders_backup.csv", "w", encoding="utf-8") as dst,
):
    dst.write(src.read())

这种写法等价于两个嵌套的with块。src先打开，最后关闭。

7. 异步上下文管理器

对于asyncio代码，对应的协议是__aenter__和__aexit__，与async with一起使用。它的生成器版本是@asynccontextmanager：

import asyncio
from contextlib import asynccontextmanager


@asynccontextmanager
async def order_queue(name: str):
    print(f"连接队列{name}")
    await asyncio.sleep(0.01)
    try:
        yield [f"order-{i}" for i in range(3)]
    finally:
        await asyncio.sleep(0.01)
        print(f"关闭队列{name}")


async def main():
    async with order_queue("payments") as messages:
        for m in messages:
            print("处理", m)

asyncio.run(main())

连接队列payments
处理order-0
处理order-1
处理order-2
关闭队列payments

注意： 对异步上下文管理器使用普通的with会导致TypeError（或在较旧的Python版本中导致AttributeError）。反之亦然。这两个协议是完全分开的。

8. 测试上下文管理器

上下文管理器应该测试成功路径和失败路径。首先安装pytest：

pip install pytest

将第3.2节的Transaction类放入transaction.py文件，然后编写测试：

# test_transaction.py
import pytest
from transaction import Transaction


def test_commit_when_no_error():
    ledger = []
    with Transaction(ledger) as tx:
        tx.append({"order_id": 1, "amount": 100})

    assert ledger == [{"order_id": 1, "amount": 100}]


def test_rollback_and_propagate_error():
    ledger = []
    with pytest.raises(ValueError, match="余额"):
        with Transaction(ledger) as tx:
            tx.append({"order_id": 2, "amount": 999})
            raise ValueError("余额不足")

    assert ledger == []

pytest -q

..                                                       [100%]
2 passed in 0.01s

第二个测试同时检查两个事情：数据被回滚，异常没有被抑制。pytest.raises本身也是一个上下文管理器，这是本文主题的一个很好的例子。

9. 常见错误和何时不应使用

忘记yield周围的try/finally。 这是使用@contextmanager最常见的错误：

@contextmanager
def broken_session(service):
    service.connect()
    yield service
    service.disconnect()  # 如果with块抛出异常不会运行

在with结束后使用资源。 变量仍然存在，但资源已被释放：

with open("orders.csv", encoding="utf-8") as f:
    pass
f.read()

ValueError: I/O operation on closed file.

在生成器中多次yield。 与@contextmanager一起使用的生成器必须恰好yield一次。如果yield多于一次，Python会抛出RuntimeError: generator didn't stop。

重新使用从@contextmanager创建的上下文管理器。 这个对象只能使用一次。对同一对象第二次调用with会出错。每次使用时，都应该再次调用函数来创建新对象。

滥用。 以下情况不应使用上下文管理器：

• 资源需要比一个代码块的生命周期更长，例如在整个应用程序生命周期中存在的连接池。在这种情况下，应该在应用程序启动和关闭层管理它。
• 逻辑没有明确的“打开/关闭”对。普通函数会更易读。
• 想用suppress(Exception)忽略所有异常。这就像except: pass一样隐藏bug。

注意： 不要依赖垃圾收集器或__del__来关闭资源。它们运行的时间没有保证，特别是在PyPy这样的实现上。请使用with。

10. 总结

上下文管理器将setup和teardown代码放在一起，并确保teardown总是运行，即使有异常也不例外。你已经看到：

• __enter__和__exit__协议，以及__exit__返回值的含义。
• 使用@contextmanager的简洁写法，yield周围的try/finally是必需的。
• 实用工具suppress、ExitStack、closing、nullcontext及其异步版本。
• 如何使用pytest测试commit和rollback路径。

当你的代码有“获取然后释放”的模式（文件、锁、事务、连接）时，应该考虑将其包装在上下文管理器中。