python @contextmanager

Python 的 @contextmanager 装饰器详解

1. 基本概念

@contextmanager 是 Python 标准库 contextlib 中的一个装饰器，用于简化上下文管理器的创建。它允许我们使用生成器函数来创建上下文管理器，而不需要实现完整的 __enter__ 和 __exit__ 方法。

2. 基本语法

from contextlib import contextmanager

@contextmanager
def managed_resource(*args, **kwargs):
    # 资源获取
    resource = acquire_resource(*args, **kwargs)
    try:
        yield resource
    finally:
        # 资源释放
        release_resource(resource)

3. 工作原理

1. 装饰器转换：

   • @contextmanager 将生成器函数转换为上下文管理器
   • 自动处理 __enter__ 和 __exit__ 方法

2. 执行流程：

   # 当使用 with 语句时：
   with managed_resource() as resource:
       # 1. 执行 yield 之前的代码（资源获取）
       # 2. yield 返回资源
       # 3. 执行 with 块中的代码
       # 4. 执行 finally 块中的代码（资源释放）
   

4. 关键组件

1. yield 语句：

   • 将资源传递给 with 语句的 as 子句
   • 分隔资源获取和释放的代码

2. try-finally 结构：

   • 确保资源总是被释放
   • 即使在发生异常的情况下

5. 实际示例

5.1 文件操作

@contextmanager
def open_file(path, mode):
    f = open(path, mode)
    try:
        yield f
    finally:
        f.close()

# 使用
with open_file('test.txt', 'r') as f:
    print(f.read())

5.2 数据库连接

@contextmanager
def db_connection(connection_string):
    conn = create_connection(connection_string)
    try:
        yield conn
    finally:
        conn.close()

# 使用
with db_connection('mysql://user:pass@localhost/db') as conn:
    cursor = conn.cursor()
    cursor.execute('SELECT * FROM table')

5.3 临时目录

@contextmanager
def temp_dir():
    import tempfile
    import shutil
    dir_path = tempfile.mkdtemp()
    try:
        yield dir_path
    finally:
        shutil.rmtree(dir_path)

# 使用
with temp_dir() as dir_path:
    # 在临时目录中工作
    with open(os.path.join(dir_path, 'test.txt'), 'w') as f:
        f.write('test')

6. 异常处理

@contextmanager
def managed_resource():
    resource = acquire()
    try:
        yield resource
    except Exception as e:
        # 处理异常
        print(f"Error occurred: {e}")
        raise
    finally:
        release(resource)

7. 嵌套使用

@contextmanager
def outer_resource():
    with inner_resource() as inner:
        yield inner

# 使用
with outer_resource() as resource:
    # 使用资源

8. 优势

1. 代码简洁：

   • 比实现完整的上下文管理器类更简洁
   • 减少样板代码

2. 易于理解：

   • 资源获取和释放的逻辑清晰可见
   • 使用熟悉的 try-finally 结构

3. 灵活性：

   • 可以处理复杂的资源管理场景
   • 支持异常处理和资源清理

9. 注意事项

1. yield 只能出现一次：

   • 生成器函数中只能有一个 yield 语句
   • 多个 yield 会导致上下文管理器行为异常

2. 资源释放：

   • 确保在 finally 块中释放资源
   • 避免资源泄漏

3. 异常传播：

   • 默认情况下，异常会传播到 with 语句之外
   • 如果需要处理异常，在 try-except 块中处理

4. 返回值：

   • yield 语句可以返回值
   • 返回值会被赋值给 with 语句的 as 子句中的变量

10. 总结

@contextmanager 装饰器是 Python 中创建上下文管理器的强大工具，它：

• 简化了上下文管理器的创建
• 提供了清晰的资源管理方式
• 确保了资源的正确获取和释放
• 支持异常处理
• 使代码更加简洁和可读

在实际开发中，@contextmanager 是处理资源管理的首选方式，特别是在需要管理多个资源或处理复杂资源管理场景时。