如何有条件拆分Python生成器管道为两分支并同时消费？

需要用Python读取大文本文件，并根据条件将数据写入两个不同的数据库表，希望通过生成器管道避免将所有数据加载到内存中。

等价逻辑：将生成整数的生成器拆分为奇数生成器和偶数生成器，再分别写入不同文件。

初始实现（使用itertools.tee） #

尝试用itertools.tee拆分生成器，代码如下：

def generate_numbers(limit):
    # 实际场景的生成器逻辑与此结构一致
    for i in range(limit):
        yield i

numbers1, numbers2 = itertools.tee(generate_numbers(10), n=2)
evens = (num for num in numbers1 if num % 2 == 0)
odds = (num for num in numbers2 if num %2 != 0)

运行后能得到预期结果：

>>> list(evens)                                                                              
[0, 2, 4, 6, 8]
>>> list(odds)                                                                               
[1, 3, 5, 7, 9]

遇到的问题 #

但itertools.tee文档警告：如果其中一个迭代器先被消费大部分或全部数据，会存储大量临时数据。而按以下方式写入文件时就会触发这个问题：

def write_to_file(filename, numbers):
    # 实际场景的消费函数接收可迭代对象作为参数
    with open(filename, 'wt') as outfile:
        for i in numbers:
            outfile.write(f"{i}\n")

write_to_file('evens.txt', evens)
write_to_file('odds.txt', odds)

疑问 #

• 是否可以同时消费这两个生成器？
• itertools.tee不是线程安全的，能否用asyncio实现？
• 有没有其他替代方案？分块处理是否有用？

核心约束：不想将单个表的所有数据保存在内存中，且消费函数期望接收可迭代对象而非单个元素。

类似问题参考 #

有类似问题的主流解决方案是两次遍历输入文件，但更希望实现单次遍历；另一种方案是同时打开两个输出文件逐个处理元素，但不符合消费函数需要读取整个迭代器的需求。

编辑补充：找到旧答案指出因内存问题无法用itertools.tee实现，求替代方法。

方案1：多线程队列分流（单次遍历+内存可控） #

用队列分发数据，让两个消费线程并行处理，既保证单次遍历输入，又通过队列大小限制内存占用，同时适配消费函数接收可迭代对象的要求。

代码示例：

import queue
import threading

def generate_numbers(limit):
    for i in range(limit):
        yield i

def queue_iter(q):
    # 封装队列迭代器，适配消费函数的可迭代对象要求
    while True:
        item = q.get()
        if item is None:  # 终止信号
            break
        yield item
        q.task_done()

def write_to_file(filename, numbers):
    with open(filename, 'wt') as outfile:
        for i in numbers:
            outfile.write(f"{i}\n")

def main():
    # 设置队列最大长度，控制内存占用
    even_q = queue.Queue(maxsize=100)
    odd_q = queue.Queue(maxsize=100)

    # 启动消费线程，传入队列迭代器
    even_thread = threading.Thread(target=write_to_file, args=('evens.txt', queue_iter(even_q)))
    odd_thread = threading.Thread(target=write_to_file, args=('odds.txt', queue_iter(odd_q)))
    even_thread.start()
    odd_thread.start()

    # 遍历生成器，分发数据到对应队列
    for num in generate_numbers(10):
        if num % 2 == 0:
            even_q.put(num)
        else:
            odd_q.put(num)

    # 发送终止信号
    even_q.put(None)
    odd_q.put(None)

    # 等待线程处理完毕
    even_thread.join()
    odd_thread.join()

if __name__ == "__main__":
    main()

方案2：协程式同步分流（无线程） #

将消费函数改造为协程生成器，在遍历输入时逐个分发数据，无需线程也能实现单次遍历、低内存占用。

代码示例：

def generate_numbers(limit):
    for i in range(limit):
        yield i

def write_to_file_gen(filename):
    # 改造消费函数为协程生成器，支持逐个接收元素
    with open(filename, 'wt') as outfile:
        while True:
            num = yield
            if num is None:
                break
            outfile.write(f"{num}\n")

def main():
    # 初始化并启动协程消费者
    even_writer = write_to_file_gen('evens.txt')
    odd_writer = write_to_file_gen('odds.txt')
    next(even_writer)
    next(odd_writer)

    # 遍历输入，分发数据
    for num in generate_numbers(10):
        if num % 2 == 0:
            even_writer.send(num)
        else:
            odd_writer.send(num)

    # 终止协程
    even_writer.send(None)
    odd_writer.send(None)

if __name__ == "__main__":
    main()

方案3：分块处理（折中适配原有消费函数） #

如果不想修改消费函数或引入线程/协程，可以分块处理输入：每次读取一小块数据，用itertools.tee拆分后分别写入，控制tee的缓存数据量不超过块大小。

代码示例：

import itertools

def generate_numbers(limit):
    for i in range(limit):
        yield i

def write_to_file(filename, numbers):
    with open(filename, 'wt') as outfile:
        for i in numbers:
            outfile.write(f"{i}\n")

def chunked_iter(iterable, chunk_size):
    # 按块拆分迭代器
    iterator = iter(iterable)
    while True:
        chunk = list(itertools.islice(iterator, chunk_size))
        if not chunk:
            break
        yield chunk

def main():
    chunk_size = 5  # 根据内存情况调整块大小
    # 保持文件打开状态，避免多次IO操作
    with open('evens.txt', 'wt') as even_f, open('odds.txt', 'wt') as odd_f:
        def write_chunk_to_file(file, numbers):
            for i in numbers:
                file.write(f"{i}\n")
        
        for chunk in chunked_iter(generate_numbers(10), chunk_size):
            chunk1, chunk2 = itertools.tee(chunk)
            evens = (num for num in chunk1 if num %2 ==0)
            odds = (num for num in chunk2 if num %2 !=0)
            write_chunk_to_file(even_f, evens)
            write_chunk_to_file(odd_f, odds)

if __name__ == "__main__":
    main()

内容的提问来源于stack exchange，提问作者Dr John A Stevenson