ARTICLE DETAIL

资讯详情

深耕郑州网站建设与运营推广的一线实战洞察。

Python生成器与yield实战:从MemoryError到内存优化神器

Python生成器与yield实战:从MemoryError到内存优化神器 回头翻我当时那个Python脚本的Git提交记录还是能感觉到那股哭笑不得的劲儿程序跑起来不到三分钟MemoryError直接把我从座位上炸起来。那天要处理的是一份 12GB 的服务器日志我的第一版代码是经典的“一把梭”with open(access.log) as f: lines f.readlines() for line in lines: parse(line)readlines()之所以叫这个名字就是因为它真的会把每一行都读出来装成一个列表。12GB 日志读进来加上Python对象头的开销直接吃掉了 40GB 内存。后来我把代码改成逐行读取按行yield同样的任务内存峰值不到 10MB。那个下午我算是真正理解了Python生成器、yield和惰性求值这三个词的分量。这篇东西不是什么入门教程的复读是我在实际项目里把生成器用出价值之后的经验总结。适合所有已经会写Python、但还没真正体会过生成器威力的朋友你会知道为什么yield能把内存占用拉低两个数量级、生成器的执行流程到底怎么走、它在文件处理和数据管道里能玩出什么花以及包括send()/yield from/异步生成器在内的一整套进阶姿势。看完你可以直接把手里的活改成生成器写法也能避开我踩过的那些坑。1. 生成器不是列表的廉价版而是一种新的计算方式1.1 MemoryError的根源一次性把所有东西装进内存先从一个反直觉的事实说起for循环能遍历的东西不一定是“全部存在内存里”的。列表推导式会生成一个真正的列表里面有每一个元素元素总数等于迭代次数生成器表达式和生成器函数则不同它们每次只“吐”一个元素出来吐完就忘。这就是惰性求值——计算被推迟到真正需要那个值的那一刻。举个最典型的对比。假设我们有 1000 万个数字求和# 全部装入内存 nums [i for i in range(10_000_000)] total sum(nums) # 内存占用约 80MB # 惰性求值 nums_gen (i for i in range(10_000_000)) total sum(nums_gen) # 内存占用几乎恒定约 1MB 以下range(10_000_000)本身也是一个惰性对象它不保存所有数字只保存起点、终点和步长。但当你把它放进列表推导式时内存就爆炸了——列表必须为每个元素分配独立的内存槽位。而生成器表达式外面套上sum()sum每要一个值生成器就算一个值算完即弃内存曲线平得像心电图。这不是“性能优化的小技巧”而是计算模型的差异。列表是“空间换时间”的典型你一次性付清内存成本换来随机访问任意元素的能力。生成器是“时间换空间”内存几乎不增长但如果你想随机访问第500万个元素对不起你得先把前499万个跑一遍。1.2 yield与return的本质差异很多人写生成器函数时会下意识把yield当成“多次return”。这个类比方向是对的但容易忽略关键细节。看下面这段def count_up_to(n): i 0 while i n: yield i i 1 gen count_up_to(5)普通函数遇到return就彻底结束函数栈帧直接被销毁局部变量i消失。但生成器函数里的yield不同它会把当前执行到的位置、局部变量的值、循环状态全部“冻结”起来返回i给调用方后整个函数并没有退出。下次调用next(gen)时代码从yield那一行继续往下跑直到再次遇到yield或函数结束。这就是为什么说生成器函数是一个“可暂停的函数”。状态保存这件事不是Python偷偷用全局变量做的而是生成器对象自己保存了完整的栈帧。你可以把生成器理解成一本带书签的书yield是你在每一页停留的位置next()是往后翻一页。这种机制带来的直接好处是生成器函数里可以写复杂的控制流for循环、while循环、异常处理、嵌套调用全都支持。它不是一个简单的“产出序列的函数”而是一段可以被外部反复唤醒的代码。1.3 惰性求值到底“惰”在哪惰性求值的“惰”不是“懒得不干活”而是“按需计算”的原则。传统求值模型eager evaluation会在表达式被赋值时就完成所有计算惰性求值则把计算推迟到值真正被消费的那一刻且通常只计算当前需要的那一步。我用吃饭打比方列表是一次性把一整桌菜全抄出来摆好你吃哪个都可以生成器是一道菜一道菜地上你说“下一道”厨房才开火。如果你吃前三道就饱了后面的菜根本没做过当然不耗时间也不耗材料。换回程序场景列表生成器如果只取前三个元素后面九千九百九十九万个元素的加法照样全跑生成器则只算三个后面的计算从未发生。这种“按需计算”对两个场景特别有价值数据量未知或极大你不确定要消费多少元素用列表就会被迫按最大可能分配内存用生成器则“消费多少计算多少”。多个阶段串联上游计算很贵但下游可能提前停止比如找到第一个满足条件的值就break中间的惰性传递能省掉大量无效计算。但惰性也意味着延迟结算。你无法提前知道生成器的长度无法索引无法重复遍历。这是一把双刃剑后文我会有专门篇幅讲怎么绕开这些限制。2. next()、send()与状态机yield背后的运行之谜2.1 生成器的执行流程遇上yield就暂停理解生成器最好的方法就是在next()之间手动推进它。写一个带打印的生成器你能清清楚楚看到每一步发生了什么def demo(): print(进入生成器函数) print(第一次 yield 前) yield 1 print(第一次 yield 后) print(第二次 yield 前) yield 2 print(函数结束) gen demo() print(next(gen)) print(--- 第一次 next 结束 ---) print(next(gen)) print(--- 第二次 next 结束 ---)输出顺序是进入生成器函数 第一次 yield 前 1 --- 第一次 next 结束 --- 第一次 yield 后 第二次 yield 前 2 --- 第二次 next 结束 ---注意第一次next(gen)的调用并没有执行完整个函数它只执行到第一个yield就停住了并把1返回出来。第二次next(gen)从第一次挂起的地方继续往下跑过两行打印再次挂起在第二个yield返回2。如果再来一次next(gen)它会打印“函数结束”然后抛StopIteration——这是生成器耗尽的标准信号for循环监听到它就会安静退出。这个“停在哪下次就从哪继续”的行为就是生成器作为状态机的直接体现。你不需要自己管理i是不是加过了、循环条件是否成立Python都帮你记着。2.2 yield表达式两侧的值产出与接收很多教程只讲yield如何产出值但yield还有另一半功能接收外部传进来的值。核心语法是def accumulator(): total 0 while True: value yield total if value is None: continue total value这里的yield total有双重身份执行到它时把total作为next()的返回值给外部同时这个表达式的值可以由外部通过send()注入。也就是说yield不只是一个“输出语句”它本身是一个表达式有返回值。acc accumulator() print(next(acc)) # 启动生成器执行到 yield total输出 0 print(acc.send(10)) # 把 10 赋值给 valuetotal 变成 10yield total输出 10 print(acc.send(5)) # total 变成 15输出 15第一行必须用next(acc)或acc.send(None)启动生成器原因很简单生成器函数从头开始执行还没有走到yield外部的send没有可以“注入”的挂起点。一旦挂起在yield处send才有目标。这个机制是生成器能当协程用的基石。一个“协程”在这个层面理解就是一段代码可以被外部反复发送数据、中断、继续执行。accumulator就是一个最简形态的协程它不需要全局变量和回调就能在多次调用之间记住total。2.3 throw()与close()异常注入和资源释放gen.throw(exc)可以在生成器挂起的位置抛入一个异常。它模拟的是“生成器内部某个环节出了故障”而不是外部调用出错。举个工程场景你有一个生成器在持续读取外部API某个响应处理到一半你觉得数据格式不对想让它立刻停止。你可以给它抛一个自定义异常让它在内部捕获并做清理def api_reader(): try: for i in range(100): data yield i if data bad: raise ValueError(bad data) except ValueError: print(内部捕获准备清理) # 释放连接等 yield cleanup_done gen api_reader() next(gen) print(gen.send(bad)) # 内部抛异常被捕获yield一个清理状态close()则更彻底它会往生成器里抛GeneratorExit要求生成器退出。如果生成器内部没有处理它正常结束。你通常只在确定不再需要这个生成器时调用close()它和在finally里做资源释放配合得很好def resource_gen(): try: yield open yield working finally: print(释放资源)调用close()后finally里的清理代码一定会被执行。这在处理自定义协议的连接、临时文件时特别有用。2.4 生成器的生命周期从GEN_CREATED到GEN_CLOSEDPython的inspect模块提供了getgeneratorstate()方法可以查看生成器当前处于哪个阶段。四个状态对应四种阶段GEN_CREATED创建后尚未启动不能send非None值必须先next()或send(None)。GEN_RUNNING正在执行中。这个状态通常只有生成器内部调用yield给子生成器时才能被观察到外部一般看不到。GEN_SUSPENDED挂起在yield处可以继续next()或send()。GEN_CLOSED执行完毕或已被close()再迭代会立刻StopIteration。写个小例子验证一下import inspect def simple_gen(): yield 1 g simple_gen() print(inspect.getgeneratorstate(g)) # GEN_CREATED next(g) print(inspect.getgeneratorstate(g)) # GEN_SUSPENDED next(g, None) # 耗尽 print(inspect.getgeneratorstate(g)) # GEN_CLOSED掌握这个状态转换对你调试复杂的管道式生成器非常有帮助。遇到“生成器怎么就是不执行”的诡异问题先看状态如果是GEN_CREATED说明你忘了先next()如果是GEN_CLOSED说明它已经被消费完再send只会抛StopIteration。3. 惰性求值实战超大文件、无限序列与数据管道3.1 分块读取大文件的正确姿势回到开头那个12GB日志的问题。正确做法不是readlines()而是逐行读取让文件对象本身作为一个迭代器用def read_lines(filepath): with open(filepath, r) as f: for line in f: yield line.rstrip(\n) for line in read_lines(access.log): parse(line)文件对象f本身惰性for line in f每次只从磁盘读一行到内存。用yield之后整个函数变成一个生成器外层for每要一行生成器才往文件系统请求一行。12GB的日志内存占用基本等于单行日志的最大长度而不是文件大小。还有一个更贴近实际的分块需求你要处理的是二进制文件或者需要按固定字节数读取。比如解析一个超大JSONL文件每行都可能几MB逐行读没问题但如果是读一个大图片文件或压缩包就需要按chunk读def read_chunks(filepath, chunk_size8192): with open(filepath, rb) as f: while True: chunk f.read(chunk_size) if not chunk: break yield chunk用法上这个生成器可以喂给任何按块消费的处理器。内存占用上限就是chunk_size你想把吞吐拉高就调大chunk想控制内存就调小。我在处理网络抓包文件时经常用这种写法512KB的chunk能把内存控制在几十MB同时兼顾磁盘IO效率。3.2 无限数列生成器写斐波那契和素数筛无限序列是列表无法表示的但生成器可以“无限”地迭代下去直到你主动停止。斐波那契是最经典的教学用例但我更想强调它的范式意义def fibonacci(): a, b 0, 1 while True: yield a a, b b, a b fib fibonacci() for _ in range(10): print(next(fib), end ) # 0 1 1 2 3 5 8 13 21 34因为函数内部a, b状态一直保留每次next()都基于上一次的结果继续推演。不用任何全局变量就能维护一个永动机式的数列。素数筛的惰性版本更有意思。传统埃拉托斯特尼筛法需要预先分配一个布尔数组受限于上限。生成器版可以边筛边产出严格说它更接近“增量筛”def primes(): known [2] yield 2 n 3 while True: is_prime True for p in known: if p * p n: break if n % p 0: is_prime False break if is_prime: known.append(n) yield n n 2这里known列表会随着生成器运行不断变大严格意义上它不是“完全惰性”——所有已发现的素数都保存在生成器栈帧里。但它的好处是接口不变你可以for _ in range(1000) next(prime_gen)取前1000个素数也可以一直取到天荒地老。内存开销只和已经产生的素数个数相关和未来的搜索范围无关。3.3 用生成器搭数据流水线一行一行处理日志生成器之间可以串联成管道每个生成器只负责一个处理步骤。这是我在生产环境里最常用的模式。假设你的日志文件里每一行是一段JSON你需要过滤出levelERROR的条目再提取request_id字段最后按小时统计import json def read_log(filepath): with open(filepath) as f: for line in f: yield line.strip() def filter_error(lines): for line in lines: record json.loads(line) if record[level] ERROR: yield record def extract_hour(records): for record in records: ts record[timestamp] hour ts[:13] # 简单取小时 yield hour, record[request_id] def aggregate_hours(entries): counts {} for hour, _ in entries: counts[hour] counts.get(hour, 0) 1 yield counts pipeline aggregate_hours( extract_hour( filter_error(read_log(app.log)) ) ) result next(pipeline) print(result)这段代码的妙处在于read_log只产生原始行filter_error只负责过滤和解析extract_hour只负责字段提取aggregate_hours只负责统计。每个函数都可以单独测试也可以替换实现。数据流像水管一样从源头流向终点中间的每一步都不会积压整个数据集只保留当前正在处理的一条记录。这种“流水线”式设计有一个隐藏优点可插拔。如果你想加一个“按用户ID去重”的步骤只需在管道里插入一个生成器函数如果你想把文件源换成HTTP流只要换掉最上游的生成器。维护起来比写一个巨大的函数体舒服十倍。3.4 yield from把子生成器接入管道当你有多个生成器需要“透传”时嵌套for循环会变得很难看。yield from就是专门用来委托生成器的语法。它接收一个子生成器把所有yield出的值从子生成器转发到外层调用者同时自动处理子生成器的StopIteration。def sub_gen(): for i in range(3): yield i def wrapper(): yield start yield from sub_gen() yield end print(list(wrapper())) # [start, 0, 1, 2, end]yield from比手写for item in sub_gen: yield item更强的地方在于它会替你把send()和throw()也转发给子生成器。这意味着如果外层调用方往主生成器里send一个值这个值能直接传输到子生成器内部挂起的位置。手写for循环做不到这一点。所以当你需要把多个协程组合成复杂工作流时yield from几乎必不可少。它的存在让生成器委托变得和函数调用一样自然一个生成器可以把自己的执行权有计划地交给另一个生成器等后者耗尽再继续执行自己的代码。4. 生成器不是银弹性能测试与日常踩坑总结4.1 生成器表达式 vs 列表推导式用数据说话很多人以为生成器一定比列表快这是个误解。我用一个简单的timeit测试来说明对100万个数字做sum()生成器表达式 vs 列表推导式结果可能让你意外。import timeit print(timeit.timeit(sum([i * 2 for i in range(1_000_000)]), number50)) print(timeit.timeit(sum((i * 2 for i in range(1_000_000))), number50))在我本机的Python 3.10上列表版本常常更快。原因很直接生成器每次yield一个值都要经历生成器函数的挂起/恢复以及Python解释器层的状态切换这个开销比从列表里取一个已经算好的值要高。列表推导式虽然一次性算好所有元素但sum遍历时拿到的是纯C层面的迭代速度。那生成器的价值在哪在于内存和延迟而不是纯粹的CPU速度。当数据量小到可以整体放入内存时用列表无可厚非当数据量大到内存吃紧或者你只需要前几个结果时生成器就赢了。我的实践经验是如果你不确定数据规模会涨到多大先用生成器保平安等实测发现生成器真的成了瓶颈再换列表也不迟。这比一开始就把内存写爆炸要划算得多。4.2 一次性消费陷阱为什么生成器不能回头生成器是“一次性”的这是和列表最大的行为差异。它内部没有保存已产出的元素你不能像列表那样反复遍历。gen (x for x in range(5)) print(list(gen)) # [0, 1, 2, 3, 4] print(list(gen)) # []第一次list(gen)已经把生成器榨干第二次得到的只能是空列表。这看起来很简单但实际项目中遇到“为什么第二次处理没有数据”时经常让人抓狂。尤其是把生成器传给多个函数串行处理时第一个函数消费完毕第二个函数收到空数据。解决方案也不是没有如果确实需要多次遍历就老老实实用列表或者调用itertools.tee()分裂出多个独立迭代器。如果只是调试可以用itertools.islice(gen, n)预览前几个元素避免污染整个生成器。如果结构允许把生成器工厂函数反复调用比保存一个生成器实例更稳妥需要用一次调一次函数生成新的生成器。itertools.tee有一个陷阱它会缓存已消费的元素以保证多个迭代器的一致性如果原始迭代器前进得比派生迭代器远很多内存会涨。所以tee只适合少量派生、消费步伐接近的场景。4.3 诱人的in操作与next()默认值生成器没有长度没有下标用in判断元素是否存在也是合法的但代价是它会把生成器一直迭代到找到目标或者耗尽为止。这个行为如果发生在你打算继续使用该生成器的时候会毁掉后续所有迭代。gen (x * 2 for x in range(10)) print(6 in gen) # True且生成器已经迭代到 x3 的位置 print(list(gen)) # [8, 10, 12, 14, 16] —— 前面的 0,2,4 已经没了如果你只是想“取下一个值没有就给个默认值”务必用next(gen, default)而不是try/except或者indata next(gen, None) if data is None: print(生成器已空或第一个值就是None)这个None默认值也有歧义如果生成器本身可能产出None你可以换成一个哨兵对象sentinel object() data next(gen, sentinel) if data is sentinel: print(确实没有下一个值了)这种写法在写健壮的消费循环时几乎天天用到。4.4 生成器表达式里的延迟绑定坑生成器表达式的外部变量是延迟求值的这个特性经常引起“离奇”问题。下面这段代码是网上流传甚广的坑能看出问题的人不多def make_counters(): return (i for i in range(3)) gens [make_counters() for _ in range(2)]这个例子还好真正阴险的是funcs [] for i in range(3): funcs.append(lambda: i) print([f() for f in funcs]) # [2, 2, 2] 而不是 [0,1,2]同样的延迟绑定也会出现在生成器表达式里比如gens [] for i in range(3): gens.append((x for x in range(i)))这里每个生成器内部的i是外层循环的变量循环结束后i的值是2所以三个生成器全都遍历range(2)。解决办法也很简单把外层值绑定成默认参数或把它放进一个立即执行的函数gens [(lambda n: (x for x in range(n)))(i) for i in range(3)]这种坑不易察觉但一旦踩到错误往往表现为“所有生成器的行为一模一样”而不是崩溃。排查时要多留个心眼生成器表达式里的自由变量实际绑定的是外层作用域当前的值。5. 从生成器到协程yield与async/await的关系5.1 用yield手工实现一个任务调度器生成器既然可以挂起和恢复天然就能用来做协作式调度。我最早接触协程概念就是通过两个互相yield的生成器实现的“伪并发”def task_a(): for i in range(3): print(fA step {i}) yield def task_b(): for i in range(3): print(fB step {i}) yield a, b task_a(), task_b() for _ in range(3): next(a) next(b)输出结果是A和B交替执行。这不是真正的并行但它证明了“执行权可以在函数之间切换”这件事可以在纯生成器层面实现。配合send()你还能把调度器的控制信号传给任务def task(): while True: step yield print(f执行步骤: {step}) t task() next(t) # 启动 t.send(第一步) t.send(第二步)这个模型就是事件循环的雏形每个任务是一个生成器主循环决定下一个调度谁。Python的asyncio本质上也借鉴了这种协作式思想只不过做成了更完整的框架。5.2 yield from如何自动转发send()和throw()前面提到yield from会把send()和throw()转发给子生成器。这到底有什么用一个经典场景是用生成器实现状态机式的解析器。如果你有两个子解析器外层主生成器只需要yield from它们外部调用方不管往哪个层次send数据最终都能到达正在等待的那个子生成器。def parser_a(): data yield await a print(a got, data) def parser_b(): data yield await b print(b got, data) def main_parser(flag): if flag: yield from parser_a() else: yield from parser_b() m main_parser(True) print(next(m)) # 返回 await a m.send(hello a) # 数据被转发给 parser_a打印 a got hello a如果没有yield from你得手动维护一个“当前子生成器”变量循环里先判断该给谁send代码会复杂得多。5.3 async生成器异步环境下的惰性求值Python 3.6之后异步生成器成了官方特性用async def定义内部用yield产出值返回的是一个异步迭代器用async for消费。它在异步代码里保留了生成器的惰性优势同时允许在每次yield之间进行await其他IO操作。import asyncio async def fetch_records(limit): for i in range(limit): await asyncio.sleep(0.1) # 模拟网络请求 yield { id: i, data: frecord-{i} } async def consume(): async for record in fetch_records(3): print(record[id]) asyncio.run(consume())这里每次yield前都执行了一次异步等待。传统生成器函数里不能await因为它是同步函数异步生成器则把两个特性融合了起来外部的async for每次从它那里取一个结果时内部可能已经执行了若干IO操作。这个模式特别适合做“分页拉取逐条处理”的场景。调用第三方API分页接口时你不想一次把全部分页都拉进内存可以用异步生成器封装分页游标async def paginate(api, page_size50): page 1 while True: batch await api.get_page(page) for item in batch: yield item if len(batch) page_size: break page 1调用方用async for逐条消费内存里永远只有一页的数据。这种做法在数据管道、消息队列消费者、爬虫任务里都很好用是异步惰性求值的典型落地。5.4 什么时候该用什么时候不该用写到这里我想补充一点“经验判断”。不是所有循环都要改成生成器。如果你的数据本身就已经在内存里而且需要反复随机访问生成器只会增加代码复杂度和运行开销。但如果你面临以下情况之一生成器基本是首选数据源头是文件、网络流、数据库游标这类不适合整体载入的对象。你只需要一个序列的前N个元素或要重复执行“取下一个”操作。你要构建多阶段数据管道每阶段都只依赖上一阶段的一个元素。你需要在多个函数之间传递执行权也就是协程需求。我个人还有一个判断标准当我发现自己在一个函数里同时维护“当前游标”“剩余数量”“下一步状态”这几个变量时我就会考虑拆成生成器。因为生成器把这些状态全部“外包”给了Python的栈帧管理我只需要关注核心逻辑。最后再分享一个工具习惯itertools模块里有大量和生成器配合的组合工具比如islice、takewhile、chain、groupby。我处理日志管道时很多时候根本不需要自己写生成器函数把itertools的函数拼一拼就是一个健壮的流式处理链。这些函数都是C实现的遍历效率比自己手写yield的纯Python循环高不少。如果你想在惰性求值这条路走得更深itertools绝对是不可错过的下一站。
返回列表