使用 @contextmanager 装饰器的上下文管理器

上下文管理器是处理资源的一种标准方式。Python 中的 with 语句可以方便地使用上下文管理器。下面将介绍如何使用 @contextmanager 装饰器来创建上下文管理器。

1. 什么是上下文管理器

在 Python 中，上下文管理器是一个对象，用于定义 with 语句要执行的代码块的启动和关闭操作。Python 内置了许多上下文管理器，如文件对象、锁对象等。如果要使用自定义资源，则需要实现自己的上下文管理器。

上下文管理器常常用于以下场景：

• 清理资源：如关闭文件、关闭数据库连接、释放锁等。
• 挂起和恢复状态：如保存并恢复变量值。

2. 上下文管理器的实现方式

在 Python 中，上下文管理器有两种实现方式：类实现和生成器实现。

2.1 类实现上下文管理器

要在 Python 中实现一个上下文管理器，需要定义一个类并实现 enter() 和 exit() 方法。这两个方法分别在 with 语句块开始前和结束后被调用。在 enter() 方法中，可以初始化资源并返回资源对象；在 exit() 方法中，可以清理资源。

下面是一个使用类实现上下文管理器的例子，用于计时：

class Timer:
    def __init__(self):
        self.start = None
        self.end = None

    def __enter__(self):
        self.start = time.time()

    def __exit__(self, exc_type, exc_val, traceback):
        self.end = time.time()

    def elapsed_time(self):
        return self.end - self.start

使用该上下文管理器：

with Timer() as timer:
    time.sleep(1)
print(timer.elapsed_time())   # 输出 1 秒

2.2 生成器实现上下文管理器

Python 还提供了使用生成器来实现上下文管理器的方法，这种方法更为简洁。使用 @contextmanager 装饰器来定义生成器，生成器的 yield 语句之前的代码相当于 enter() 方法，yield 语句之后的代码相当于 exit() 方法。

下面是一个使用生成器实现上下文管理器的例子，用于打印函数运行时间：

from contextlib import contextmanager

@contextmanager
def time_it():
    start = time.time()
    yield
    end = time.time()
    print('elapsed time:', end-start)

使用该上下文管理器：

def func():
    time.sleep(1)

with time_it():
    func()   # 输出 elapsed time: 1.0008xxxxxx

3. 总结

使用 @contextmanager 装饰器可以简化上下文管理器的实现方式，使代码更加简洁易懂。注意在上下文管理器中处理异常，以确保资源正确地被清理。