如何同步迭代依赖生成器以同时查看输入输出流？

问题核心分析 #

生成器是一次性、状态化的迭代器。你用zip(g2, g1)时，zip会交替从两个生成器取元素：每次从g2取数，会先驱动g1生成一个元素并完成转换；接着从g1取数时，g1又会生成下一个元素，最终导致原始值和转换值错位，且仅能拿到一半数据。

最优解决方案 #

使用itertools.tee复制原始生成器，让转换生成器和测试用的原始视图共享同一套生成器输出，无需重复创建复杂的底层迭代器。tee会自动缓存生成器的输出，供多个迭代器同步使用，既保留流式处理特性，又能精准匹配原始值与转换后的值。

修正后代码示例 #

import itertools

def gen1(a):                                                                                                            
    for item in a:                                                                                                      
        yield item                                                                                                      
                                                                                                                        
def gen2(a):                                                                                                            
    for item in a:                                                                                                      
        yield item**2                                                                                                   
                                                                                                        
a = range(10)                
g1 = gen1(a)                                                                                                            
# 复制原始生成器为两个独立迭代器，共享底层输出
g1_for_transform, g1_for_test = itertools.tee(g1)
g2 = gen2(g1_for_transform)                                                                                                           
# 同步迭代转换后的生成器和复制的原始生成器
for transformed, original in zip(g2, g1_for_test):                                                                                                                                                                            
    print(transformed, original)

执行结果 #

0 0
1 1
4 2
9 3
16 4
25 5
36 6
49 7
64 8
81 9

注意事项 #

• 若生成器输出数据量极大，tee的缓存可能占用较多内存，但测试场景下一般无需担忧。
• 需保持两个复制后的迭代器同步迭代，避免交叉操作导致缓存失效、数据错位。

内容的提问来源于stack exchange，提问作者Ben Farmer