初识with和上下文管理器

1. 上下文管理器是什么？

先简单说一嘴 with 和 上下文管理器之间的关系，with语句块是用来处理 上下文管理器的，就如 for 语句是用来处理迭代器的。

那上下文管理器具体是个啥？说不明白。可以这样理解，实现了 __enter__ 和 __exit__ 方法的类，就是一个上下文管理器。就如实现了 __iter__ 和 __next__ 方法的类，是一个迭代器。再换句话说，一个类若实现了 __enter__ 和 __exit__ 方法，则称这个类实现了上下文管理器协议。

上下文管理器一般用来干什么？

顾名思义，上下文管理器用来管理一段代码的上下文状态。使用上下文管理器可以方便的帮助我们管理，在处理一段代码前进行相应的资源获取、配置设置等，在处理一段代码后进行相应的资源释放、配置写入文件等操作。以打开文件的 open 函数为例。

# 使用 open 函数打开文件，并进行一些处理
f = open(file)     # 使用 open 函数打开一个文件，并将打开的文件绑定到变量 f 上
# 一些数据处理的操作，省略细节
for line in f:
		pass
f.close()          # 调用 close 函数，关闭文件

# 上述代码等价于下面代码
with open(file) as f:
		for line in f:
				pass

如下图分析所示：

上下文管理器中的 __enter__ 和 __exit__ 方法

要想自定义一个上下文管理器，最直接的方法就是定义一个类，该类实现了 __enter__ 和 __exit__ 方法。示例如下：

class ContextManager:
		def __enter__(self):
				pass

		def __exit__（self, exc_type, exc_value, traceback):
				pass 

__enter__ 不需要传入其他参数，在执行 with 语句块内的语句时，会先执行 __enter__ 方法，with 后面的 as 语句（可选）用来接收 __enter__ 方法的返回值。

__exit__ 方法接收以下三个参数，如果在执行 with 语句块内的语句时发生异常，这三个参数用来传递异常数据；如果一切正常（未发生异常），则传入的参数值为 None，None，None；

• exc_type 异常类
• exc_value 异常实例
• traceback traceback 对象

在执行完with语句块内的语句时，会执行 __exit__ 方法，如果 exit 方法返回 None，或者True之外的值，with块中的任何异常都会向上冒泡。

2. 使用上下文管理计算一段代码的运行时长

代码示例如下：

import time

class Timer:

    def __enter__(self):
        self.start_time = time.time()
        return self

    def __exit__(self, exc_type, exc_value, traceback):
        self.total_time = time.time() - self.start_time
        print(f"total_time = {self.total_time}")

def cal(nums):
    ans = 0
    for num in nums:
        ans += num * num + num
    return ans

if __name__ == "__main__":
    import random

    nums = [random.randint(0, 100) for _ in range(1000000)]

    with Timer() as t:
        ans = cal(nums)

    print(f"ans = {ans}")
    
    print(f"total_time = {t.total_time}")

# 运行结果
total_time = 0.04856252670288086
ans = 3399433022
total_time = 0.04856252670288086

几点说明：

• with 语句块内的语句和语句块外的语句具有同等的作用域，因此在with语句块后的语句可以访问with语句块内的变量；
• 在上述示例中，with后的as语句后的变量用来接收 __enter__ 的返回值，上述示例中 __enter__ 的返回值是 Timer 的类对象，因此可以使用 t.total_time 访问对象属性；

显示地调用 __enter__ 和 __exit__ 方法，可以发现使用with语句管理上下文管理器非常方便；

if __name__ == "__main__":
    import random

    nums = [random.randint(0, 100) for _ in range(1000000)]

    t = Timer()
    t.__enter__()

    ans = cal(nums)

    t.__exit__(None, None, None)

    print(f"ans = {ans}")
    
    print(f"total_time = {t.total_time}")

3. 使用 @contextmanager 创建上下文管理器

使用 contextlib.contextmanager 装饰器可以方便地创建一个上下文管理器，使用该装饰器实现上述 Timer 类相似的功能；

from contextlib import contextmanager

@contextmanager
def timer():
    import time
    start_time = time.time()
    yield 
    total_time = time.time() - start_time
    print(f"total_time = {total_time}")

if __name__ == "__main__":
    import random

    nums = [random.randint(0, 100) for _ in range(1000000)]

    with timer():
        ans = cal(nums)

    print(f"ans = {ans}")

如何使用 @contextmanager 装饰器创建上下文管理器

如上述实例所示，在使用 @contextmanager 装饰器创建上下文管理器时，需要实现一个有 yield 语句的生成器。在生成器函数中，yield 语句前的所有代码在执行 with 语句块前执行，对应于 __enter__ 方法内的语句；yield 语句后的所有代码在执行完 with 语句块内的所有语句后执行，对应于 __exit__ 方法内的语句。若 yield 语句返回一个变量，则 as 语句后的变量用来接收该 yield 语句返回的变量。

@contextmanager 装饰器实现上下文管理器的原理的简单分析

一个上下文管理器类实现了 __enter__ 和 __exit__ 方法，@contextmanager装饰器实现上下文管理器的基本原理是把函数包装成实现了 __enter__ 和 __exit__ 方法的类，这个类的名称为 _GeneratorContextManager 。（源码可以参考 ../lib/contextlib.py ，截取源码片段如下）

def contextmanager(func):
    @wraps(func)
    def helper(*args, **kwds):
        return _GeneratorContextManager(func, args, kwds)
    return helper

class _GeneratorContextManager(
    _GeneratorContextManagerBase,
    AbstractContextManager,
    ContextDecorator,
):
    """Helper for @contextmanager decorator."""

    def __enter__(self):
        del self.args, self.kwds, self.func
        try:
            return next(self.gen)
        except StopIteration:
            raise RuntimeError("generator didn't yield") from None

    def __exit__(self, typ, value, traceback):
        if typ is None:
            try:
                next(self.gen)
            except StopIteration:
                return False
            else:
                raise RuntimeError("generator didn't stop")
        else:
            if value is None:  
                value = typ()
            try:
                self.gen.throw(typ, value, traceback)
            except StopIteration as exc:
                return exc is not value
            except RuntimeError as exc:
                if exc is value:
                    return False
                if (
                    isinstance(value, StopIteration)
                    and exc.__cause__ is value
                ):
                    return False
                raise
            except BaseException as exc:      
                if exc is not value:
                    raise
                return False
            raise RuntimeError("generator didn't stop after throw()")

**_GeneratorContextManager 类的 enter 方法的步骤和作用如下：**

• 首先调用生成器函数，即我们自己实现的被装饰器 @contextmanager 装饰的实现了 yield 的函数，并保存该生成器对象 （self.gen）。这一步的实现是在 _GeneratorContextManager 的父类 _GeneratorContextManagerBase 中实现的，如下给出源码截图。
  

• 调用 next(self.gen) 方法，执行到 yield 关键字所在的位置，返回 yield 关键字产出的值，以便把产出的值绑定到 as 语句中的目标变量中。

**_GeneratorContextManager 类的 enter 方法的步骤和作用如下：**

• 首先检测是否有异常发生，如果有，调用 self.gen.throw(typ, value, traceback) 在生成器函数定义体中 yield 关键字那行抛出异常，（因此在生成器函数中需要进行异常捕获和处理，否则不会执行yield后面的语句（即上下文管理器中实现的 __exit__ 方法），即我们需要保证 __enter__ 方法中能够正确处理异常）；
• 若没有异常发生，则调用 next(self.gen)，继续执行生成器函数中 yield 语句之后的代码。

在《流畅的python》书中，给出了一个很好的例子，这里就不摘录了。

参考资料

• 《流畅的python》
• 《python CookBook》