Question

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

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

下面举个例子。示例 15-5 使用一个生成器函数代替示例 15-3 中定义的 LookingGlass 类。

示例 15-5　mirror_gen.py：使用生成器实现的上下文管理器

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  ➍
  yield 'JABBERWOCKY'  ➎
  sys.stdout.write = original_write  ➏

❶ 应用 contextmanager 装饰器。

❷ 贮存原来的 sys.stdout.write 方法。

❸ 定义自定义的 reverse_write 函数；在闭包中可以访问 original_write。

❹ 把 sys.stdout.write 替换成 reverse_write。

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

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

示例 15-6 是使用 looking_glass 函数的例子。

示例 15-6　测试 looking_glass 上下文管理器函数

  >>> from mirror_gen import looking_glass
  >>> with looking_glass() as what:  ➊
  ...    print('Alice, Kitty and Snowdrop')
  ...    print(what)
  ...
  pordwonS dna yttiK ,ecilA
  YKCOWREBBAJ
  >>> what
  'JABBERWOCKY'

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

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

5类的名称是 _GeneratorContextManager。如果想了解具体的工作方式，可以阅读 Python 3.4 发行版中 Lib/contextlib.py 文件里的源码。

这个类的 __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 语句之后的代码。

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

示例 15-7 添加了一些代码，特别用于处理 ZeroDivisionError 异常；这样，在功能上它就与示例 15-3 中基于类的实现等效了。

示例 15-7　mirror_gen_exc.py：基于生成器的上下文管理器，而且实现了异常处理——从外部看，行为与示例 15-3 一样

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'
  except ZeroDivisionError:  ➋
    msg = 'Please DO NOT divide by zero!'
  finally:
    sys.stdout.write = original_write  ➌
    if msg:
      print(msg)  ➍

❶ 创建一个变量，用于保存可能出现的错误消息；与示例 15-5 相比，这是第一处改动。

❷ 处理 ZeroDivisionError 异常，设置一个错误消息。

❸ 撤销对 sys.stdout.write 方法所做的猴子补丁。

❹ 如果设置了错误消息，把它打印出来。

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

6把异常发给生成器的方式是使用 throw 方法，参见 16.5 节。

7这样约定的原因是，创建上下文管理器时，生成器无法返回值，只能产出值。不过，现在可以返回值了，如 16.6 节所述。届时你会看到，如果在生成器中返回值，那么会抛出异常。

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

8这条提示直接引用 Leonardo Rochael 的评论，他是本书的技术审校之一。说得好，Leo ！

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

示例 15-8　用于原地重写文件的上下文管理器

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 函数；顺便说一下，这个函数也提供了一个上下文管理器）易于使用。

如果想学习 Martijn 实现 inplace 的源码（列在这篇文章中），找到 yield 关键字，在此之前的所有代码都用于设置上下文：先创建备份文件，然后打开并产出 __enter__ 方法返回的可读和可写文件句柄的引用。yield 关键字之后的 __exit__ 处理过程把文件句柄关闭；如果什么地方出错了，那么从备份中恢复文件。

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