python contextlib模块 实现自定义上下文管理器

contextlib模块工具

自己定义上下文管理器类之前，我们先看看contextlib模块的实用工具。先参照官方文档： https://docs.python.org/zh-cn/3/library/contextlib.html， 除了 前面提到的 redirect_stdout 函数，contextlib 模块中还有一些类和其他函数，使用范围 更广。

closing
如果对象提供了 close() 方法，但没有实现 __enter__/__exit__ 协议，那么可以使用这 个函数构
建上下文管理器。

suppress
构建临时忽略指定异常的上下文管理器。

@contextmanager
这个装饰器把简单的生成器函数变成上下文管理器，这样就不用创建类去实现管理器协 议了。

ContextDecorator
这是个基类，用于定义基于类的上下文管理器。这种上下文管理器也能用于装饰函数， 在受管理的上下文中运行整个函数。

ExitStack
这个上下文管理器能进入多个上下文管理器。with 块结束时，ExitStack 按照后进先出的 顺序调
用栈中各个上下文管理器的 __exit__ 方法。如果事先不知道 with 块要进入多少个 上下文管理
器，可以使用这个类。例如，同时打开任意一个文件列表中的所有文件。

显然，在这些实用工具中，使用最广泛的是 @contextmanager 装饰器，因此要格外留心。 这个
装饰器也有迷惑人的一面，因为它与迭代无关，却要使用 yield 语句。由此可以引出 协程

使用@contextmanager
@contextmanager 装饰器能减少创建上下文管理器的样板代码量，因为不用编写一个完整的 类，定义 __enter__ 和 __exit__ 方法，而只需实现有一个 yield 语句的生成器，生成想让 __enter__ 方法返回的值。

在使用 @contextmanager 装饰的生成器中，yield 语句的作用是把函数的定义体分成两部 分：yield 语句前面的所有代码在 with 块开始时（即解释器调用 __enter__ 方法时）执行， yield 语句后面的代码在 with 块结束时（即调用 __exit__ 方法时）执行。

"""
使用生成器实现的上下文管理器 
"""
import contextlib


# 应用 contextmanager 装饰器
@contextlib.contextmanager
def looking_glass():
    import sys
    # 贮存原来的 sys.stdout.write 方法
    original_write = sys.stdout.write

    # 定义自定义的 reverse_write 函数；在闭包中可以访问 original_write
    def reverse_write(text):
        original_write(text[::-1])

    # 把 sys.stdout.write 替换成 reverse_write
    sys.stdout.write = reverse_write

    # 产出一个值，这个值会绑定到 with 语句中 as 子句的目标变量上。
    # 执行 with 块中的代 码时，这个函数会在这一点暂停。
    yield 'JABBERWOCKY'

    # 控制权一旦跳出 with 块，继续执行 yield 语句之后的代码；
    # 这里是恢复成原来的 sys. stdout.write 方法
    sys.stdout.write = original_write


if __name__ == '__main__':
    with looking_glass() as what:
        print('Alice, Kitty and Snowdrop')
        print(what)

    print(what)

运行结果：

与之前示例唯一的区别是上下文管理器的名字：LookingGlass 变成了looking_glass。

其实，contextlib.contextmanager 装饰器会把函数包装成实现 __enter__ 和 __exit__ 方法 的类。

这个类的 __enter__ 方法有如下作用。
(1) 调用生成器函数，保存生成器对象（这里把它称为 gen）。
(2) 调用 next(gen)，执行到 yield 关键字所在的位置。
(3) 返回 next(gen) 产出的值，以便把产出的值绑定到 with/as 语句中的目标变量上。

with 块终止时，__exit__ 方法会做以下几件事。
(1) 检查有没有把异常传给 exc_type；如果有，调用 gen.throw(exception)，在生成器函数 定义体中包含 yield 关键字的那一行抛出异常。
(2) 否则，调用 next(gen)，继续执行生成器函数定义体中 yield 语句之后的代码。

示例中， 有一个严重的错误：如果在 with 块中抛出了异常，Python 解释器会将其捕获， 然后在 looking_glass 函数的 yield 表达式里再次抛出。但是，那里没有处理错误的代码， 因此 looking_glass 函数会中止，永远无法恢复成原来的 sys.stdout.write 方法，导致系 统处于无效状态。

"""
使用生成器实现的上下文管理器,优化
"""
import contextlib


@contextlib.contextmanager
def looking_glass():
    import sys
    original_write = sys.stdout.write

    def reverse_write(text):
        original_write(text[::-1])

    sys.stdout.write = reverse_write
    # 创建一个变量，用于保存可能出现的错误消息
    msg = ''
    try:
        yield 'JABBERWOCKY'

    # 处理 ZeroDivisionError 异常，设置一个错误消息
    except ZeroDivisionError:
        msg = 'Please DO NOT divide by zero!'
    finally:
        # 撤销对 sys.stdout.write 方法所做的猴子补丁
        sys.stdout.write = original_write
        if msg:
            # 如果设置了错误消息，把它打印出来
            print(msg)


if __name__ == '__main__':
    with looking_glass() as what:
        print('Alice, Kitty and Snowdrop')
        print(what)
    print(what)

运行结果：

前面说过，为了告诉解释器异常已经处理了，__exit__ 方法会返回 True，此时解释器会压制 异常。如果 __exit__ 方法没有显式返回一个值，那么解释器得到的是 None，然后向上冒泡异 常。使用 @contextmanager 装饰器时，默认的行为是相反的：装饰器提供的 __exit__ 方法假定 发给生成器的所有异常都得到处理了，因此应该压制异常。 如果不想让 @contextmanager 压 制异常，必须在被装饰的函数中显式重新抛出异常。

使用@contextmanager 装饰器时，要把yield 语句放在try/finally 语句中 （或者放在 with 语句中），这是无法避免的，因为我们永远不知道上下文管理器 的用户会在 with 块中做什么。

除了标准库中举的例子之外，Martijn Pieters 实现的原地文件重写上下文管理器，是@contextmanager 不错的使 用实例。用法如示例 所示：

# 用于原地重写文件的上下文管理器 

 	
import csv 


with inplace(csvfilename, 'r', newline='') as (infh, outfh): 
	reader = csv.reader(infh) 
	writer = csv.writer(outfh) 
	
	for row in reader: 
		row += ['new', 'columns'] 
		writer.writerow(row)

inplace 函数是个上下文管理器，为同一个文件提供了两个句柄（这个示例中的 infh 和 outfh），以便同时读写同一个文件。这比标准库中的 fileinput.input（这个函数也提供了一个上下 文管理器） 函数易于使用。

注意，在 @contextmanager 装饰器装饰的生成器中，yield 与迭代没有任何关系。在本节所 举的示例中，生成器函数的作用更像是协程：执行到某一点时暂停，让客户代码运行，直 到客户让协程继续做事。