如何让Python生成器向两个分支输出且不重复计算状态？

方法1：异步协程+异步队列（轻量替代线程队列） #

用asyncio的异步队列和协程实现，无需多线程，单线程内即可完成生产者-消费者的分发，每个元素仅生成一次，且不会缓存全部结果（可通过队列最大长度控制内存占用）。

示例代码：

import asyncio

async def source_generator():
    # 模拟原始生成器，每次生成一个元素，仅计算一次
    for i in range(5):
        await asyncio.sleep(0.1)  # 模拟耗时计算
        print(f"生成元素: {i}")
        yield i

async def consumer(name, queue):
    while True:
        item = await queue.get()
        if item is None:  # 终止信号
            break
        # 模拟分支处理逻辑
        print(f"消费者{name}处理元素: {item}")
        queue.task_done()

async def main():
    # 创建两个异步队列，用于分发元素
    queue1 = asyncio.Queue(maxsize=1)  # 限制队列长度，避免内存堆积
    queue2 = asyncio.Queue(maxsize=1)

    # 启动两个消费者任务
    task1 = asyncio.create_task(consumer("A", queue1))
    task2 = asyncio.create_task(consumer("B", queue2))

    # 生产者分发元素到两个队列
    async for item in source_generator():
        await queue1.put(item)
        await queue2.put(item)

    # 发送终止信号
    await queue1.put(None)
    await queue2.put(None)

    # 等待消费者完成
    await asyncio.gather(task1, task2)

asyncio.run(main())

该方案中，source_generator每生成一个元素就立刻放入两个队列，消费者按需取出处理，元素不会被全量缓存（仅队列内暂存未处理的少量元素），完全符合“单次计算、无全量缓存”的要求。

方法2：用itertools.tee（同步场景下的轻量方案） #

如果后续分支都是同步迭代，且两个分支的消费速度相差不大，可以用itertools.tee。注意：若两个迭代器消费速度差异大，tee会缓存未被慢迭代器消费的元素，但对称消费（分支处理速度接近）场景下，这个方案最简洁。

示例代码：

from itertools import tee

def source_generator():
    for i in range(5):
        print(f"生成元素: {i}")
        yield i

# 拆分生成器为两个独立迭代器
gen1, gen2 = tee(source_generator())

# 分支1处理
for item in gen1:
    print(f"分支1处理元素: {item}")

# 分支2处理
for item in gen2:
    print(f"分支2处理元素: {item}")

如果需要并行同步处理，tee也能配合线程使用，但消费速度差异大时仍会产生缓存，因此更适合同步、消费速度接近的场景。

方法3：自定义生成器分发器（同步场景下的无缓存方案） #

如果不想用队列和tee，可自定义分发器，利用回调函数实时传递元素，确保每个元素仅生成一次且无缓存：

def source_generator(callback1, callback2):
    for i in range(5):
        print(f"生成元素: {i}")
        # 实时分发给两个回调
        callback1(i)
        callback2(i)

def process_branch1(item):
    print(f"分支1处理元素: {item}")

def process_branch2(item):
    print(f"分支2处理元素: {item}")

# 执行分发
source_generator(process_branch1, process_branch2)

该方案完全无缓存，元素生成后直接传给两个分支处理，但缺点是分支处理必须同步阻塞，无法异步或延迟处理。

内容的提问来源于stack exchange，提问作者s-m-e