无法修改昂贵生成器时，如何关联gen与trans的对应输出？

问题描述 #

我有一个生成器gen用来生成数据项，另一个生成器trans用来转换输入项，每个输入项对应一个输出项。已知这两个生成器运行成本很高，而且无法修改，它们还可能带有额外参数。gen的输出会传入trans，但我遍历trans的结果时，需要同时获取gen对应的输出。我目前用tee(gen())，再把它和trans的结果zip起来，这个方案能用，但想知道有没有更优的替代方案？

示例代码：

from itertools import tee

# 这两个生成器只是示例，假设它们运行成本高且无法修改
def gen():
    yield from range(3)
def trans(inp):
    for x in inp:
        yield chr(x + ord("A"))

# 我的问题是：有没有更好的方式实现下面两行代码的功能？
g1, g2 = tee(gen())
for i, o in zip(g1, trans(g2)):
    print(f"{i} -> {o}")

解决方案分析 #

你的方案其实已经是这类场景下的标准最优解之一，原因如下：

• tee的设计初衷就是为了复制迭代器/生成器，避免重复执行高成本的gen()——如果直接写zip(gen(), trans(gen()))，会让gen()运行两次，完全不符合你的需求。
• 这种写法简洁、符合Python迭代器的设计哲学，而且不需要修改原生成器，完美适配你“无法修改gen和trans”的限制。

如果非要找替代方案，只有两种思路，但都不如你的方案优雅：

1. 缓存gen的输出：先把gen()的所有结果一次性存入列表，再用zip和trans处理。但这种方式会把所有数据加载到内存里，如果gen生成的数据量很大，会直接导致内存问题，完全不如tee的惰性迭代高效。
   示例代码：

   gen_results = list(gen())
   for i, o in zip(gen_results, trans(iter(gen_results))):
       print(f"{i} -> {o}")
   

2. 包装生成器同步输出：写一个包装函数，每次从gen取一个项，同时拿到trans的对应输出，再把两者一起yield出来。但这需要额外写代码，而且本质上和tee+zip的逻辑一致，只是把逻辑封装了起来，并没有性能或简洁性上的优势。
   示例代码：

   def sync_gen_trans(gen_func, trans_func, *gen_args, **gen_kwargs):
       g = gen_func(*gen_args, **gen_kwargs)
       t = trans_func(g)
       for item, transformed in zip(g, t):
           yield item, transformed
   
   for i, o in sync_gen_trans(gen, trans):
       print(f"{i} -> {o}")
   

总结来说，你当前的tee+zip方案已经是最适合的选择，不需要替换。

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