Python上下文管理器with块及@contextmanager的用法

上下文管理器和with块

上下文管理器对象存在的目的是为了管理with块，就像迭代器的存在是为了管理for循环一样

with 语句存在的意义是对一些常用的 try/finally 结构予以简化。这种结构能够保障一段代码在运行完成后实施某项操作，就算该段代码因为 return 语句、异常或者 sys.exit() 调用而停止，也依旧执行指定的操作。finally 子句中的代码往往用于释放重要的资源，亦或是还原临时变动的状态。

上下文管理器接口包含了__enter__以及__exit__这两个方法。当 with 语句开始执行的时候，Python 会在上下文管理器对象上启用__enter__方法。待 with 块运行完毕，或者因某些缘由终止之后，Python 会在上下文管理器对象上调用__exit__方法。

• __enter()__方法：当进入 with 语句时，__enter__() 方法会被调用。该方法通常用于执行一些准备工作，并返回一个资源对象，这个对象会被绑定到 as 后的变量名上。
• __exit()__方法：当 with 语句块执行结束时，无论是否发生了异常，__exit__() 方法都会被调用。该方法通常用于执行一些清理操作，比如关闭文件、释放资源等。如果在 with 语句块中发生了异常，异常信息会作为参数传递给 __exit__() 方法。

下面是一个精心设计的上下文管理器执行操作案例

import sys


class LookingGlass:
    def __enter__(self):
        # 将原始的 sys.stdout.write 函数保存起来
        self.original_wirte = sys.stdout.write
        # 将 sys.stdout.write 函数替换为 reverse_write 方法
        sys.stdout.write = self.reverse_wirte
        # 返回一个字符串作为上下文管理器的上下文对象
        return 'QWERTYUIOP'

    def reverse_wirte(self, text):
        # 对传入的文本进行反转后输出
        return self.original_wirte(text[::-1])

    def __exit__(self, exc_type, exc_val, exc_tb):  # 如果一切正常，那么Python 调⽤__exit__ ⽅法时传⼊的参数是None, None, None
        # 在退出上下文管理器时，恢复原始的 sys.stdout.write 函数
        sys.stdout.write = self.original_wirte
        # 如果发生的异常是 ZeroDivisionError，输出相应提示信息
        if exc_type is ZeroDivisionError:
            print('Please DO NOT divide by zero!')
            return True  # 返回True，告诉解释器异常已处理


if __name__ == '__main__':
    with LookingGlass() as what:
        # print(obj)实质就是调用sys.stdout.write(obj+’\n’)
        print("床前明月光")  # 光月明前床
        print("疑是地上霜")  # 霜上地是疑
        print(what)  # POIUYTREWQ
    print(what)  # QWERTYUIOP 处理上下文管理器，输出的文本正常，不再翻转

解释器调用__enter__方法时，除了隐式的self，不需要传入其他参数，传給__exit__方法的三个参数如下：

• exc_type:异常类。
• exc_value:异常实例。有时候会有参数传递给异常构造函数，例如错误消息，这些参数能够通过 exc_value.args获取。
• exc_tb:trackback对象。

@contextmanager的用法

@contextmanager是一个十分实用且巧妙的工具。他将Python3个功能结合在一起：函数装饰器，生成器以及上下文管理器。

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

在使用@contextmanager装饰的生成器中，yiled把函数主体分成两部分，yield前面的代码在with块开始时（解释器调用__enter__方法时）执行。yield后面的代码在with块结束时（解释器调用__exit__方法时）执行。

继续上面的示例：

import sys
import contextlib


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

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

    sys.stdout.write = reverse_write
    msg = ''
    # 这里处理异常是因为如果with块有异常抛出，Python会将其捕获，然后再looking_glass函数的yield表达式再次抛出
    # 但是那里又没有处理错误的代码，因此looking_glass函数将终止执行，永远无法恢复成原来的sys.stdout.write方法，导致系统处于无效状态。
    try:
        yield 'QWERTYUIOP'
    except ZeroDivisionError:
        msg = 'Please DO NOT divide by zero!'
    finally:
        sys.stdout.write = original_text
        if msg:
            print(msg)


if __name__ == '__main__':
    with looking_glass() as what:
        print("床前明月光")  # 光月明前床
        print("疑是地上霜")  # 霜上地是疑
        print(what)  # POIUYTREWQ
    print(what)  # QWERTYUIOP

其实@contextlib.contextmanager把函数包装成了实现__enter__和__exit__的类。

那个类的__enter__方法：

1. 调用生成器函数，获取生成器对象，姑且称之为gen。
2. 调用next(gen)，驱动生成器对象执行到yield关键字位置。
3. 返回next(gen)产生的值，以便把产生的值绑定到as语句的目标变量上。

with语句终止，__exit__方法有以下作用

1. 检查有没有异常传给exc_type，如果有则调用gen.throw(exception) ，在生成器函数主体的yield关键字所在的行抛出异常。
2. 否则调用next(gen)，恢复执行生成器函数主体中yiled后面的代码。

因此在使用@contextmanager装饰器是，应该把yield放在try/finally语句中，这是不可避免的，因为我们永远不知道上下文管理器的用户会在with块中做什么。

@contextmanager还有一个鲜为人知的功能：它装饰的生成器也可以用做装饰器。例如：

@looking_glass()
def varse():
    print("举头望明月低头思故乡")

if __name__ == '__main__':
    varse() # 乡故思头低月明望头举