Python中 使用 contextlib 简化上下文管理器

在 Python 中，contextlib 模块提供了一种更简洁的方式来创建自定义上下文管理器，通过 @contextmanager 装饰器可以将一个生成器函数快速转换为上下文管理器，避免手动编写 __enter__ 和 __exit__ 方法的样板代码。

核心方法：使用 @contextmanager 装饰器

1. 定义生成器函数：

   • yield 之前的代码相当于 __enter__ 的逻辑（资源分配）。
   • yield 之后的代码相当于 __exit__ 的逻辑（资源释放）。

2. 异常处理：

   • 使用 try...finally 或 except 确保退出逻辑始终执行。
   • yield 的值会传递给 with 语句中的 as 变量。

示例：文件操作的上下文管理器

传统类实现方式

class FileManager:
    def __init__(self, filename, mode):
        self.filename = filename
        self.mode = mode

    def __enter__(self):
        self.file = open(self.filename, self.mode)
        return self.file

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.file.close()

# 使用
with FileManager("test.txt", "w") as f:
    f.write("Hello")

使用 @contextmanager 简化

from contextlib import contextmanager

@contextmanager
def file_manager(filename, mode):
    file = open(filename, mode)
    try:
        yield file  # yield 的值传递给 as 后的变量
    finally:
        file.close()  # 确保无论是否异常都会关闭文件

# 使用
with file_manager("test.txt", "w") as f:
    f.write("Hello")

高级用法

1. 临时修改全局状态

@contextmanager
def temporary_redirect(new_target):
    original = sys.stdout
    sys.stdout = new_target
    try:
        yield
    finally:
        sys.stdout = original

# 使用：临时重定向输出到文件
with open("output.txt", "w") as f, temporary_redirect(f):
    print("This goes to the file")

2. 自动处理异常

@contextmanager
def handle_exception():
    try:
        yield
    except Exception as e:
        print(f"Caught error: {e}")
        raise  # 可选：重新抛出异常

# 使用
with handle_exception():
    1 / 0  # 触发异常但被捕获

contextlib 的其他工具

• closing(thing)：自动调用 thing.close()。

  from contextlib import closing
  with closing(urllib.request.urlopen("http://example.com")) as page:
      html = page.read()
  

• ExitStack：管理多个上下文管理器动态入栈。

  from contextlib import ExitStack
  with ExitStack() as stack:
      files = [stack.enter_context(open(fname)) for fname in filenames]
  

总结

通过 @contextmanager 装饰器，只需编写一个生成器函数即可替代传统的类实现方式，代码更简洁且逻辑更清晰。结合 try...finally 确保资源释放，同时支持异常处理和状态恢复，是 Python 上下文管理的首选方案。