Python 中的生成器实现原理

1. 如何生成一个巨大的序列

1.1 需求描述

要求生成一个包含很多元素的序列,假设:

  • 存储 1 个整数需要 4 个字节
  • 现在要创建一个包含 1 G 个整数的序列,从 0 到 1 * 1024 * 1024 * 1024 - 1
  • 如果需要为序列中的每个整数分配内存,则需要分配的内存为 1G * 4 = 4G

1.2 通过列表推导

Python 提供了列表推导用于生成列表,下面使用列表推导生成一个包含 0 到 4 之间所有整数的列表,代码如下:

>>> list = [i for i in range(4)]
>>> list
[0, 1, 2, 3]
  • 在第 1 行,使用列表推导创建一个包含 4 个元素的列表
  • 在第 2 行,显示新创建的列表
  • 在第 3 行,创建了一个包含 0、1、2、3 等 4 个元素的列表

如果生成一个从 0 到 1G 的列表,代码如下:

>>> N = 1024 * 1024 * 1024
>>> list = [i for i in range(N)]
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "<stdin>", line 1, in <listcomp>
MemoryError
  • 在第 1 行,设定 N 为 1G
  • 在第 2 行,使用列表推导创建一个包含 N 个元素的列表
  • 在第 6 行,程序运行出错,提示 MemoryError

使用列表推导创建包含 1G 个整数的列表时,需要为这 1G 个整数分配至少 4G 的内存,需要消耗大量的内存,超出了 Python 的限制,因此出现了 MemoryError 的错误。

另外,创建这个巨大的列表需要消耗大量的时间,因此执行第 2 行的语句后,系统失去响应,大约 10 多秒后才出现错误信息。

1.3 通过动态计算

列表推导需要一次性的为 1G 个整数分配内存空间,带来了两个问题:

  1. 列表占用了大量的物理内存
  2. 创建列表的时间过长

Python 提供了一种动态计算的思路解决以上问题,它的思想如下:

  1. 要生成的序列是有规则的,在这个例子中,要求生成连续递增的序列
  2. 使用一个特殊的对象 generator,该对象被称为生成器 generator,生成器按照规则依次输出该序列
  3. Python 提供了内置方法 next(generator),该方法通知生成器产生下一个数据并返回该数据
  4. 不需要为 generator 预先分配内存,通过调用 next(generator) 可以动态获取序列的下一个数据

创建一个输出从 0 到 1G 的生成器,代码如下:

>>> N = 1024 * 1024 * 1024
>>> generator = (i for i in range(N))
>>> next(generator)
0
>>> next(generator)
1
>>> next(generator)
2
  • 在第 1 行,设定 N 为 1G
  • 在第 2 行,使用类似于列表推导的语法创建一个生成器,它输出从 0 到 1G 的序列
    • 注意:创建生成器的语法采用小括号 (),创建列表的语法采用方括号 []
  • 在第 3 行,使用 next(generator),通知 generator 生产一个数据
    • 在第 4 行,generator 输出从 0 到 1G 序列中的第 0 个整数
  • 在第 5 行,使用 next(generator),通知 generator 生产一个数据
    • 在第 6 行,generator 输出从 0 到 1G 序列中的第 1 个整数
  • 在第 7 行,使用 next(generator),通知 generator 生产一个数据
    • 在第 8 行,generator 输出从 0 到 1G 序列中的第 2 个整数

注意:在第 2 行,创建一个输出从 0 到 1G 的序列的生成器,因为不需要分配内存,创建生成器的速度非常快,几乎是瞬间完成的。与之相比,在上一节中创建一个输出从 0 到 1G 的序列的列表,因为需要分配内存,创建列表的速度非常慢,并且导致了 MemoryError。

2. 生成器概述

2.1 生成器的定义

在 Python 中,生成器是一个特殊的对象,它按照一定的规则依次输出数据。Python 的内置函数 next(generator) 通知生成器输出一个新的数据,当生成器输出全部数据后,产生一个特殊的异常 StopIteration,用于标记生成器输出结束。

下面的代码创建一个产生 0 到 3 之间所有整数的生成器:

>>> generator = (i for i in range(3))
>>> next(generator)
0
>>> next(generator)
1
>>> next(generator)
2
>>> next(generator)
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
StopIteration
  • 在第 1 行,创建一个产生 0 到 3 之间所有整数的生成器
    • 注意:创建生成器的语法采用小括号 (),创建列表的语法采用方括号 []
  • 在第 2 行,生成器产生第 0 个整数
  • 在第 4 行,生成器产生第 1 个整数
  • 在第 6 行,生成器产生第 2 个整数
  • 在第 8 行,生成器产生第 3 个整数
    • 在第 11 行,因为生成器生成的序列只包含 3 个整数,此时已经生成全部的整数,因此抛出异常 StopIteration

2.2 使用 while 循环访问生成器

根据生成器的原理,可以循环的调用 next(generator) 输出全部的序列,示例如下:

generator = (i for i in range(3))

while True:
    try:
        item = next(generator)
        print(item)
    except StopIteration:
        break
  • 在第 1 行,创建一个产生 0 到 3 之间所有整数的生成器
  • 在第 3 行,创建一个循环
    • 在第 5 行,调用 next(generator) 通知生成器返回一个数据
    • 在第 7 行,当生成器输出结束后,抛出异常 StopIteration

运行程序,输出结果如下:

0
1
2

2.3 使用 for 循环访问生成器

通常使用 for 循环访问生成器,示例如下:

generator = (i for i in range(3))

for item in generator:
    print(item)
  • 在第 1 行,创建一个产生 0 到 3 之间所有整数的生成器
  • 在第 3 行,使用 for 循环访问生成器

运行程序,输出结果如下:

0
1
2

3. 创建生成器

3.1 通过推导创建生成器

可以使用类似于列表推导的语法创建一个生成器,语法如下:

(expression for i in iterable)

该生成器遍历对象 iterable,依次产生数据 expression,它的工作流程如下:

for i in iterable:
    generate expression

注意:创建生成器的语法与列表推导的语法相似,不同之处在于,创建生成器的语法采用小括号 (),创建列表的语法采用方括号 []

通过推导创建生成器的示例如下:

generator = (i*2 for i in range(5))
for i in generator:
    print(i)
  • 循环变量 i 从 0 变化到 4
  • 生成器每次产生数据 i*2

运行程序,输出结果如下:

0
2
4
6
8

3.2 通过复杂的推导创建生成器

可以使用类似于列表推导的语法创建一个生成器,语法如下:

(expression for i in iterable if condition)

该生成器遍历对象 iterable,如果条件 condition 为真,则产生数据 expression,它的工作流程如下:

for i in iterable:
    if condition:
        generate expression

通过复杂推导创建生成器的示例如下:

generator = (i for i in range(10) if i % 2 == 0)
for i in generator:
    print(i)
  • 循环变量 i 从 0 变化到 9
  • 如果 i % 2 == 0,表示 i 是偶数
  • 生成器每次产生从 0 到 9 之间的偶数

运行程序,输出结果如下:

0
2
4
6
8

3.3 通过 yield 创建生成器

在生成器的生命周期中,生成器根据一定的规则产生一系列的数据,生成器可以使用 yield 关键字产生一个数据。例如,一个生成特定范围内的奇数序列的函数:

def generateOddNumbers(n):
    for i in range(n):
        if i % 2 == 1:
            yield i

generator = generateOddNumbers(10)
for i in generator:
    print(i)
  • 在第 1 行,定义了函数 generateOddNumbers(n),它返回一个生成器,该生成器产生从 0 到 n 范围内的奇数
  • 在第 2 行到第 4 行,使用 for 循环生成从 0 到 n 范围内的奇数
    • 在第 3 行,如果 i % 2 == 1 为真,表示 i 是奇数
    • 在第 4 行,使用 yield i 生成一个数据 i
  • 在第 6 行,generateOddNumbers(10) 返回一个生成器,该生成器产生从 0 到 10 范围内的奇数
  • 在第 7 行,使用 for 循环遍历该生成器

运行该程序,输出如下:

1
3
5
7
9

注意:包含 yield 关键字的函数被称为生成器函数,调用生成器函数会返回一个生成器。在上面的例子中,函数 generateOddNumbers(n) 包含 yield 关键字,是一个生成器函数,它返回一个生成器,该生成器产生从 0 到 n 范围内的奇数。

4. 使用 yield 实现遍历堆栈的生成器

4.1 通过单链表实现堆栈

通过单链表实现堆栈,图示如下:

![图片描述](//img.mukewang.com/wiki/5ea92d460974644d07000133.jpg)

在上图中,每个节点有两个字段: item 和 next,item 用于存储数据,next 指向下一个节点,head 指针指向堆栈的顶部。描述堆栈的 Python 代码如下:

class Node:
    def __init__(self, item):
        self.item = item
        self.next = None

class Stack:
    def __init__(self):
        self.head = None

    def push(self, item):
        node = Node(item)
        node.next = self.head
        self.head = node

stack = Stack()
stack.push('a')
stack.push('b')
stack.push('c')
  • 在第 1 行,定义了类 Node 用于描述链表中的节点
  • 在第 6 行,定义了类 Stack 描述堆栈
    • 在第 8 行,定义了头指针 head,指向链表中的首个节点
    • 在第 10 行,定义了成员方法 push,将元素压如到堆栈中
      • 在第 11 行,创建一个新节点 node
      • 在第 12 行,新节点 node 的 next 指向头结点
      • 在第 13 行,头结点指向新节点
  • 在第 15 行,创建一个对象 stack
  • 在第 16 行到第 18 行,依次压入 3 个元素 ‘a’、‘b’、‘c’

4.2 使用 yield 关键字实现生成器函数

def stackGenerate(stack):
    cursor = stack.head
    while cursor != None:
        yield cursor.item
        cursor = cursor.next
  • 在第 1 行,定义函数 stackGenerate(stack)
    • 该函数包含 yield 关键字,是一个生成器函数,它返回一个生成器
    • 生成器遍历堆栈,按出栈的顺序输出数据
  • 在第 2 行,变量 cursor 指向了当前正在遍历的元素,初始化被设置为链表的头结点
  • 在第 3 行,使用循环遍历堆栈
    • 如果变量 cursor 等于 None,表示已经到达链表的尾部,即遍历完全部的元素了
    • 在第 4 行,使用 yield 输出当前正在遍历的元素
    • 在第 5 行,将 cursor 指向下一个元素

4.3 通过 while 循环遍历堆栈

使用 while 循环显式的使用 next、StopIteration 完成对 stack 的遍历,代码如下:

generator = stackGenerate(stack)
while True:
    try:
        item = next(generator)
        print(item)
    except StopIteration:
        break
  • 在第 1 行,stackGenerate(stack) 返回一个遍历堆栈的生成器
  • 在第 4 行,next(generator) 获取生成器的输出
  • 在第 6 行,当生成器输出结束后,抛出异常 StopIteration

程序依次压入 ‘a’、‘b’、‘c’,遍历时以压入相反的顺序输出,结果如下:

c
b
a

4.4 通过 for … in 循环遍历堆栈

通过 for … in 循环对生成器进行遍历,代码如下:

generator = stackGenerate(stack)
for item in generator:
    print(item)

与上一节的代码相比,代码要简洁很多,程序输出相同的结果如下:

c
b
a