生成器进化到协程 Part 2

• Part 1 传送门
• David Beazley 的博客
• PPT 下载地址

在 Part 1 我们已经介绍了生成器的定义和生成器的操作，现在让我们开始使用生成器。Part 2 主要描述了如何使用 yield 和 contextmanager 创建一个上下文管理器，并解释了原理。

---

上下文管理器

理解上下文可以联想我们做阅读理解时要解读文章某处的意思需要阅读该处前后段落，正是前后文提供了理解的“背景”。而程序的运行的上下文也可以理解为程序在运行时的某些变量，正是这些变量构成了运行环境，让程序可以完成操作。

Python 中的上下文管理器提供这样一种功能，为你的程序运行时提供一个特定“空间”，当进入这个空间时 Python 上下文管理器 为你做一些准备工作。这个“空间”中一般含有特殊的变量，你在这个“空间”中进行一些操作，然后离开。在你离开时 Python 上下文管理器又会帮你做一些收尾工作，保证不会污染运行环境。

下面是一些常见的代码模式

# 读取文件
f = open()
# do something
f.close()


# 使用锁
lock.acquire()
# do somethin
lock.release()


# 进行数据库操作
db.start_transaction()
# do something
db.commit()


# 对某段代码进行计时
start = time.time()
# do something
end = time.time()

这些代码进行的都是“先做这个（准备工作，比如获取一个数据库连接），然后做这个（比如写入数据），最后整理工作环境（如提交改动，关闭链接，释放资源等）。

如果使用 with 可以这样写：

witn open(filename) as f:
    # do something
    pass
    
with lock():
    # do something
    pass

with 语句实际上使用了实现了 __enter__ 和 __exit__ 方法的上下文管理器类。一个典型的上下文管理器类如下：

clss ContextManager:
    def __enter__(self):
        return value
    def __exit__(self, exc_type, val, tb):
        if exec_type is None:
            return
        else:
            # 处理异常
            return True if handled else False

正如方法名明确告诉我们的，__enter__ 方法负责进入上下的准备工作，如果有需要可以返回一个值，这个值将会被赋值给 with ContextManager() as ret_value 中的 ret_value 。__exit__ 则负责收尾工作，这包括了异常处理。

对于这样一段代码

with ContextManager() as var:
    # do something

相当于

ctxmanager = ContextManager()
var = ctxmanager.__enter__()
# do somethin
ctxmanager.__exit__()

一个可用的例子：

import tempfile
import shutil

class TmpDir:
    def __enter__(self):
        self.dirname = tempfile.mkdtemp()
        return self.dirname
    
    def __exit__(self, exc, val, tb):
    shutil.rmtree(self.dirname)

这个上下文管理提供临时文件的功能，在 with 语句结束后会自动删除临时文件夹。

with TempDir() as dirname:
    # 使用临时文件夹进行一些操作
    pass

关于上面两个特殊方法的文档可以在 Python 文档的 Context Manager Types 找到。另外关于 with 关键字的详细说明参考 PEP 343，不过这篇 PEP　不是很好读，Good Luck :simple_smile:！

使用 yield 和 contextmanager

能看到这里的都应该对上下文管理器有所了解，准备好把 yield 加入我们的上下文管理器代码中。

先看一个例子

import tempfile, shutil
from contextlib import contextmanager

@contextmanager
def tempdir():
    outdir = tempfile.mkdtemp()
    try:
        yield outdir
    finally:
        shutil.rmtree(outdir)

与使用上下文管理器类的实现方式不同，这里我们没有显式实现 __enter__ 和 __exit__，而是通过 contextmanager 装饰器和 yield 实现，你可以试试这两种方式是等价的。

要理解上面的代码，可以把 yield 想象为一把剪刀，把这个函数一分为二，上部分相当于 __enter__，下部分相当于 __exit__。我这样说大家应该明白了吧。

import tempfile, shutil
from contextlib import contextmanager

@contextmanager
def tempdir():
    outdir = tempfile.mkdtemp() #
    try:                        # __enter__
        yield outdir            #
--cut---╳-----------------------------------
    finally:                    #
        shutil.rmtree(outdir)   # __exit__

实现“剪刀”功能关键在于 contextmanager 。对于上面的代码，我们来一步一步地结构它：

contextmanager 装饰器

contextmanager 其实使用了一个上下文管理器类，这个类在在初始化时需要提供一个生成器。

class GeneratorCM:
    def __init__(self, gen):
        self.gen = gen
    
    def __enter__(self):
       ...
        
    def __exit__(self, exc, val, tb):
        ...

contextmanager 的实现如下

def contextmanager(func):
    def run(*args, **kwargs):
        return GeneratorCM(func(*args, **kwargs))
    return run

由于 contextmanger 所装饰的函数里有 yield 所以我们在调用 func(*args, **kwargs) 时返回的是一个生成器。要使这个生成器前进，我们需要调用 next 函数

让生成器前进

def __enter__(self):
    return next(self.gen)

GeneratorCM 的 __ente__ 方法会让生成器前进到 yield 语句处，并返回产出值。

收尾

def __exit__(self, exc, val, tb):
    try:
        if exc is None:
            next(self.gen)
        else:
            self.gen.throw(exc, val, tb)
        raise RuntimeError('Generator didn\'t stop')
    except StopIteration:
        return True
    except:
        if sys.exc_info()[1] is not val: raise

__exit__ 函数的逻辑比较复杂，如果没有传入异常，首先它会尝试对生成器调用 next，正常情况下这会抛出 StopIteration ，这个异常会被不做并返回 True ，告诉解释器正常退出；如果传入异常，会使用 throw 在 yield 处抛出这个异常；如果有其他未捕捉的错误，就重新抛出该错误。

实际的代码实现会更加复杂，还有一些异常情况没有处理

• 没有相关值的异常
• 在 with 语句块中抛出的 StopIteration
• 在上下文管理器中抛出的异常

如果你对怎么实现感兴趣，你可以阅读代码或者再一次阅读 PEP 343。

总结

Part 2 都是关于上下文管理器的内容，与协程关系不大。但通过这部分我们可以看到 yield 完全不同的用法，也熟悉了控制流 (control-flow) ，这与 Part 3 的异步处理流程有很大关系。让我们 Part 3 再见。