一个真实的翻车现场
上个月同事写了个日志分析脚本,跑我们线上那台 8G 内存的机器,脚本一启动 OOM Killer 就把它干掉了。我过去一看,代码大概是这么写的:
lines = open("/var/log/app.log").readlines() # 先把 2 亿行全读进内存
errors = [l for l in lines if "ERROR" in l] # 再复制一份做过滤
result = [l.strip() for l in errors] # 又复制一份 strip
print(len(result))
一个 4G 的日志文件,经过三次列表推导,峰值内存直接冲到十几 G。机器不死才怪。
我当时只改了一行——把方括号换成圆括号,内存从十几 G 掉到几百 KB,效果立竿见影。这个圆括号背后藏着的,就是 Python 生成器(generator)这个东西。今天把它讲透。
生成器到底是什么
一句话:函数体里只要出现了 yield 关键字,它就不再是普通函数,而是「生成器函数」。调用它不会执行函数体,而是返回一个生成器对象。真正的执行,是每次你调 next() 的时候才发生。
def count_up(n):
i = 0
while i < n:
yield i # 产出一个值,然后「暂停」在这里
i += 1
g = count_up(3)
print(next(g)) # 0
print(next(g)) # 1
print(next(g)) # 2
print(next(g)) # StopIteration 异常
你可能会问,这不就是个能暂停的函数吗,有什么稀奇的?稀奇的就在「暂停」这两个字上。普通函数要么一口气跑完,要么抛异常退出,中间状态你拿不到。而生成器在 yield 处暂停时,它的局部变量、循环进度、执行位置全都原封不动地保存在生成器对象里。下次 next() 它会从暂停点接着跑,而不是从头再来。
这本质上是个状态机。你每调一次 next(),它就往前推一格,推完停下来等你下一次调用。因为它一次只算一个值,所以内存占用是 O(1) 级别的——这就是惰性求值(lazy evaluation)的威力。
生成器表达式:把方括号换成圆括号
前面说的翻车现场,核心就是列表推导和生成器表达式的区别。看这个对比:
# 列表推导:先建好整个列表,再交给 sum 去加
total = sum([i * i for i in range(10_000_000)])
# 生成器表达式:sum 要一个,它就现算一个
total = sum(i * i for i in range(10_000_000))
第一行会先在内存里铺开一个 1000 万个元素的列表,第二行则是一次只生成一个平方数、加完就丢。结果一模一样,内存却是天壤之别。
我用 tracemalloc 实测了不同数据规模下,列表和生成器的峰值内存,差别大到有点不真实:
1 万元素时列表要 0.38MB,100 万元素时飙到 38.57MB;而生成器无论数据多大,始终只占 0.4KB 左右——就一个生成器对象的大小。内存节省率 99.9% 起步。这就是为什么大数据量场景下,能用生成器就别用列表。
进阶:send() 让数据双向流动
如果生成器只是「往外吐值」,那它跟迭代器也没太大区别。真正的分水岭是 send()——它能在生成器暂停时,把值「塞回」生成器内部。
def accumulator():
total = 0
while True:
x = yield total # 先产出 total,同时等着接收塞进来的值
if x is None:
break
total += x
acc = accumulator()
next(acc) # 必须先启动一次,产出 0
print(acc.send(10)) # 塞入 10,产出 10
print(acc.send(20)) # 塞入 20,产出 30
print(acc.send(30)) # 塞入 30,产出 60
注意 yield total 这行的读法:它先做两件事——把 total 吐出去、然后卡住,等外部 send() 一个值进来赋给 x。所以生成器不是单向的管道,它是双向的。这个能力后来直接催生了 Python 的协程(coroutine)——asyncio 里 await 的老祖宗就是 yield from。
有个新手必踩的坑:send() 之前必须先 next() 一次(或 send(None)),把生成器推进到第一个 yield 处。上来就 send() 会直接抛 TypeError。
throw() 和 close():往生成器里塞异常
既然能塞值,自然也能塞异常。这就是 throw() 和 close() 的用处——在生成器暂停的位置注入异常,让它有机会做清理。
def worker():
try:
while True:
yield "working"
except Exception as e:
yield f"caught: {e}"
w = worker()
print(next(w)) # "working"
print(w.throw(ValueError("boom"))) # "caught: boom"
throw() 注入的是你指定的异常;close() 注入的则是特殊的 GeneratorExit 异常,专门用来通知生成器「结束了,赶紧收尾」。看这个带 finally 的经典例子:
def gen():
try:
yield 1
yield 2
finally:
print("清理资源完成")
g = gen()
next(g)
g.close() # 在 yield 处抛 GeneratorExit,触发 finally
# 输出:清理资源完成
如果你的生成器打开了文件句柄、持有数据库连接,finally + close() 这套组合就是你保证资源不泄漏的兜底方案。
yield from:把生成器委托出去
当你的生成器需要调用另一个生成器时,新手会写嵌套的 for 循环:
def main_gen():
yield "start"
for item in sub_gen():
yield item
yield "end"
这段完全可以用 yield from 一行搞定:
def sub_gen():
yield "a"
yield "b"
def main_gen():
yield "start"
yield from sub_gen() # 委托给子生成器
yield "end"
print(list(main_gen())) # ['start', 'a', 'b', 'end']
但 yield from 远不只是省两行代码。它建立的是一条双向隧道:你对 main_gen 调 send()、throw()、close(),这些操作会原封不动地穿透到 sub_gen 内部。它也是协程能层层组合的关键机制。
实战:用生成器管道流式处理大日志
回到最开始的翻车现场,正确的写法是搭一条生成器管道,让数据像流水一样流过每个处理环节,全程不落地:
def read_log(path):
with open(path) as f:
for line in f: # 文件对象本身惰性迭代,逐行读
yield line
def filter_errors(lines):
for line in lines:
if "ERROR" in line:
yield line.strip()
def extract_timestamps(lines):
for line in lines:
yield line[:19] # 只取前 19 个字符的时间戳
# 三个生成器串成一条流水线,任何时刻内存里只有一行
pipe = extract_timestamps(filter_errors(read_log("/var/log/app.log")))
for ts in pipe:
print(ts)
这段代码里 read_log、filter_errors、extract_timestamps 都是生成器,一个套一个。数据被「拉」着走:最外层的 for 要一个值,这个需求一路传导到最底层去读一行文件。任何一个时刻,内存里都只有一行日志,加上几个生成器对象。这就是生成器管道(generator pipeline)的典型用法。
三个必须知道的陷阱
生成器好用,但它不是没有代价。这三个坑我挨个踩过:
- 一次性消费:生成器只能从头到尾走一遍,走到头就「烧完了」,再遍历就是空的。需要反复访问同一份数据,或者需要
len()、随机下标访问时,老老实实用列表。 - 延迟报错:因为值是惰性算出来的,你写代码时不会报错,真正迭代的时候错误才会炸出来。生成器表达式里如果有个除零,错误可能延迟到几层管道之后才暴露,排查时要顺着管道往回找。
- 内存换 CPU:生成器省内存,但每次
next()都有一层函数调用的开销。数据量小的时候(比如几百个元素),生成器反而可能比列表慢。小数据直接上列表,别为了「优雅」硬上生成器。
常见问题(FAQ)
Q: 生成器和迭代器(iterator)是什么关系?
生成器是迭代器的一种,也是最省事的一种。任何实现了 __iter__ 和 __next__ 方法的对象都是迭代器;而生成器函数/生成器表达式会自动生成这两个方法,你一行都不用自己写。反过来说,迭代器不一定是生成器——你可以手写一个 __next__ 的类,但十有八九不如 yield 来得简洁。
Q: 什么情况下应该用列表,而不是生成器?
三个信号:一是数据量小(几百到几千个元素),生成器的函数调用开销反而拖后腿;二是你需要多次遍历同一份数据,生成器一次性用完就没了;三是你需要 len()、切片、随机访问等列表专属能力。除此之外,尤其是处理「读一次、算一次、丢一次」的流式数据,无脑选生成器。
Q: 生成器和 asyncio 协程有关系吗?
有,而且是直接的传承关系。Python 的协程演进路线就是:yield(生成器)→ yield from(委托)→ async/await(原生协程)。await 的底层机制和 yield from 的「暂停 + 双向传值」是同一套思想。理解了生成器,再去看 asyncio 会顺很多。
总结
生成器这东西,说白了就是给你一个「随时暂停、随时恢复」的函数。它的价值不在语法,而在思维方式——把「一次性算完一整批数据」换成「需要时算一个」。这一换,内存占用从 O(n) 变成 O(1),十几 G 的翻车现场变成几百 KB 的丝滑流水线。
如果你对惰性求值还有兴趣,可以接着看看Python itertools 惰性管道的实战,它讲的是怎么用标准库把生成器管道玩出花;想进一步压内存,memoryview 零拷贝那篇会告诉你连字节都不复制是什么体验;而生成器的 yield from 一路走下去就是协程,asyncio.TaskGroup 结构化并发里有异步世界的完整玩法。
