Menu

Python 生成器:yield、惰性迭代与生成器表达式

Python 中的生成器如何惰性地产生值:yield 关键字、生成器表达式,以及什么时候它比普通列表更合适。

本页包含可运行的编辑器。编辑、运行,立即查看输出。

会暂停的函数

生成器看起来像普通函数,但它不会算出完整结果再返回,而是一次产出(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 出去”。

什么时候该用生成器

三个信号说明生成器是合适的工具:

  1. 序列很大、可能无限,或者一次性完整生成的代价很高。
  2. 使用方可能在中途停止(例如找到第一个匹配项就 break)。
  3. 你想把多个变换串联起来(过滤、映射、取前几个),又不想构建中间列表。

以及什么时候不该用:

  • 你需要随机访问(seq[42])。生成器只能向前走。
  • 你需要多次遍历同一个序列。用列表。
  • 序列很小,而且你已经有了它。列表推导式更简单。

生成器、列表推导式和普通列表各自适合不同的任务。真正的本事是不假思索地选出合适的那个。培养这种直觉最快的方法是:每写一次迭代,都留意一下是“先全部生成”更合适,还是“一次生成一个”更合适。

下一篇:深入上下文管理器

到这里,你已经见过 Python 用于迭代的大部分惯用写法。接下来是上下文管理器,也就是 with 语句。它和生成器配合起来,非常适合从文件和网络连接中以流的方式读取数据。

常见问题

Python 中的生成器是什么?

生成器是一种一次产生一个值、在值与值之间暂停的函数。它和普通函数一样用 def 来写,只是用 yield 代替 return。调用它会返回一个生成器对象;for 的每次迭代或每次调用 next(),都会让函数运行到下一个 yield 为止。

列表和生成器有什么区别?

列表一次性把所有元素都放在内存中。生成器按需计算元素,元素被使用后就会被丢弃。对于很大或无限的序列,生成器只占用很少且固定的内存;对于需要反复使用的小结果,列表更合适。

生成器可以遍历两次吗?

不可以。生成器在完整遍历一遍后就耗尽了,对它进行第二次 for 循环什么也得不到。如果需要多次遍历,可以再次调用生成器函数得到一个新的生成器,或者把结果转换成列表。

Coddy 编程语言插图

和 Coddy 一起学编程

开始