Python itertools.tee将一个迭代器拆分为多个独立迭代器

Python 中的迭代器是一次性资源，一旦遍历结束，就无法重新开始。这导致在需要多次遍历同一数据流，或在不同逻辑分支中处理同一序列时，直接使用原始迭代器会变得非常棘手。itertools.tee 函数正是为了解决这一痛点而设计，它能够将一个迭代器拆分为多个独立的迭代器，使得每个分支都能独立地读取数据。

基础用法：拆分迭代器

itertools.tee 接收一个可迭代对象作为输入，并返回一个包含多个独立迭代器的元组。默认情况下，它会生成 2 个迭代器，但你可以通过参数指定数量。

1. 导入 itertools 模块。
2. 定义一个简单的生成器或迭代器。这里使用生成器表达式模拟一个只能读取一次的数据流。
3. 调用 itertools.tee(iterable, n=2) 将原始迭代器拆分为 n 个独立副本。
4. 分别遍历 返回的迭代器，验证它们是否独立工作。

import itertools

# 1. 创建一个简单的生成器（模拟数据流）
def data_generator():
    for i in range(3):
        yield i

original_gen = data_generator()

# 2. 拆分为 2 个独立的迭代器
it1, it2 = itertools.tee(original_gen, 2)

# 3. 遍历第一个迭代器
print("迭代器 1 的内容:")
for val in it1:
    print(f"  - {val}")

# 4. 遍历第二个迭代器（验证它没有因为步骤3而失效）
print("迭代器 2 的内容:")
for val in it2:
    print(f"  - {val}")

执行上述代码后，你会发现 it1 和 it2 都能完整地输出 0, 1, 2。如果直接尝试对 original_gen 遍历两次，第二次将不会有任何输出。

进阶应用：成对处理数据

在实际开发中，经常需要查看“当前元素”和“下一个元素”来进行对比或计算（例如计算相邻数值的差值）。tee 可以非常优雅地实现这一功能，而无需手动索引。

1. 复制 迭代器得到两份副本 a 和 b。
2. 推进 副本 b 的一个位置，使其指向下一个元素。
3. 组合 使用 zip 函数将 a 和 b 打包在一起。

import itertools

def pairwise(iterable):
    # 1. 拆分为两个迭代器
    a, b = itertools.tee(iterable)

    # 2. 推进第二个迭代器 b 的指针，跳过第一个元素
    next(b, None)

    # 3. 返回组合后的迭代器
    return zip(a, b)

data = [10, 20, 30, 40]

# 使用自定义的 pairwise 函数
for current_item, next_item in pairwise(data):
    # 计算差值
    difference = next_item - current_item
    print(f"当前: {current_item}, 下一个: {next_item}, 差值: {difference}")

通过这种方式，你无需将列表加载到内存中通过下标 i 和 i+1 访问，这同样适用于无限长度的数据流。

核心机制与内存陷阱

使用 itertools.tee 时，必须理解其背后的数据缓冲机制，否则极易导致内存溢出。tee 并不是简单地“复制”数据，而是维护了一个单向链表作为缓冲区。

当一个迭代器（假设为 A）领先于另一个迭代器（假设为 B）时，A 消耗掉的数据会被 tee 暂时保存起来。因为 B 还没有读取这些数据，如果 B 后续需要读取，tee 必须能提供这些值。只有当 B 也读取了这些数据，它们才能从缓冲区中释放。

以下流程展示了数据在迭代器拆分后的流向与存储逻辑：

为了避免内存问题，请遵循以下规则：

1. 保持消费速度一致：尽量让拆分出来的迭代器以相近的速度消耗数据。
2. 及时消费：如果其中一个迭代器不再需要，应尽量让其运行完毕或显式销毁，以便让缓冲区释放数据。
3. 避免过度领先：绝对避免让一个迭代器跑得非常远（例如遍历了 100 万条），而另一个迭代器却停留在起点。这会导致 tee 在内存中缓存这 100 万条数据。

如果你发现处理数据时内存飙升，请检查是否存在一个迭代器“卡”在前面，而另一个迭代器在疯狂读取数据的情况。

实战场景：多路输出处理

假设你需要从同一个数据源中，同时筛选出满足不同条件的子集，并分别写入不同的文件或发送到不同的网络接口。

1. 创建 一个生产数据的源迭代器（例如读取大文件的行）。
2. 拆分 源迭代器为三份：raw_stream, check_stream, archive_stream。
3. 连接 第一个迭代器到原始日志存储。
4. 过滤 第二个迭代器，只保留包含 ERROR 的行并报警。
5. 过滤 第三个迭代器，只保留包含 Transaction 的行并存入数据库。

import itertools

# 模拟日志流
log_stream = [
    "INFO: System start",
    "ERROR: Disk full",
    "Transaction: User login",
    "Transaction: Payment",
    "ERROR: Network timeout"
]

# 1. 拆分为 3 个独立迭代器
iter_logs, iter_errors, iter_trans = itertools.tee(log_stream, 3)

# 2. 处理原始日志
print("--- 存储所有日志 ---")
for log in iter_logs:
    print(f"存储: {log}")

# 3. 处理错误日志
print("\n--- 提取错误告警 ---")
for log in iter_errors:
    if "ERROR" in log:
        print(f"告警: {log}")

# 4. 处理交易日志
print("\n--- 提取交易记录 ---")
for log in iter_trans:
    if "Transaction" in log:
        print(f"入库: {log}")

总结关键参数与行为

下表总结了 itertools.tee 的核心行为特征，帮助你在不同场景下快速决策。