2.31. 迭代器与生成器

for 循环所做的工作比它看起来要多。本页将介绍它所依赖的迭代器协议，以及让你能够构建自己迭代器的 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

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）会在两次恢复之间被保留。

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）决定何时停止。

消费者上的每次 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 点的循环时，使用生成器；当你确实需要把整个序列放在内存中时，使用普通的列表推导式。