生成器和协程

Posted by YuanBao on April 4, 2016

在学习 Python 的过程中,总会遇到与其他语言不太相同的元素,yield 关键字就是其中之一。随着 Python 2.5 加入 yield,一直到 PEP 380 为 Python 3.3 加入 yield from,这个关键字的能力在 Python 中不断地增强。今天就介绍一下强大的 yield 究竟是什么,为什么使用以及究竟该怎么用。

生成器

提到 yield 就不得不首先讲到生成器。生成器的概念在 Python 2.2 中首次出现,之所以引入生成器,是为了实现一个在计算下一个值时不需要浪费空间的结构。这么说可能不好理解,但是我们可以通过一个例子来说明问题。

假如你需要返回一个从 1 到 $N$ 的列表,从而将其作为自己的 range 函数,那么你在 Python 中很可能会这么写:

def my_range(N):
    index = 0
    my_list = []
    while index < N:
        my_list.append(index)
        index += 1
    return my_list

当然,上面这段代码看起来毫无问题也能满足要求。但是,如果你需要一个非常非常大的列表,使得 $N$ 值设定的非常大。使用上面的 my_range 返回的结果可能会超出计算机的内存,从而使其不能成功运行。

为了避免超出内存的问题,我们需要一个更加智能的 return 功能,它能够每次返回一个结果,并且在函数的内部记录返回结果时的状态,使得下一次能够在此基础上运行。这样的话,我们就可以在节省内存的前提下实现上述的 my_range 函数。yield 关键字的引入正是为了解决这种问题。

我们首先看一下如何使用 yield 关键字重写上面的函数:

def my_range_2(N):
    index = 0
    print "Start to yield!"
    while index < N:
        yield index
        print "Now index is %d" % index
        index += 1
r = my_range_2(5)
print r.next()

执行结果如下:

############################
Start to yield!
0

我们在函数体中使用 yield 关键字返回了一个结果,根据 Python 的实现,任何一个在函数体内使用了 yield 的函数都将变成一个生成器。我们可以将这个生成器理解成为迭代器,因为它完全实现了迭代器的功能。当我们调用 r=my_range_2(5) 的时候,函数 my_range_2(5) 并没有直接执行,而是返回了一个迭代器(注意这个我们传统的 C-like 语言完全不同)。当调用 r.next() 的时候,迭代器中的代码(也就是函数中的代码)开始向前顺序运行,并且在执行完 yield 语句之后暂停(注意是执行完 yield 语句之后),而 yield 语句的返回值就是其后面表达式的值(这里就是 index )。如果我们继续调用 next 语句,函数将会从暂停的位置继续运行,直到再次执行完 yield 语句之后停止。

尽管我们可以使用 next 来一步一步驱动生成器的执行,但是在大多数情况,我们都是依赖于循环语句或者支持迭代器的语句来使用它,例如:

for index in my_range_2(5):
    print index
total = sum(my_range_2(5))

注意在 Python 中生成器的应用非常频繁,Python 也天然的具有很多支持迭代器的语法。下面我们通过一个简单的例子来说明生成器便捷之处。

假如我们有一个整数序列,需要将序列中重复的数字删除掉,并保持剩下元素的次序不变,我们就可以使用一个简单的生成器来完成,代码如下:

def delete_dup(l):
    items = set()
    for num in l:
        if num not in items
            yield num
            items.add(num)
l = [1,5,2,1,9,1,5,10]
print list(delete_dup(l))

###########################
[1,5,2,9,10]

协程

通过上面的介绍,我们直到生成器为我们引入了暂停函数执行的功能。当有了暂停的功能之后,人们就想能不能在生成器暂停的时候向其发送一点东西。这种向暂停的生成器发送信息的功能通过 PEP 342 进入 Python 2.5 中,并催生了 Python 中协程的诞生。根据 wikipedia 中的定义,

协程是为非抢占式多任务产生子程序的计算机程序组件,协程允许不同入口点在不同位置暂停或开始执行程序。

注意从本质上而言,协程并不属于语言中的概念,而是编程模型上的概念。协程不同于线程,其不会像线程一样并行的在多个 CPU 核心上执行。多个协程之间只会交叉串行执行,这决定了使用协程可以避免多线程编程所带来的复杂的加锁问题。通过使用 yield 暂停生成器,可以将程序的执行流程交给其他的子程序,从而实现不同子程序的之间的交替执行。下面的程序简要的演示了,如何向生成器发送数据:

def jumping_range(N):
    index = 0
    while index < N:
    jump = yield index
        if jump is None:
            jump = 1
        index += jump

if __name__ == '__main__':
    itr = jumping_range(5)
    print(itr.next())
    print(itr.send(2))
    print(itr.next())
    print(itr.send(-1))
###################################
0
2
3
2

在这里,协程执行完 yield index 之后可以获取从外部程序发送过来的值,并存入到 jump 中。其中 itr.send(value) 实现了发送数据的功能。与调用 next 函数一样,调用 send 函数也会驱动协程继续执行到下一个 yield 语句之后。事实上 itr.next() 的作用等同于调用 itr.send(None)

需要注意的是,每一个协程初始时都需要使用 next 语句驱动一次,在上述代码11行的位置处的 next 函数驱动了该协程的第一次执行。同时,不同于上一小节中生成器的使用,协程是永远存在并且无限期的,当不需要使用的时候,需要调用 close 函数进行关闭。

在理解了 Python 中上面的代码之后,我们可以使用协程来编写一个简单的生产者消费者程序,这个程序演示了两个子程序是如何串行的交替执行的:

def consumer():
    n = 0
    while True:
        n = yield n
        if n is not None and n > 0:
            n -= 1
            print "Consume 1 object, %d remained!" % n

def producer(c):
    n = 0
    c.next()
    while n < 6:
        n += 2
        print "Produce 2 objects, %d remained!" % n
        c.send(n)
    c.close()

if __name__ == '__main__':
    c = consumer()
    producer(c)        

执行之后可以看到如下结果:

Produce 2 objects, 2 remained!
Consume 1 object, 1 remained!
Produce 2 objects, 4 remained!
Consume 1 object, 3 remained!
Produce 2 objects, 6 remained!
Consume 1 object, 5 remained!

yield from 语句

PEP 380 为 Python 3.5 引入了 yield from 语句。这允许我们从从迭代器(生成器刚好也是迭代器)中返回任何值,从而可以干净利索的方式重构生成器

这里需要解释一下什么是”重构迭代器”。举个简单的例子,我们知道每个迭代器都可以在循环的驱动下产生一个序列,如果我们有一个迭代器列表(一个 list,其中每一个元素都是一个迭代器)。我们现在要把所有的迭代器连接起来,生成一个大的迭代器,那么我们应该怎么办呢?

有了 yield from 语句,这项工作完成起来就非常简单了,我们可以这么写:

def concate_iterators(iterators):
    for itr in iterators:
        yield from iterators

其中 iterators 是一个迭代器列表。当然你也可以使用 itertools 模块中的 chain 函数实现上述功能。有了这个简单的例子,我想应该能够明白 yield from 的作用了,它可以将迭代器的功能转移到另一个迭代器,从而能够实现迭代器之间的各种组合包装。

需要说明的是,在 asyncio 模块中,yield from 关键字的使用非常普遍(后续文章中介绍)。

Python的协程是怎么工作的

我想大部分人都能理解函数调用的过程,以下面的代码为例:

def foo():
    bar()
    
def bar():
    pass

当函数 foo 调用子例程 bar 时,bar 将会保持对代码流的控制,直到其返回函数值或者抛出一个异常。需要明白的是,标准的 Python 解释器使用 C 语言编写,在解释器中负责执行 Python 中函数的 C 函数名为 PyEval_EvalFrameEx ,它获取一个栈帧对象,并且在这个栈帧中执行 Python 程序的字节码。我们将 foo 的字节码展示出来如下:

>>> import dis
>>> dis.dis(foo)
  2           0 LOAD_GLOBAL              0 (bar)
              3 CALL_FUNCTION            0 (0 positional, 0 keyword pair)
              6 POP_TOP
              7 LOAD_CONST               0 (None)
             10 RETURN_VALUE

从这段字节码中可以看到的是,foo 函数加载 bar 函数到它的栈上并且调用该函数,之后将它的返回值从栈中弹出,加载一个常量值 None 作为自身的返回值并且返回。

当函数 PyEval_EvalFrameEx 遇到指令 CALL_FUNCTION 时,它新建一个栈帧,并且递归的调用 PyEval_EvalFrameEx, 这样 PyEval_EvalFrameEx 就能够在新的栈帧中执行函数 bar 的字节码。

最关键的地方来了,Python 的栈帧并不是真正存在于栈上,而是存在于堆上。事实上使用 C 语言编写的 Python解释器的栈确实是运行在操作系统分配的栈上,然而 Python 语言中的栈帧却是分配在堆上。这意味着什么?这意味着即使 Python 函数的调用结束了,我们仍然可以持有函数调用时分配的栈帧。我们使用如下交互的方式来展示这一点:

>>> import inspect
>>> frame = None
>>> def foo():
...     bar()
...
>>> def bar():
...     global frame
...     frame = inspect.currentframe()
...
>>> foo()
>>> # The frame was executing the code for 'bar'.
>>> frame.f_code.co_name
'bar'
>>> # Its back pointer refers to the frame for 'foo'.
>>> caller_frame = frame.f_back
>>> caller_frame.f_code.co_name
'foo'

现在你可能有一点明白了,事实上 coroutine 有跟子例程一样的代码对象和栈帧,并且其栈帧同样存在于堆上。假如我们有如下的生成器代码:

>>> def gen_fn():
...     result = yield 1
...     print('result of yield: {}'.format(result))
...     result2 = yield 2
...     print('result of 2nd yield: {}'.format(result2))
...     return 'done'
...     

当 Python 编译函数 gen_fn 时看到关键字 yield, 它将意识到该函数需要当做生成器函数来处理,因此它设置一个标志来记住该事实:

>>> # The generator flag is bit position 5.
>>> generator_bit = 1 << 5
>>> bool(gen_fn.__code__.co_flags & generator_bit)
True

当调用一个生成器函数时,并不直接执行函数中的代码,而是返回一个生成器类型:

>>> gen = gen_fn()
>>> type(gen)
<class 'generator'>

一个 Python 的生成器类型同时封装了一个栈帧和一个指向生成器代码的指针。所有的对于函数 gen_fn 的调用共享一个代码指针,但是却各自拥有各自的栈帧。前面已经提到,该栈帧并没有真正在栈上,而是分配在堆上。

每一个栈帧都拥有一个指向上一条执行指令的指针 f_lasti,初始的时候,f_lasti 的值为 -1, 意味着该生成器还没有开始执行:

>>> gen.gi_frame.f_lasti
-1

当我们执行 gen.send(None) 的时候,生成器开始执行到第一条 yield 语句处,停止并且返回值 1。这个时候我们打印 f_lasti 的值,发现指令已经执行到第三条 bytecode 处:

>>> gen.gi_frame.f_lasti
3
>>> len(gen.gi_code.co_code)
56

由于该生成器的栈帧并没有真正存在栈上,其可以在任何函数中随时地被恢复。它的栈帧并不像函数调用一样具有固定的层次化,并且不满足函数调用时先进先出的顺序。它是完全独立的,可以存在于堆上的任何一块空间内。如果我们继续向该生成器发送数据,

>>> gen.send('hello')
result of yield: hello
2

现在我们可以看到生成器的栈帧中包含着局部变量 result

>>> gen.gi_frame.f_locals
{'result': 'hello'}

每一个有 函数 gen_fn 创建的生成器都具有自己的栈帧和自己的局部变量。接下来我们继续驱动生成器运行,将会看到生成器抛出一个异常:

>>> gen.send('goodbye')
result of 2nd yield: goodbye
Traceback (most recent call last):
  File "<input>", line 1, in <module>
StopIteration: done

这个异常返回了生成器的返回值: “done”.