Python面试题：在 Python 中，如何实现上下文管理器（context manager）？

在 Python 中，实现上下文管理器（context manager）有两种常见的方法：使用类和使用装饰器（contextlib 模块中的 contextmanager 装饰器）。上下文管理器用于管理资源，例如文件、网络连接等，确保在使用资源时，资源能够被正确地获取和释放。

使用类实现上下文管理器

要使用类实现上下文管理器，需要实现两个特殊方法：__enter__ 和 __exit__。

示例代码

class MyContextManager:
    def __enter__(self):
        # 获取资源的代码
        print("Entering the context")
        return self

    def __exit__(self, exc_type, exc_value, traceback):
        # 释放资源的代码
        print("Exiting the context")
        if exc_type:
            print(f"An exception occurred: {exc_value}")
        return True  # 如果要抑制异常传播，返回 True，否则返回 False 或省略

# 使用上下文管理器
with MyContextManager() as manager:
    print("Inside the context")
    # 如果这里发生异常，__exit__ 也会被调用

使用 contextlib 模块中的 contextmanager 装饰器

使用 contextlib 模块中的 contextmanager 装饰器可以更简洁地实现上下文管理器。这种方法使用生成器来管理资源。

示例代码

from contextlib import contextmanager

@contextmanager
def my_context_manager():
    try:
        # 获取资源的代码
        print("Entering the context")
        yield
    finally:
        # 释放资源的代码
        print("Exiting the context")

# 使用上下文管理器
with my_context_manager():
    print("Inside the context")
    # 如果这里发生异常，finally 块也会被执行

示例解释

1. 使用类实现上下文管理器：

   • __enter__ 方法在进入 with 语句块时被调用。它负责获取资源，并且可以返回一个对象，该对象可以在 with 语句块中使用。
   • __exit__ 方法在退出 with 语句块时被调用。它负责释放资源。如果 with 语句块中发生异常，异常信息会被传递给 __exit__ 方法。

2. 使用 contextlib 模块中的 contextmanager 装饰器：

   • 使用 @contextmanager 装饰器定义一个生成器函数。
   • yield 语句之前的代码在进入上下文时执行，yield 语句之后的代码在退出上下文时执行（无论是否发生异常）。

完整示例

使用类实现文件上下文管理器

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_value, traceback):
        self.file.close()

# 使用文件上下文管理器
with FileManager('example.txt', 'w') as f:
    f.write('Hello, world!')

使用 contextlib 实现文件上下文管理器

from contextlib import contextmanager

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

# 使用文件上下文管理器
with file_manager('example.txt', 'w') as f:
    f.write('Hello, world!')

通过这两种方法，可以方便地管理资源，确保资源在使用后能够被正确地释放。