Python 中的生成器實現(xiàn)原理

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

1.1 需求描述

要求生成一個包含很多元素的序列，假設(shè)：

• 存儲 1 個整數(shù)需要 4 個字節(jié)
• 現(xiàn)在要創(chuàng)建一個包含 1 G 個整數(shù)的序列，從 0 到 1 * 1024 * 1024 * 1024 - 1
• 如果需要為序列中的每個整數(shù)分配內(nèi)存，則需要分配的內(nèi)存為 1G * 4 = 4G

1.2 通過列表推導(dǎo)

Python 提供了列表推導(dǎo)用于生成列表，下面使用列表推導(dǎo)生成一個包含 0 到 4 之間所有整數(shù)的列表，代碼如下：

>>> list = [i for i in range(4)]
>>> list
[0, 1, 2, 3]

• 在第 1 行，使用列表推導(dǎo)創(chuàng)建一個包含 4 個元素的列表
• 在第 2 行，顯示新創(chuàng)建的列表
• 在第 3 行，創(chuàng)建了一個包含 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 行，設(shè)定 N 為 1G
• 在第 2 行，使用列表推導(dǎo)創(chuàng)建一個包含 N 個元素的列表
• 在第 6 行，程序運(yùn)行出錯，提示 MemoryError

使用列表推導(dǎo)創(chuàng)建包含 1G 個整數(shù)的列表時，需要為這 1G 個整數(shù)分配至少 4G 的內(nèi)存，需要消耗大量的內(nèi)存，超出了 Python 的限制，因此出現(xiàn)了 MemoryError 的錯誤。

另外，創(chuàng)建這個巨大的列表需要消耗大量的時間，因此執(zhí)行第 2 行的語句后，系統(tǒng)失去響應(yīng)，大約 10 多秒后才出現(xiàn)錯誤信息。

1.3 通過動態(tài)計算

列表推導(dǎo)需要一次性的為 1G 個整數(shù)分配內(nèi)存空間，帶來了兩個問題：

1. 列表占用了大量的物理內(nèi)存
2. 創(chuàng)建列表的時間過長

Python 提供了一種動態(tài)計算的思路解決以上問題，它的思想如下：

1. 要生成的序列是有規(guī)則的，在這個例子中，要求生成連續(xù)遞增的序列
2. 使用一個特殊的對象 generator，該對象被稱為生成器 generator，生成器按照規(guī)則依次輸出該序列
3. Python 提供了內(nèi)置方法 next(generator)，該方法通知生成器產(chǎn)生下一個數(shù)據(jù)并返回該數(shù)據(jù)
4. 不需要為 generator 預(yù)先分配內(nèi)存，通過調(diào)用 next(generator) 可以動態(tài)獲取序列的下一個數(shù)據(jù)

創(chuàng)建一個輸出從 0 到 1G 的生成器，代碼如下：

>>> N = 1024 * 1024 * 1024
>>> generator = (i for i in range(N))
>>> next(generator)
0
>>> next(generator)
1
>>> next(generator)
2

• 在第 1 行，設(shè)定 N 為 1G
• 在第 2 行，使用類似于列表推導(dǎo)的語法創(chuàng)建一個生成器，它輸出從 0 到 1G 的序列
  • 注意：創(chuàng)建生成器的語法采用小括號 ()，創(chuàng)建列表的語法采用方括號 []
• 在第 3 行，使用 next(generator)，通知 generator 生產(chǎn)一個數(shù)據(jù)
  • 在第 4 行，generator 輸出從 0 到 1G 序列中的第 0 個整數(shù)
• 在第 5 行，使用 next(generator)，通知 generator 生產(chǎn)一個數(shù)據(jù)
  • 在第 6 行，generator 輸出從 0 到 1G 序列中的第 1 個整數(shù)
• 在第 7 行，使用 next(generator)，通知 generator 生產(chǎn)一個數(shù)據(jù)
  • 在第 8 行，generator 輸出從 0 到 1G 序列中的第 2 個整數(shù)

注意：在第 2 行，創(chuàng)建一個輸出從 0 到 1G 的序列的生成器，因為不需要分配內(nèi)存，創(chuàng)建生成器的速度非?？?，幾乎是瞬間完成的。與之相比，在上一節(jié)中創(chuàng)建一個輸出從 0 到 1G 的序列的列表，因為需要分配內(nèi)存，創(chuàng)建列表的速度非常慢，并且導(dǎo)致了 MemoryError。

2. 生成器概述

2.1 生成器的定義

在 Python 中，生成器是一個特殊的對象，它按照一定的規(guī)則依次輸出數(shù)據(jù)。Python 的內(nèi)置函數(shù) next(generator) 通知生成器輸出一個新的數(shù)據(jù)，當(dāng)生成器輸出全部數(shù)據(jù)后，產(chǎn)生一個特殊的異常 StopIteration，用于標(biāo)記生成器輸出結(jié)束。

下面的代碼創(chuàng)建一個產(chǎn)生 0 到 3 之間所有整數(shù)的生成器：

>>> 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 行，創(chuàng)建一個產(chǎn)生 0 到 3 之間所有整數(shù)的生成器
  • 注意：創(chuàng)建生成器的語法采用小括號 ()，創(chuàng)建列表的語法采用方括號 []
• 在第 2 行，生成器產(chǎn)生第 0 個整數(shù)
• 在第 4 行，生成器產(chǎn)生第 1 個整數(shù)
• 在第 6 行，生成器產(chǎn)生第 2 個整數(shù)
• 在第 8 行，生成器產(chǎn)生第 3 個整數(shù)
  • 在第 11 行，因為生成器生成的序列只包含 3 個整數(shù)，此時已經(jīng)生成全部的整數(shù)，因此拋出異常 StopIteration

2.2 使用 while 循環(huán)訪問生成器

根據(jù)生成器的原理，可以循環(huán)的調(diào)用 next(generator) 輸出全部的序列，示例如下：

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

while True:
    try:
        item = next(generator)
        print(item)
    except StopIteration:
        break

• 在第 1 行，創(chuàng)建一個產(chǎn)生 0 到 3 之間所有整數(shù)的生成器
• 在第 3 行，創(chuàng)建一個循環(huán)
  • 在第 5 行，調(diào)用 next(generator) 通知生成器返回一個數(shù)據(jù)
  • 在第 7 行，當(dāng)生成器輸出結(jié)束后，拋出異常 StopIteration

運(yùn)行程序，輸出結(jié)果如下：

0
1
2

2.3 使用 for 循環(huán)訪問生成器

通常使用 for 循環(huán)訪問生成器，示例如下：

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

for item in generator:
    print(item)

• 在第 1 行，創(chuàng)建一個產(chǎn)生 0 到 3 之間所有整數(shù)的生成器
• 在第 3 行，使用 for 循環(huán)訪問生成器

運(yùn)行程序，輸出結(jié)果如下：

0
1
2

3. 創(chuàng)建生成器

3.1 通過推導(dǎo)創(chuàng)建生成器

可以使用類似于列表推導(dǎo)的語法創(chuàng)建一個生成器，語法如下：

(expression for i in iterable)

該生成器遍歷對象 iterable，依次產(chǎn)生數(shù)據(jù) expression，它的工作流程如下：

for i in iterable:
    generate expression

注意：創(chuàng)建生成器的語法與列表推導(dǎo)的語法相似，不同之處在于，創(chuàng)建生成器的語法采用小括號 ()，創(chuàng)建列表的語法采用方括號 []。

通過推導(dǎo)創(chuàng)建生成器的示例如下：

generator = (i*2 for i in range(5))
for i in generator:
    print(i)

• 循環(huán)變量 i 從 0 變化到 4
• 生成器每次產(chǎn)生數(shù)據(jù) i*2

運(yùn)行程序，輸出結(jié)果如下：

0
2
4
6
8

3.2 通過復(fù)雜的推導(dǎo)創(chuàng)建生成器

可以使用類似于列表推導(dǎo)的語法創(chuàng)建一個生成器，語法如下：

(expression for i in iterable if condition)

該生成器遍歷對象 iterable，如果條件 condition 為真，則產(chǎn)生數(shù)據(jù) expression，它的工作流程如下：

for i in iterable:
    if condition:
        generate expression

通過復(fù)雜推導(dǎo)創(chuàng)建生成器的示例如下：

generator = (i for i in range(10) if i % 2 == 0)
for i in generator:
    print(i)

• 循環(huán)變量 i 從 0 變化到 9
• 如果 i % 2 == 0，表示 i 是偶數(shù)
• 生成器每次產(chǎn)生從 0 到 9 之間的偶數(shù)

運(yùn)行程序，輸出結(jié)果如下：

0
2
4
6
8

3.3 通過 yield 創(chuàng)建生成器

在生成器的生命周期中，生成器根據(jù)一定的規(guī)則產(chǎn)生一系列的數(shù)據(jù)，生成器可以使用 yield 關(guān)鍵字產(chǎn)生一個數(shù)據(jù)。例如，一個生成特定范圍內(nèi)的奇數(shù)序列的函數(shù)：

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

generator = generateOddNumbers(10)
for i in generator:
    print(i)

• 在第 1 行，定義了函數(shù) generateOddNumbers(n)，它返回一個生成器，該生成器產(chǎn)生從 0 到 n 范圍內(nèi)的奇數(shù)
• 在第 2 行到第 4 行，使用 for 循環(huán)生成從 0 到 n 范圍內(nèi)的奇數(shù)
  • 在第 3 行，如果 i % 2 == 1 為真，表示 i 是奇數(shù)
  • 在第 4 行，使用 yield i 生成一個數(shù)據(jù) i
• 在第 6 行，generateOddNumbers(10) 返回一個生成器，該生成器產(chǎn)生從 0 到 10 范圍內(nèi)的奇數(shù)
• 在第 7 行，使用 for 循環(huán)遍歷該生成器

運(yùn)行該程序，輸出如下：

1
3
5
7
9

注意：包含 yield 關(guān)鍵字的函數(shù)被稱為生成器函數(shù)，調(diào)用生成器函數(shù)會返回一個生成器。在上面的例子中，函數(shù) generateOddNumbers(n) 包含 yield 關(guān)鍵字，是一個生成器函數(shù)，它返回一個生成器，該生成器產(chǎn)生從 0 到 n 范圍內(nèi)的奇數(shù)。

4. 使用 yield 實現(xiàn)遍歷堆棧的生成器

4.1 通過單鏈表實現(xiàn)堆棧

通過單鏈表實現(xiàn)堆棧，圖示如下：

在上圖中，每個節(jié)點(diǎn)有兩個字段: item 和 next，item 用于存儲數(shù)據(jù)，next 指向下一個節(jié)點(diǎn)，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 用于描述鏈表中的節(jié)點(diǎn)
• 在第 6 行，定義了類 Stack 描述堆棧
  • 在第 8 行，定義了頭指針 head，指向鏈表中的首個節(jié)點(diǎn)
  • 在第 10 行，定義了成員方法 push，將元素壓如到堆棧中
    • 在第 11 行，創(chuàng)建一個新節(jié)點(diǎn) node
    • 在第 12 行，新節(jié)點(diǎn) node 的 next 指向頭結(jié)點(diǎn)
    • 在第 13 行，頭結(jié)點(diǎn)指向新節(jié)點(diǎn)
• 在第 15 行，創(chuàng)建一個對象 stack
• 在第 16 行到第 18 行，依次壓入 3 個元素 ‘a(chǎn)’、‘b’、‘c’

4.2 使用 yield 關(guān)鍵字實現(xiàn)生成器函數(shù)

def stackGenerate(stack):
    cursor = stack.head
    while cursor != None:
        yield cursor.item
        cursor = cursor.next

• 在第 1 行，定義函數(shù) stackGenerate(stack)
  • 該函數(shù)包含 yield 關(guān)鍵字，是一個生成器函數(shù)，它返回一個生成器
  • 生成器遍歷堆棧，按出棧的順序輸出數(shù)據(jù)
• 在第 2 行，變量 cursor 指向了當(dāng)前正在遍歷的元素，初始化被設(shè)置為鏈表的頭結(jié)點(diǎn)
• 在第 3 行，使用循環(huán)遍歷堆棧
  • 如果變量 cursor 等于 None，表示已經(jīng)到達(dá)鏈表的尾部，即遍歷完全部的元素了
  • 在第 4 行，使用 yield 輸出當(dāng)前正在遍歷的元素
  • 在第 5 行，將 cursor 指向下一個元素

4.3 通過 while 循環(huán)遍歷堆棧

使用 while 循環(huán)顯式的使用 next、StopIteration 完成對 stack 的遍歷，代碼如下：

generator = stackGenerate(stack)
while True:
    try:
        item = next(generator)
        print(item)
    except StopIteration:
        break

• 在第 1 行，stackGenerate(stack) 返回一個遍歷堆棧的生成器
• 在第 4 行，next(generator) 獲取生成器的輸出
• 在第 6 行，當(dāng)生成器輸出結(jié)束后，拋出異常 StopIteration

程序依次壓入 ‘a(chǎn)’、‘b’、‘c’，遍歷時以壓入相反的順序輸出，結(jié)果如下：

c
b
a

4.4 通過 for … in 循環(huán)遍歷堆棧

通過 for … in 循環(huán)對生成器進(jìn)行遍歷，代碼如下：

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

與上一節(jié)的代碼相比，代碼要簡潔很多，程序輸出相同的結(jié)果如下：

c
b
a