如果每次都要定义一个类、实现 enter 和 exit 方法来创建一个上下文管理器,确实有些繁琐,尤其对于简单的场景。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!')
这种写法比手动实现 enter 和 exit 要清晰很多,而且由于 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():一个不执行任何操作的上下文管理器,用于条件性地进入上下文。
什么时候用类,什么时候用装饰器?
- 如果上下文管理器的逻辑比较复杂,需要在
enter和exit中做大量状态维护,或者需要复用、继承等面向对象特性,定义类更合适。 - 对于大多数简单资源管理(文件、锁、临时修改环境变量等),
@contextmanager装饰器写法更直观、更短小。
contextlib 让我们能够用函数思维快速生产上下文管理器,这是 Python “内置电池”理念的又一体现——标准库总在替我们省掉样板代码。