2.31. 迭代器與產生器

for 迴圈所做的工作比它表面上看起來的還要多。本頁將介紹它所依賴的迭代器協定(iterator protocol),以及讓你能建構自己迭代器的 yield 關鍵字。

2.31.1. 迭代器協定

每一個可以被迴圈走訪的物件都實作了兩個方法:

  • __iter__() -- 傳回一個走訪物件項目的迭代器

  • __next__() -- 在迭代器上呼叫,傳回下一個項目,或在沒有更多項目時引發 StopIteration

內建的 iter() 會呼叫 __iter__next() 會呼叫 __next__。手動逐步走訪一個串列:

it = iter([10, 20, 30])
print(next(it))    # 10
print(next(it))    # 20
print(next(it))    # 30
print(next(it))    # raises StopIteration
A for loop calls __iter__ once to get an iterator, then calls __next__ repeatedly until StopIteration ends the loop.

for 是「呼叫一次 __iter__,然後在 __next__ 上不斷迴圈,直到出現 StopIteration」的語法糖。

for x in items: 實際上做的事:

_it = iter(items)
while True:
    try:
        x = next(_it)
    except StopIteration:
        break
    # ... loop body ...

每一個串列、元組、字串、字典、集合、檔案物件與產生器都已經實作了 __iter____next__ -- 這正是它們全都能搭配 for 使用的原因。

2.31.2. yield 與產生器函式

包含 yield 陳述式的函式是產生器函式。呼叫它並不會執行函式主體,而是傳回一個產生器物件(一種迭代器),它每次只執行函式主體中的一個 yield

def count_up_to(n):
    i = 0
    while i < n:
        yield i
        i += 1

for value in count_up_to(3):
    print(value)

輸出:

0
1
2

每一次對 next() 的呼叫都會恢復函式的執行,直到下一個 yield,將該值交給呼叫者,然後在那裡暫停。區域狀態(此例中為 i)會在多次恢復之間被保留。

Sequence diagram with two lifelines (caller and generator body). The caller calls count_up_to(3), which creates a generator without running the body. Each subsequent next() runs the body until the next yield, returns the yielded value, and pauses. The fourth next() falls off the end and raises StopIteration. The variable i is preserved across pauses.

next() 會執行函式主體直到下一個 yield,把值交回,然後暫停。區域狀態會在暫停期間存續。

產生器是惰性產生序列最簡單的方式 -- 不會建構任何串列,項目只在消費者要求時才會被計算,而且函式想要的話可以永無止境地產出項目。

2.31.3. 惰性管線

產生器可以良好地組合。一個產生器的輸出可以餵入另一個產生器:

def numbers():
    i = 0
    while True:
        yield i
        i += 1

def squares(source):
    for x in source:
        yield x * x

pipeline = squares(numbers())

for v in pipeline:
    if v > 100:
        break
    print(v)

這些值一次一個地流經管線 -- 沒有中間串列,numbers 中也沒有內建上限,由消費者(for v in pipeline)決定何時停止。

Three boxes left-to-right: numbers(), squares(source), and the for-v-in-pipeline consumer. Three cycles are drawn beneath. In each cycle, the consumer sends a pull request leftward to squares, which sends a pull leftward to numbers; numbers yields a value rightward to squares, which yields its squared value rightward to the consumer.

消費者上的每一次 next() 都會觸發一次貫穿整條鏈的拉取;值只有在有東西要求時才會存在。

2.31.3.1. yield from

從另一個可迭代物件拉取項目並逐一產出,這種迴圈夠常見,因此 Python 提供了一個捷徑。yield from iter 運算式會依序產出該可迭代物件所產生的每一個值 -- 就如同產生器內嵌了一個 for x in iter: yield x 迴圈一般:

def chain(*sources):
    for source in sources:
        yield from source

for v in chain([1, 2, 3], (4, 5), "abc"):
    print(v)

輸出:

1
2
3
4
5
a
b
c

yield from 與明確的 for 迴圈完全等價,只是更簡短,而且它能將內層可迭代物件的 StopIteration 乾淨俐落地向上傳遞到外層產生器 -- 在將數個產生器首尾相連串接時非常有用。

2.31.3.2. yield 用盡時

執行到產生器函式的結尾之外(或遇到明確的 return)會自動引發 StopIteration。不需要手動引發它;外圍的 for 迴圈會偵測到它並結束。

當產生程式碼自然地寫成帶有幾個 yield 點的迴圈時,請使用產生器;當你真的需要將整個序列保留在記憶體中時,請使用單純的串列生成式。