Python itertools.tee 分流:两个迭代器为什么仍会一起占用内存
一份数据既要生成预览,又要计算汇总,最自然的想法是把迭代器分成两路。调用 tee 后,确实可以分别推进两个结果,但这不意味着复制了两个独立数据源。它们仍从同一个上游取值,只是各自记录阅读位置,暂存尚未被其他分支追上的内容。
这种设计适合速度接近的消费者。若预览只取几条便停住,而汇总一路读到末尾,落后的那一路就可能让大量数据继续被保留。排查内存增长时,不能只看当前循环一次处理一条,也要检查所有分支之间累计拉开的距离。
AI概念配图:用抽象物件说明本文讨论的关系,不是实际运行界面或测量结果。
先观察上游究竟被调用了几次
下面的完整程序可保存为 demo.py,用 Python 3 运行。数据源每产生一个字典,就把编号写进 produced。这个列表只用于观察取值次数,不把对象大小当作内存测量。实验刻意限制为五条记录,可以清楚检查快路前进和慢路追赶。
创建分支时还没有数据产生。快路取三条后,上游前进三次;慢路再取第一条时,拿到的是已有对象,上游不再运行。代码随后修改快路读到的字典,再检查慢路所见内容,说明分流复制的是读取位置,而不是每条记录的深层内容。
from itertools import tee
produced = []
def records():
for i in range(5):
produced.append(i)
yield {'id': i, 'tags': []}
fast, slow = tee(records())
assert produced == []
head = [next(fast) for _ in range(3)]
assert produced == [0, 1, 2]
head[0]['tags'].append('preview')
first = next(slow)
assert first is head[0] and first['tags'] == ['preview']
assert produced == [0, 1, 2]
assert [x['id'] for x in fast] == [3, 4]
assert produced == [0, 1, 2, 3, 4]
assert [x['id'] for x in slow] == [1, 2, 3, 4]
assert list(fast) == list(slow) == []
print('source produced once:', produced)
print('shared first object:', first['tags'])
empty_a, empty_b = tee(())
assert list(empty_a) == list(empty_b) == []
source = iter(range(4))
a, b = tee(source)
assert next(a) == 0
assert next(source) == 1
assert list(a) == [2, 3]
assert list(b) == [0, 2, 3]
print('empty and bypass-source checks passed')缓冲的压力来自读者之间的距离
断言证明了数据读取行为,但没有声称五条记录恰好占多少字节。具体实现可能按块管理缓存,还包含对象引用与运行时开销。实际预算应结合记录大小、分支数量和最大落后距离测量,不能把一个固定常数写成所有版本都成立的公式。
当慢路暂停而快路继续时,后续对象必须保持可用,否则慢路恢复就无法得到完整序列。尤其要留意日志预览、抽样检查和被提前退出的循环:代码表面上已经停止工作,只要仍持有相应分支,它就可能继续影响其他消费者的资源使用。
创建 tee 后不要再直接推进原迭代器。最后一组实验故意从原始 source 偷读一个数,于是两条分支都失去这个值。这不是数据源自动分配给第三个消费者的承诺,而是绕开了分流器的缓存管理;业务代码应只保留分支入口。
把消费计划写进接口约定
如果任务本来就是先完整分析一次、再完整输出一次,而且数据量已知且有限,直接收集为列表通常更容易理解和估算。这个选择明确承认需要保存完整输入,也让重复遍历成为显式需求,而不会把存储成本藏在落后的分支后面。
如果数据量很大,可以优先把两项计算放在一次循环里,逐条更新两个结果;若消费者确实必须独立推进,则应在更上层设计有容量限制的队列、暂停机制或落盘方案。tee 本身不会替调用方建立这些流量控制规则。
不要把这些分支交给多个线程并发调用。官方文档明确说明同一次 tee 得到的迭代器不保证线程安全,即使原始数据源自身允许并发读取也不例外。这里的实验是单线程的交错消费,不是并发正确性的证明。
验收时至少保留空输入、单路先走完、两路交替、可变记录与意外直读源这几类样本。真正要确认的是每条数据被谁读取、对象由谁修改、落后多久仍会保留,以及某个消费者退出后,剩下的处理流程能否满足资源预算。


