会暂停的函数
生成器看起来像普通函数,但它不会算出完整结果再返回,而是一次产出(yield)一个值,在每次产出之间暂停,直到使用它的一方要下一个值。
最简单的生成器:
注意这里用的是 yield 而不是 return。for 第一次要值时,Python 运行函数体,直到遇到 yield 1。函数就在那里暂停,把 1 交给循环,并准确记住自己停在哪里,包括所有变量。下一次迭代从上次停下的地方继续:current += 1,回到 while,yield 2。如此反复,直到循环条件不成立,生成器就会直接结束。
这种暂停再恢复的机制,就是生成器的全部诀窍。
为什么不直接构建一个列表?
因为列表版本会预先分配好所有的值:
5 个元素没问题。现在想象你需要 5000 万个整数,而你只关心第一个满足某个条件的数。列表版本会分配 5000 万个 int,然后你把其中大部分都扔掉。生成器版本则是调用方用多少就创建多少。当 for 循环找到想要的值并 break 时,生成器就直接停下了。
值得记住的模式是:生成器让你在编写迭代代码时,不必预先决定需要多少结果。
生成器表达式
如果你写过列表推导式,你已经知道它的语法了:把方括号换成圆括号就行:
squares_gen 此时还没有计算任何东西,它只是一份“配方”。遍历它时,配方才会一步一步地执行。
生成器表达式非常适合作为参数,传给接收可迭代对象的函数:
没有中间列表。sum、max 和 any 一次读取一个值,这正是它们需要的。
逐行读取大文件
这是生成器在实际中最典型的用法:处理一个大到无法整个载入内存的文件:
def parse_log_lines(path):
with open(path) as f:
for line in f:
if line.startswith("ERROR"):
yield line.rstrip()
for error in parse_log_lines("app.log"):
print(error)
文件是惰性读取的。每次向生成器要值,它都从磁盘读取一行、过滤,然后 yield。无论文件多大,内存占用都保持平稳。
只能用一次
生成器只能遍历一遍。遍历到末尾后,它就耗尽了:
第二个循环什么也不会打印,因为生成器里已经什么都不剩了。
如果需要多次遍历,可以再次调用生成器函数得到一个新的生成器,或者用 list(...) 把序列转换成列表,再反复遍历这个列表。根据成本来选择:如果生成的工作很轻量,重新生成就行;如果序列很小,用列表就行。
next() 与手动迭代
你不一定要用 for 循环。next() 一次取出一个值:
StopIteration 是生成器表示“我结束了”的方式。for 循环会悄悄地捕获它。在手动迭代的代码中,你可以给 next(gen, default) 传一个默认值来避免这个异常。
无限生成器
由于值是按需产生的,生成器可以表示一个没有尽头的序列,只要使用方会在某个时候停止索取:
里面带有 yield 的 while True 并不会让程序卡住,它只是表示“只要有人继续要,就继续产出”。何时停止由使用方决定。
这种模式常见于流式数据、事件循环,以及任何从没有确定长度的数据源中取值的场景。
yield from:委托给另一个可迭代对象
如果你的生成器想产出另一个可迭代对象中的所有值,用 yield from 一行就能搞定:
不用 yield from 的话,你就得写一个嵌套的 for 循环,在里面 yield x。如果你用到 send() 和 throw(),它也会正确地转发这些调用。不过在日常代码中,你可以把它理解为“把这个东西里的每个值都 yield 出去”。
什么时候该用生成器
三个信号说明生成器是合适的工具:
- 序列很大、可能无限,或者一次性完整生成的代价很高。
- 使用方可能在中途停止(例如找到第一个匹配项就
break)。 - 你想把多个变换串联起来(过滤、映射、取前几个),又不想构建中间列表。
以及什么时候不该用:
- 你需要随机访问(
seq[42])。生成器只能向前走。 - 你需要多次遍历同一个序列。用列表。
- 序列很小,而且你已经有了它。列表推导式更简单。
生成器、列表推导式和普通列表各自适合不同的任务。真正的本事是不假思索地选出合适的那个。培养这种直觉最快的方法是:每写一次迭代,都留意一下是“先全部生成”更合适,还是“一次生成一个”更合适。
下一篇:深入上下文管理器
到这里,你已经见过 Python 用于迭代的大部分惯用写法。接下来是上下文管理器,也就是 with 语句。它和生成器配合起来,非常适合从文件和网络连接中以流的方式读取数据。
常见问题
Python 中的生成器是什么?
生成器是一种一次产生一个值、在值与值之间暂停的函数。它和普通函数一样用 def 来写,只是用 yield 代替 return。调用它会返回一个生成器对象;for 的每次迭代或每次调用 next(),都会让函数运行到下一个 yield 为止。
列表和生成器有什么区别?
列表一次性把所有元素都放在内存中。生成器按需计算元素,元素被使用后就会被丢弃。对于很大或无限的序列,生成器只占用很少且固定的内存;对于需要反复使用的小结果,列表更合适。
生成器可以遍历两次吗?
不可以。生成器在完整遍历一遍后就耗尽了,对它进行第二次 for 循环什么也得不到。如果需要多次遍历,可以再次调用生成器函数得到一个新的生成器,或者把结果转换成列表。