Python itertools模块实现高效迭代器链式操作

Python 标准库中的 itertools 模块是一组用于处理迭代器的快速、内存高效的工具。使用 Python 原生循环处理大量数据时，往往会因为生成中间列表而占用过多内存。itertools 通过“惰性计算”解决了这个问题，即只在需要时才生成数据，从而构建出高效的数据处理流水线。

导入 itertools 模块以开始使用。

import itertools

1. 使用无限迭代器生成序列

当需要生成无限序列时，不要使用 while True 循环配合 list，这会撑爆内存。itertools 提供了三种无限迭代器，它们生成的是一个个按需产生的值。

调用 itertools.count(start=0, step=1) 生成一个从 start 开始，以 step 为步长的无限计数器。它类似于 range()，但没有上限。

counter = itertools.count(start=10, step=-2)
print(next(counter))  # 输出 10
print(next(counter))  # 输出 8

使用 itertools.cycle(iterable) 对传入的可迭代对象进行无限循环。当迭代器耗尽时，它会从头开始重复。

colors = cycle(['red', 'green', 'blue'])
print(next(colors))  # 输出 red
print(next(colors))  # 输出 green
print(next(colors))  # 输出 blue
print(next(colors))  # 输出 red (重新开始)

应用 itertools.repeat(object, times=None) 重复生成某个对象。如果指定了 times，则重复指定次数；如果不指定，则无限重复。这在需要向 map 函数传递常量参数时非常有用。

for i in repeat("OK", 3):
    print(i)
# 输出:
# OK
# OK
# OK

2. 合并与截断迭代器

处理多个数据源或只需部分数据时，直接使用 + 号或切片操作 [:] 会强制创建副本。使用以下函数可以避免内存开销。

利用 itertools.chain(*iterables) 将多个迭代器连接起来，形成一个连续的序列。它不会复制数据，而是按顺序从各个迭代器中取出元素。

list1 = [1, 2, 3]
list2 = ['a', 'b']
combined = chain(list1, list2)
print(list(combined))  # 输出: [1, 2, 3, 'a', 'b']

采用 itertools.islice(iterable, start, stop[, step]) 对迭代器进行切片。这与列表切片类似，但支持惰性计算，且适用于没有索引的迭代器（如生成器）。

# 只取 count() 生成器的前 5 个数
limited = islice(itertools.count(), 5)
print(list(limited))  # 输出: [0, 1, 2, 3, 4]

3. 高效过滤与筛选

在复杂的逻辑判断中，内置的 filter() 可能不够用，或者写起来比较啰嗦。itertools 提供了更灵活的过滤工具。

使用 itertools.compress(data, selectors) 根据布尔值列表筛选数据。selectors 是一个包含 True/False 的可迭代对象，只有对应位置为 True 的数据才会被保留。

data = ['a', 'b', 'c', 'd']
selectors = [True, False, True, False]
result = compress(data, selectors)
print(list(result))  # 输出: ['a', 'c']

调用 itertools.dropwhile(predicate, iterable) 丢弃 predicate 函数返回 True 的元素，直到第一次返回 False 为止，然后保留后面所有的元素。

# 丢弃所有小于 5 的数，直到遇到第一个 >= 5 的数
nums = [1, 2, 3, 4, 5, 6, 1, 2]
result = dropwhile(lambda x: x < 5, nums)
print(list(result))  # 输出: [5, 6, 1, 2]

调用 itertools.takewhile(predicate, iterable) 保留 predicate 函数返回 True 的元素，一旦返回 False，立即停止迭代。

nums = [1, 2, 3, 4, 5, 6, 1, 2]
result = takewhile(lambda x: x < 5, nums)
print(list(result))  # 输出: [1, 2, 3, 4]

4. 排列与组合

在需要进行全排列或组合计算时，手写递归算法既容易出错又效率不高。

计算 itertools.product(*iterables, repeat=1) 生成笛卡尔积。相当于嵌套的 for 循环。

# 生成 ['A', 'B'] 和 [1, 2] 的笛卡尔积
result = product(['A', 'B'], [1, 2])
print(list(result))
# 输出: [('A', 1), ('A', 2), ('B', 1), ('B', 2)]

生成 itertools.permutations(iterable, r=None) 生成指定长度的所有排列（顺序不同视为不同结果）。

result = permutations('AB', 2)
print(list(result))  # 输出: [('A', 'B'), ('B', 'A')]

获取 itertools.combinations(iterable, r) 生成指定长度的所有组合（顺序不同视为相同结果）。

result = combinations('ABC', 2)
print(list(result))  # 输出: [('A', 'B'), ('A', 'C'), ('B', 'C')]

以下表格总结了这三个函数的区别：

5. 数据分组

注意：在使用分组功能前，必须 先对数据进行排序。itertools.groupby 不会自动排序，它只会把相邻的相同 key 的元素归为一组。

执行 itertools.groupby(iterable, key=None) 对连续的相同元素进行分组。

from itertools import groupby

# 错误示范：数据未排序，导致分组混乱
data = ['A', 'B', 'A', 'A']
# 正确做法：先排序
data.sort()

# 分组
grouped = groupby(data)
for key, group in grouped:
    print(f"Key: {key}, List: {list(group)}")

# 输出:
# Key: A, List: ['A', 'A', 'A']
# Key: B, List: ['B']

6. 实战案例：构建数据处理流水线

假设需要处理一个巨大的日志文件，我们要过滤出错误日志，提取特定字段，并只保留前 10 条。

定义 一个模拟的日志生成器。

import itertools

def log_generator():
    logs = [
        "INFO: System started",
        "ERROR: Disk full",
        "INFO: User logged in",
        "ERROR: Connection timeout",
        "ERROR: Null pointer",
        "WARNING: Low memory",
        "ERROR: File not found"
    ]
    for log in logs:
        yield log

# 1. 生成数据源
logs = log_generator()

# 2. 过滤：只保留包含 "ERROR" 的行
errors = filter(lambda x: "ERROR" in x, logs)

# 3. 映射：提取错误码或简单处理（去掉了 "ERROR: " 前缀）
cleaned_errors = map(lambda x: x.replace("ERROR: ", ""), errors)

# 4. 截断：只取前 3 个
first_three_errors = itertools.islice(cleaned_errors, 3)

# 5. 消费数据
print("First 3 errors found:")
for e in first_three_errors:
    print(f"- {e}")

通过 itertools，上述操作在执行过程中从未创建包含所有日志的中间列表，即使面对数百万行的日志文件，内存占用依然极低。