contextlib 是 Python 标准库中的一个模块,专门用于帮助开发者更方便地创建和处理 上下文管理器(Context Managers)。
介绍
通常我们使用 with 语句来管理资源(如文件打开/关闭、锁的获取/释放、数据库连接等)。contextlib 提供了实用的工具,让我们不必每次都编写带有 __enter__ 和 __exit__ 方法的繁琐类。
核心主角:@contextmanager
这是 contextlib 中最常用的工具。它是一个装饰器,可以将一个**生成器函数(generator function)**直接转化为一个上下文管理器。
原理
yield 之前 的代码相当于 __enter__ 方法(进入 with 块前执行)。
yield 的值 会被赋给 with ... as var 中的 var。
yield 之后 的代码相当于 __exit__ 方法(离开 with 块后执行)。
⚠️ 重要注意事项
如果在 with 代码块中发生了异常,该异常会在生成器的 yield 语句处被抛出。因此,为了保证资源一定能被释放(即执行 yield 后面的代码),通常必须配合 try...finally 结构使用。
@contextmanager 用法示例
示例 A:最简单的用法(不带错误处理)
这个例子展示了执行顺序。
from contextlib import contextmanager
@contextmanager
def simple_context(name):
print(f"1. [进入] 正在初始化: {name}")
# yield 出来的值会被 as 接收
yield f"Hello, {name}"
print(f"3. [退出] 清理完成: {name}")
# 使用
with simple_context("World") as msg:
print(f"2. [执行] 在 with 块内部, 收到: {msg}")
# 输出:
# 1. [进入] 正在初始化: World
# 2. [执行] 在 with 块内部, 收到: Hello, World
# 3. [退出] 清理完成: World
示例 B:模拟资源管理(标准写法,带 try...finally)
这是生产环境中最常见的模式,用于确保资源(如文件、连接)总是关闭。
from contextlib import contextmanager
class MockDatabase:
def query(self):
print(" -> 执行数据库查询")
def close(self):
print(" -> 关闭数据库连接")
@contextmanager
def open_database():
print("1. [Setup] 建立连接")
db = MockDatabase()
try:
yield db # 将 db 对象交给 with 块使用
except Exception as e:
print(f" [Error] 发生错误: {e}")
raise e # 可以选择重新抛出异常,或者吞掉它
finally:
# 无论是否发生异常,这行代码都会执行
print("3. [Teardown] 正在清理资源...")
db.close()
# 正常情况
print("--- 正常运行 ---")
with open_database() as db:
print("2. [Action] 使用数据库")
db.query()
# 异常情况
print("\n--- 发生异常 ---")
try:
with open_database() as db:
print("2. [Action] 准备搞破坏")
raise ValueError("这就尴尬了")
except ValueError:
print("--- 捕获到了外层异常 ---")
输出结果:
--- 正常运行 ---
1. [Setup] 建立连接
2. [Action] 使用数据库
-> 执行数据库查询
3. [Teardown] 正在清理资源...
-> 关闭数据库连接
--- 发生异常 ---
1. [Setup] 建立连接
2. [Action] 准备搞破坏
[Error] 发生错误: 这就尴尬了
3. [Teardown] 正在清理资源...
-> 关闭数据库连接
--- 捕获到了外层异常 ---
可以看到,即使中间抛出了 ValueError,finally 块中的 db.close() 依然被执行了。
contextlib 其他常用工具
除了 @contextmanager,这个库还有很多好用的工具:
closing(thing)
用于那些有 close() 方法但本身不是上下文管理器的对象(比如 urllib.request.urlopen 返回的对象)。
from contextlib import closing
from urllib.request import urlopen
# urlopen 返回的对象没有 __enter__/__exit__,但有 close()
# 使用 closing 可以确保退出时自动调用 page.close()
with closing(urlopen('https://www.python.org')) as page:
for line in page:
print(line)
break
suppress(*exceptions)
用于显式地忽略指定的异常。这比写 try...except pass 更优雅。
import os
from contextlib import suppress
# 传统写法
try:
os.remove('somefile.tmp')
except FileNotFoundError:
pass
# 使用 suppress 写法 (更易读)
with suppress(FileNotFoundError):
os.remove('somefile.tmp')
redirect_stdout / redirect_stderr
用于临时将输出重定向到文件或类文件对象中。
from contextlib import redirect_stdout
import io
f = io.StringIO()
with redirect_stdout(f):
print("这句话不会打印在屏幕上")
print("它会进入 f 对象中")
print(f"捕获到的内容: {f.getvalue()}")
nullcontext (Python 3.7+)
有时候我们需要根据条件决定是否使用上下文管理器。如果不需要,可以使用 nullcontext 占位,它什么都不做。
from contextlib import nullcontext
def process_file(file_path, use_lock=False):
# 如果 use_lock 为 True,使用 my_lock;否则使用 nullcontext (即不加锁)
lock_context = my_lock if use_lock else nullcontext()
with lock_context:
# 做一些操作
pass
ExitStack (进阶神器)
当你需要管理数量不确定的上下文管理器时(例如同时打开列表中的 10 个文件),ExitStack 非常有用。
from contextlib import ExitStack
filenames = ['file1.txt', 'file2.txt', 'file3.txt']
with ExitStack() as stack:
# 动态打开所有文件
files = [stack.enter_context(open(fname)) for fname in filenames]
# 在这里可以同时操作这 3 个文件
# ...
# 离开 with 块时,stack 会自动按相反顺序关闭所有文件
总结
@contextmanager:是将生成器函数改写为上下文管理器的最快方式。记得利用 try...finally 来处理异常和资源清理。
contextlib 模块:是编写 Pythonic 代码(优雅、简洁)的利器,尤其是 suppress 和 closing 能大幅简化样板代码。