contextlib.contextmanager 是 Python 标准库中一个非常实用的装饰器，它可以将一个生成器函数转换为上下文管理器，避免了手动实现 __enter__ 和 __exit__ 方法的繁琐过程。

基本概念 #

什么是上下文管理器？ #

上下文管理器是实现了上下文管理协议的对象，即包含 __enter__() 和 __exit__() 方法的对象。它们通常与 with 语句一起使用，用于资源的获取和释放。

传统实现方式 #

传统实现上下文管理器需要定义一个类：

class MyContext:
    def __enter__(self):
        print("Entering the context")
        return "some value"

    def __exit__(self, exc_type, exc_val, exc_tb):
        print("Exiting the context")
        return False  # 不抑制异常

with MyContext() as value:
    print(f"Inside context, got: {value}")

使用 contextmanager 简化 #

@contextmanager 装饰器可以简化这一过程：

from contextlib import contextmanager

@contextmanager
def my_context():
    print("Entering the context")
    yield "some value"  # yield 的值会赋给 as 后面的变量
    print("Exiting the context")

with my_context() as value:
    print(f"Inside context, got: {value}")

工作原理详解 #

1. 生成器函数：被装饰的函数必须是一个生成器函数（包含 yield 语句）

2. 执行流程：

   • 调用被装饰函数时，它会运行到 yield 语句处暂停
   • yield 之前的部分相当于 __enter__ 方法
   • yield 的值会作为 as 后面的值
   • with 块执行完毕后，生成器会继续执行 yield 之后的代码
   • yield 之后的部分相当于 __exit__ 方法

3. 异常处理：

   • 如果在 with 块中发生异常，它会在生成器中 yield 语句处重新抛出
   • 你可以用 try/finally 或 try/except 来处理异常

高级用法 #

1. 异常处理 #

@contextmanager
def error_handling_context():
    try:
        yield  # 不需要返回值时可以不带参数
    except Exception as e:
        print(f"Caught exception: {e}")
        raise  # 可以选择重新抛出或抑制异常

with error_handling_context():
    1 / 0  # 这会触发异常

2. 资源管理 #

@contextmanager
def managed_resource(*args, **kwargs):
    resource = acquire_resource(*args, **kwargs)
    try:
        yield resource
    finally:
        release_resource(resource)

with managed_resource(timeout=10) as r:
    r.do_something()

3. 嵌套上下文 #

@contextmanager
def nested_contexts():
    with context1() as c1, context2() as c2:
        yield (c1, c2)

with nested_contexts() as (c1, c2):
    # 同时使用 c1 和 c2

实际应用示例 #

1. 计时器 #

@contextmanager
def timer(name):
    start = time.time()
    yield
    duration = time.time() - start
    print(f"{name} took {duration:.2f} seconds")

with timer("database query"):
    # 执行数据库查询
    time.sleep(1)

2. 临时目录 #

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

with temp_dir() as tmp:
    # 在临时目录中工作
    with open(f"{tmp}/test.txt", "w") as f:
        f.write("hello")

3. 数据库事务 #

@contextmanager
def db_transaction(session):
    try:
        yield session
        session.commit()
    except:
        session.rollback()
        raise

with db_transaction(session) as s:
    s.add(User(name="Alice"))

注意事项 #

1. 生成器只能 yield 一次：被装饰的函数必须且只能包含一个 yield 语句
2. 资源释放：确保在 finally 块中释放资源，以防异常导致资源泄漏
3. 异常传播：默认情况下，异常会传播出去，如果需要抑制异常，可以在 except 块中返回 True
4. 性能：对于性能敏感的场景，类形式的上下文管理器可能更快

与类形式对比 #

特性	@contextmanager	类形式
代码量	更简洁	更冗长
可读性	简单场景更易读	复杂场景更清晰
异常处理	需要手动 try/except	在 exit 中自动处理
状态保持	较困难	更容易
性能	稍慢（生成器开销）	稍快

总结 #

contextlib.contextmanager 是一个强大的工具，它：

• 简化了上下文管理器的创建
• 使资源管理代码更加清晰
• 减少了样板代码
• 保持了 Python 的优雅风格

对于大多数日常使用场景，@contextmanager 都能提供简洁高效的解决方案，特别是在需要快速创建一次性上下文管理器时。对于更复杂的需求（如需要保持大量状态或高性能场景），可能需要考虑传统的类实现方式。