Python 上下文管理器详解

摘要

上下文管理器是 Python 实现资源管理和“设置/清理”模式的优雅方案。通过 with 语句，代码不仅更清晰简洁，而且能可靠地处理资源的释放，即使在异常情况下也不会遗漏。掌握 __enter__/__exit__ 和 contextlib 两种实现方式，可以让你的代码更加健壮、Pythonic。上下文管理器（Context Manager）是 Python 中一个非常实用的特性，它让你能够精确地控制资源的准备与清理，常用于管理文件、网络连接、锁等需要“打开/关闭”操作的资源。最典型的用法就是 with 语句。

1. 为什么需要上下文管理器？

当我们读写文件时，传统写法是：

f = open('data.txt', 'r')
data = f.read()
f.close()

如果 f.read() 发生异常，f.close() 就不会执行，导致文件句柄泄露。虽然可以用 try...finally 解决：

f = open('data.txt', 'r')
try:
    data = f.read()
finally:
    f.close()

但这样写不够简洁。Python 的 with 语句和上下文管理器则优雅地封装了这一模式：

with open('data.txt', 'r') as f:
    data = f.read()
# 离开 with 块后文件会自动关闭，即使发生异常也一样

with 语句确保了进入时自动调用准备逻辑，退出时自动调用清理逻辑，无论代码块是否抛出异常。

2. `with` 语句的工作原理

with EXPRESSION as TARGET:
    SUITE

Python 实际做了以下事情：

计算表达式 EXPRESSION，得到上下文管理器对象。
调用该对象的 __enter__() 方法，返回值（如果有）赋给 as 后面的变量 TARGET。
执行 SUITE 代码块。
无论代码块是否抛出异常，最后都会调用上下文管理器的 __exit__() 方法。
- 如果没有异常，__exit__ 的三个参数 exc_type, exc_val, exc_tb 都是 None。
- 如果发生了异常，这三个参数会携带异常信息，__exit__ 可以决定是否抑制该异常（返回 True 则抑制，返回 False 或 None 则异常会继续向外传播）。

3. 如何自定义上下文管理器

有两种方式来实现：类方式（定义 __enter__ 和 __exit__）和生成器方式（使用 contextlib 模块）。

3.1 类方式实现

只要类定义了 __enter__ 和 __exit__ 方法，它的实例就可以用作上下文管理器。

class DatabaseConnection:
    def __init__(self, db_name):
        self.db_name = db_name

    def __enter__(self):
        # 建立连接
        print(f"连接数据库 {self.db_name}...")
        self.conn = f"connection_to_{self.db_name}"  # 模拟连接
        return self.conn  # 返回的资源会被 as 后的变量接收

    def __exit__(self, exc_type, exc_val, exc_tb):
        # 关闭连接
        print("关闭数据库连接...")
        # 在此可以做异常处理
        if exc_type:
            print(f"发生异常：{exc_val}")
        # 返回 False 让异常继续传播，返回 True 则吞掉异常
        return False

# 使用
with DatabaseConnection("my_db") as conn:
    print(f"使用 {conn} 查询数据")
    # 如果这里抛出异常，连接依然会被关闭

运行输出：

连接数据库 my_db...
使用 connection_to_my_db 查询数据
关闭数据库连接...

__exit__ 的参数含义：

exc_type：异常类，如 ZeroDivisionError
exc_val：异常实例
exc_tb：traceback 对象

如果你想在 __exit__ 中处理异常并阻止其向上传播，可以返回 True。

3.2 使用 `contextlib` 模块（生成器方式）

标准库 contextlib 提供了更简单的写法：用 @contextmanager 装饰一个生成器函数。yield 之前的代码相当于 __enter__，yield 返回的值给 as 变量；yield 之后的代码相当于 __exit__。生成器函数必须只能 yield 一次。

from contextlib import contextmanager

@contextmanager
def managed_file(name, mode):
    print("打开文件...")
    f = open(name, mode)
    try:
        yield f          # 将文件对象交给 with 代码块
    finally:
        print("关闭文件...")
        f.close()

with managed_file('test.txt', 'w') as f:
    f.write('Hello')
    # 如果发生异常，依然会执行 finally 块，安全关闭

这种写法不需要定义类，更加简洁，且利用 try...finally 天然保证了清理逻辑的执行。注意：如果生成器内的 yield 被异常中断，异常会被重新抛出到 yield 那一行，你可以用 try 捕获它（就像上面那样），然后决定如何处理。

4. 常用的内置上下文管理器

文件操作：open() 是最经典的例子。
线程锁：threading.Lock() 等支持上下文管理，自动获取和释放。
decimal 精度：decimal.localcontext() 临时修改小数精度。
unittest.mock.patch：测试中打补丁。
contextlib.suppress：忽略指定异常。如 with suppress(FileNotFoundError): os.remove('somefile')。
contextlib.redirect_stdout：重定向标准输出。
contextlib.ExitStack：动态管理多个上下文管理器，特别是在不知道需要多少个资源时非常有用。

5. 应用场景

上下文管理器适用于任何需要成对操作的场景，如：

文件打开/关闭
网络连接（socket）建立/断开
数据库事务开始/提交或回滚
加锁/解锁
临时修改环境变量或工作目录
性能计时（进入记录时间，退出打印耗时）
日志上下文（进入时添加上下文信息，退出时清除）

示例：简单的计时器

import time
from contextlib import contextmanager

@contextmanager
def timer(description):
    start = time.perf_counter()
    yield
    elapsed = time.perf_counter() - start
    print(f"{description}: {elapsed:.4f} 秒")

with timer("计算耗时"):
    sum(range(1000000))

6. 异常处理细节

如果 __exit__ 返回 False（或者 None），with 代码块中产生的异常会正常传播。
如果返回 True，异常会被吞掉，继续执行 with 之后的语句。
在 @contextmanager 装饰的生成器中，如果你捕获了异常并继续执行（没有重新 raise），就相当于返回了 True，异常会被抑制。若希望异常继续传播，需要重新 raise。