Python 生成器和 yield

生成器就是用普通函数写出来的迭代器。它在 yield 处暂停，从停下的地方接着执行，而且从不构建你原本要构建的那个列表。

上一篇结尾，为了从三倒数到一，写了二十行的类。下面是同样的功能。

三行

def countdown(n):
    while n > 0:
        yield n
        n -= 1

print(list(countdown(3)))

for i in countdown(3):
    print(i)

输出：

[3, 2, 1]
3
2
1

函数体里只要有 yield，就是生成器函数。调用它并不会执行函数体——而是返回一个生成器，生成器就是迭代器。__iter__ 和 __next__ 由 Python 替你写好。

你不要，它就不动

这部分值得仔细看。

def noisy():
    print('  starting')
    yield 1
    print('  between')
    yield 2
    print('  finishing')

g = noisy()
print('created — the body has not run')
print('got', next(g))
print('got', next(g))
try:
    next(g)
except StopIteration:
    print('done')

输出：

created — the body has not run
  starting
got 1
  between
got 2
  finishing
done

注意顺序。starting 出现在 created 之后，而不是之前。只有调用 next 时，函数体才会执行，一碰到 yield 就又停下。

return 结束函数，丢掉所有局部变量。yield 则把函数挂起，变量原样保留，下一次 next() 从下一行接着执行。函数体最终结束时，Python 抛出 StopIteration——正是上一篇手动捕获的那个信号。

惰性到底有什么用

import sys

squares_list = [n * n for n in range(1_000_000)]
squares_gen  = (n * n for n in range(1_000_000))

print('list:', sys.getsizeof(squares_list), 'bytes')
print('gen: ', sys.getsizeof(squares_gen), 'bytes')
print('same total:', sum(squares_list) == sum(squares_gen))

输出：

list: 8448728 bytes
gen:  200 bytes
same total: True

列表里存着一百万个数。生成器里存的只是一个暂停的函数。两者求和结果相同，但只有一个需要整个装进内存。

把方括号换成圆括号，区别就只有这一点。这就是生成器表达式——永远不会构建出列表的列表推导式。

nums = [1, 2, 3, 4, 5, 6]

print(sum(n * n for n in nums if n % 2))
print(sorted((n * n for n in nums), reverse=True)[:3])

输出：

35
[36, 25, 16]

生成器表达式是唯一的参数时，可以省掉它自己的括号——所以 sum(n * n for n in nums) 才能写成这样。

现在可以无限了

包含所有自然数的列表，你不可能拥有。包含所有自然数的生成器却完全没问题，因为它每次只保存当前这一个。

def naturals():
    n = 1
    while True:
        yield n
        n += 1

out = []
for n in naturals():
    if n * n > 500:
        break
    out.append(n * n)

print(out)

输出：

[1, 4, 9, 16, 25, 36, 49, 64, 81, 100, 121, 144, 169, 196, 225, 256, 289, 324, 361, 400, 441, 484]

while True 永远不会结束。结束循环的是 break，生成器就这么停在原地，一直暂停着，最后被回收。

一级一级往前要

生成器可以串起来。每一级只向前一级要一个元素。

lines = ['17,ada,engineer', '', '18,grace,admiral', 'not a row', '19,alan,logician']

good   = (l for l in lines if l.count(',') == 2)
fields = (l.split(',') for l in good)
people = ((int(i), name) for i, name, _ in fields)

for pid, name in people:
    print(pid, name)

输出：

17 ada
18 grace
19 alan

for 循环开口要之前，什么都没发生。之后，一行数据依次穿过三级，被打印出来，然后才轮到下一行。整个过程中从来没有出现过完整的中间列表。把 lines 换成 4GB 的文件，代码一行都不用改。

yield from

如果一个生成器的任务就是把另一个生成器产出的东西转交出去，用 yield from 一行就能写清楚。

def flatten(items):
    for item in items:
        if isinstance(item, list):
            yield from flatten(item)
        else:
            yield item

print(list(flatten([1, [2, [3, 4]], 5, [[6]]])))

输出：

[1, 2, 3, 4, 5, 6]

没有它，你就得写个循环把每个元素再 yield 一遍。能用，但读起来差一些。

你要放弃的两样东西

生成器是迭代器，所以只能跑一次——上一篇讲的全都适用。

g = (n for n in range(4))
print(list(g))
print(list(g))

输出：

[0, 1, 2, 3]
[]

而且它没有长度，也不能索引，因为还没产出的东西，它自己也不知道。

g = (n for n in range(5))
try:
    print(len(g))
except TypeError as err:
    print('TypeError:', err)

try:
    print(g[0])
except TypeError as err:
    print('TypeError:', err)

输出：

TypeError: object of type 'generator' has no len()
TypeError: 'generator' object is not subscriptable

如果你需要长度或索引，那你需要的其实是列表。对它调用 list()，把内存花出去。

要点

• 函数体里有 yield，函数就成了生成器函数。调用它什么都不会执行。
• 每次 next() 会把函数体执行到下一个 yield，然后冻结在那里，连同局部变量一起。
• 圆括号得到生成器表达式：就是你平时写的推导式，只是不构建列表。
• 生成器可以是无限的，还能串成管道，每一级最多只持有一个元素。
• 只能遍历一遍，没有 len，不能索引。这就是代价，而且通常值得付。