Python contextlib 实战：用 with 管好资源的 5 个实用技巧

写 Python 时，with open(...) 几乎人人都会用。但很多人不知道，标准库 contextlib 能让你用极少的代码写出自己的上下文管理器，并优雅地处理计时、临时状态、多资源清理等场景。本文给出 5 个可直接运行的技巧（Python 3.8+，代码均已实际运行验证）。

1. 用 @contextmanager 写一个计时器

不必再写带 __enter__/__exit__ 的类，一个生成器函数就够了：

import time
from contextlib import contextmanager

@contextmanager
def timer(name):
    start = time.perf_counter()
    try:
        yield
    finally:
        print(f"{name} 耗时 {(time.perf_counter() - start) * 1000:.0f} ms")

with timer("sleep"):
    time.sleep(0.05)   # 输出：sleep 耗时 50 ms 左右

yield 之前相当于 __enter__，之后相当于 __exit__。

2. 一定要用 try/finally 包住 yield

这是最常见的坑。如果 with 块里抛出异常，异常会在 yield 处被重新抛出；如果你没有 try/finally，yield 之后的清理代码就不会执行：

@contextmanager
def bad():
    print("enter")
    yield
    print("cleanup")      # with 块内抛异常时，这一行不会执行

try:
    with bad():
        raise RuntimeError
except RuntimeError:
    print("cleanup 没有打印")

运行输出 enter 和 cleanup 没有打印。凡是需要保证释放的资源（锁、文件、临时目录），都要写 try/finally。

3. suppress：优雅地忽略指定异常

删除一个可能不存在的文件，比起 try/except/pass 更直观：

import os
from contextlib import suppress

with suppress(FileNotFoundError):
    os.remove("none.txt")

注意只忽略你明确列出的异常类型，不要写 suppress(Exception)，否则真正的 bug 会被吞掉。

4. ExitStack：动态数量的资源

如果要同时打开的文件数量在运行时才知道，没法写固定层数的 with a, b, c。ExitStack 可以解决：

from contextlib import ExitStack

names = ["0.txt", "1.txt", "2.txt"]   # 假设这些文件已存在
with ExitStack() as stack:
    files = [stack.enter_context(open(n)) for n in names]
    print([f.read() for f in files])
print([f.closed for f in files])      # [True, True, True]

退出时所有文件按后进先出的顺序关闭，即使中途某个 open 失败，已打开的也会被正确关闭。

5. redirect_stdout：临时捕获输出

测试或调用只会 print 的第三方函数时很有用：

import io
from contextlib import redirect_stdout

buf = io.StringIO()
with redirect_stdout(buf):
    print("hello")
print(repr(buf.getvalue()))   # 'hello\n'

附：类写法与事务语义

需要保存状态或可重复使用时，可以用类。__exit__ 的返回值决定异常是否被吞掉：返回 True 吞掉，返回 False（或 None）则继续向外抛出。

class Txn:
    def __enter__(self):
        print("begin")
        return self

    def __exit__(self, exc_type, exc, tb):
        print("rollback" if exc_type else "commit")
        return False

with Txn():
    pass            # begin / commit

try:
    with Txn():
        raise KeyError
except KeyError:
    pass            # begin / rollback

这正是数据库事务的典型模式：正常结束提交，异常则回滚，同时不隐藏异常本身。

小结

• 简单场景用 @contextmanager，并且 yield 外面务必加 try/finally；
• 忽略特定异常用 suppress，不要忽略宽泛的 Exception；
• 资源数量不固定时用 ExitStack；
• 需要控制异常是否外抛时，看 __exit__ 的返回值。

把这些小工具用熟，代码里的 try/finally 会少很多，资源泄漏的风险也会更低。