人人都会AI编程

contextlib 装饰器简化实现

更新时间:2026-07-12

如果每次都要定义一个类、实现 enterexit 方法来创建一个上下文管理器,确实有些繁琐,尤其对于简单的场景。contextlib 模块提供了 @contextmanager 装饰器,可以用生成器函数快速生成一个上下文管理器,让代码更加简洁。

基本原理

@contextmanager 装饰一个生成器函数:

  • yield 之前的代码相当于 enter,可以进行资源初始化、获取锁等操作,并可以将需要管理的对象通过 yield 抛出,赋值给 as 后的变量。
  • yield 之后的代码相当于 exit,负责清理资源、释放锁等,无论 with 块中是否发生异常,yield 之后的代码都会被执行(就像 finally 一样)。
  • 如果在 with 块中抛出了异常,生成器内部可以用 try/except/finally 来捕获或处理异常。

基础示例:替代类的上下文管理器

假设我们需要一个简单的计时器上下文管理器,用来统计一段代码的执行时间:

import time
from contextlib import contextmanager

@contextmanager
def timer():
    start = time.time()
    try:
        yield                # 将控制权交给 with 块
    finally:
        end = time.time()
        print(f"耗时:{end - start:.2f} 秒")

# 使用
with timer():
    # 模拟耗时操作
    total = sum(range(1000000))

这里 yield 之前的代码记录开始时间,yield 之后用 finally 确保无论是否出现异常,都会计算并输出耗时。相较于定义一个 TimerContext 类,这种写法简洁得多。

带返回值的上下文管理器

如果需要在 with 语句中使用 as 获取一个资源对象,只需让 yield 抛出一个对象即可,这是最常见的用法:

from contextlib import contextmanager

@contextmanager
def managed_file(filename, mode='r'):
    f = open(filename, mode)
    try:
        yield f        # f 会赋值给 with ... as f 中的 f
    finally:
        f.close()

# 使用
with managed_file('test.txt', 'w') as f:
    f.write('Hello, contextlib!')

这种写法比手动实现 enterexit 要清晰很多,而且由于 finally 的存在,即使 write 时发生异常,文件也会被正确关闭。

异常处理与抑制

如果想在上下文管理器内部捕获某些异常并按照特定逻辑处理,可以在生成器里加入 except 语句。contextlib 还提供了现成的 suppress 快捷方法,用于忽略指定类型的异常:

from contextlib import suppress

with suppress(FileNotFoundError):
    os.remove('some_file.txt')

这等价于手动 try/except 然后 pass,但意图更明确。

其他便捷方法

  • closing(thing):接受一个带有 close() 方法的对象,with 块结束时自动调用 thing.close(),常用于处理网络连接、数据库游标等。
  • nullcontext():一个不执行任何操作的上下文管理器,用于条件性地进入上下文。

什么时候用类,什么时候用装饰器?

  • 如果上下文管理器的逻辑比较复杂,需要在 enterexit 中做大量状态维护,或者需要复用、继承等面向对象特性,定义类更合适。
  • 对于大多数简单资源管理(文件、锁、临时修改环境变量等),@contextmanager 装饰器写法更直观、更短小。

contextlib 让我们能够用函数思维快速生产上下文管理器,这是 Python “内置电池”理念的又一体现——标准库总在替我们省掉样板代码。