Python contextlib.contextmanager简化上下文管理器的写法

编写 Python 代码时，经常需要管理资源（如文件句柄、数据库连接、锁等），确保在使用完毕后正确释放。传统方式需要创建一个类并实现 __enter__ 和 __exit__ 方法，代码结构冗长且逻辑分散。Python 标准库 contextlib 提供的 @contextmanager 装饰器，允许你只用一个函数就实现同样的功能，将“准备”和“清理”逻辑紧密写在一起。

1. 理解传统写法的痛点

在开始简化之前，先通过一个文件操作的例子回顾传统方式。你需要定义一个类，并处理两个魔术方法。

观察 以下代码，注意其结构：

class ManagedFile:
    def __init__(self, filename):
        self.filename = filename

    def __enter__(self):
        print("正在打开文件...")
        self.file = open(self.filename, 'w')
        return self.file

    def __exit__(self, exc_type, exc_val, exc_tb):
        if self.file:
            self.file.close()
        print("文件已关闭")
        if exc_type:
            print(f"发生错误: {exc_val}")
        return True  # 抑制异常

# 使用
with ManagedFile('test.txt') as f:
    f.write('Hello World')
    # raise ValueError("测试异常") 

这种写法虽然严谨，但为了仅仅打开和关闭文件，需要编写大量样板代码，且逻辑被分割在不同的方法中，阅读不够连贯。

2. 使用 @contextmanager 重构

contextlib.contextmanager 的核心思想是将 __enter__ 的逻辑放在 yield 之前，将 __exit__ 的逻辑放在 yield 之后。

执行 以下步骤进行重构：

1. 导入 contextlib 模块。
2. 定义 一个生成器函数，函数名即为你想要的上下文管理器名称。
3. 添加 @contextmanager 装饰器到该函数上方。
4. 编写 资源获取代码（原 __enter__ 内容）。
5. 使用 yield 关键字返回需要赋值给 as 变量的对象。
6. 编写 资源释放代码（原 __exit__ 内容）。

查看 重构后的代码：

from contextlib import contextmanager

@contextmanager
def managed_file(filename):
    # --- 进入阶段 (对应 __enter__) ---
    print("正在打开文件...")
    f = open(filename, 'w')
    try:
        yield f  # 暂停执行，将 f 传给 with 语句中的变量
    finally:
        # --- 退出阶段 (对应 __exit__) ---
        if f:
            f.close()
        print("文件已关闭")

# 使用方式完全一致
with managed_file('test.txt') as f:
    f.write('Hello World')

3. 掌握执行流程与 yield 机制

理解 yield 在这里的运作方式至关重要。它不仅仅是返回值，更是程序执行的“暂停点”。

为了直观展示控制流，参考 以下的执行顺序图：

注意 以下几点细节：

• yield 之前的代码在进入 with 块时立即执行。
• yield 的值会被赋值给 as 后面的变量。
• 无论 with 块中是否发生异常，yield 之后的代码（通常包裹在 finally 中）保证会被执行。

4. 妥善处理异常

在使用 @contextmanager 时，异常处理不能像类方法那样通过 __exit__ 的参数来控制，必须依赖 try...finally 结构。

遵循 以下规则确保健壮性：

1. 包裹 yield 语句在 try 块中。
2. 放置 清理代码在 finally 块中，确保资源必释放。
3. 捕获 异常（如果需要），或者让异常自然抛出。

分析 以下处理异常的示例：

@contextmanager
def db_connection_handler():
    conn = "Database Connection Object" # 模拟获取连接
    print("数据库已连接")
    try:
        yield conn
    except ValueError as e:
        # 如果 with 块中抛出 ValueError，会在这里被捕获
        print(f"捕获到业务异常，进行回滚操作: {e}")
        # 可以选择不 raise，则异常被吞掉；或者 raise 继续向上抛出
    finally:
        # 无论是否异常，都会断开连接
        print("数据库连接已关闭")

# 测试异常捕获
with db_connection_handler() as db:
    print(f"使用连接: {db}")
    raise ValueError("数据写入失败")

在这个例子中，如果 with 块内发生了 ValueError，控制权会直接跳转到 except ValueError 块，执行回滚逻辑，随后进入 finally 块关闭连接。如果发生的是其他类型的异常（如 TypeError），则不会被 except 捕获，会直接跳过 except 进入 finally，之后异常继续向上传播。

5. 对比两种实现方式

为了更清晰地展示差异，对照 下表中的核心区别：

6. 实战应用：计时器上下文管理器

为了巩固知识，我们构建 一个实用的工具：一个用于测量代码块执行时间的上下文管理器。

1. 导入 time 模块和 contextmanager。
2. 定义 函数 timer()，不接受参数。
3. 记录 yield 之前的时间戳 start。
4. yield 一个控制对象（这里只是为了暂停，不需要返回值）。
5. 计算 yield 之后的时间差并打印。

编写 代码如下：

import time
from contextlib import contextmanager

@contextmanager
def timer():
    start = time.time()
    # yield 前不返回具体值，或者返回 None
    yield 
    end = time.time()
    print(f"代码执行耗时: {end - start:.4f} 秒")

# 使用示例
with timer():
    # 模拟耗时操作
    sum([i**2 for i in range(1000000)])

运行 这段代码，你将看到精确的执行耗时输出，且无需手动在代码前后写获取时间的逻辑，代码主逻辑更加纯净。

通过以上步骤，你已经掌握了如何利用 contextlib.contextmanager 将繁琐的类定义简化为直观的生成器函数，既保证了代码的可读性，又完美处理了资源的清理工作。