第4章：迭代器与生成器

一、导读

1.1 本章学习目标

第4章讲解 Python 的迭代协议和生成器：yield、生成器表达式、itertools。

本章核心学习目标：

掌握迭代协议：__iter__ 和 __next__
理解生成器：yield 和 yield from
掌握 itertools：高效的迭代工具

1.2 为什么本章如此重要

迭代器是 Python 中最强大的抽象之一。生成器让惰性求值变得简单，itertools 提供了高效的组合工具。

1.3 核心问题引导

迭代器和可迭代对象有什么区别？

yield from 如何简化嵌套生成器？

itertools 的哪些工具最常用？

二、核心概念详解

2.1 迭代协议

python

class CountDown:
    def __init__(self, start):
        self.current = start
    
    def __iter__(self):
        return self
    
    def __next__(self):
        if self.current <= 0:
            raise StopIteration
        val = self.current
        self.current -= 1
        return val

for num in CountDown(5):
    print(num)  # 5, 4, 3, 2, 1

2.2 生成器

python

# 生成器函数
def countdown(n):
    while n > 0:
        yield n
        n -= 1

# 生成器表达式
squares = (x*x for x in range(10))
print(sum(squares))  # 285

# yield from
def flatten(nested):
    for item in nested:
        if isinstance(item, list):
            yield from flatten(item)
        else:
            yield item

print(list(flatten([1, [2, [3, 4]], 5])))  # [1, 2, 3, 4, 5]

2.3 itertools

python

import itertools

# chain: 连接多个迭代器
print(list(itertools.chain([1, 2], [3, 4], [5])))
# [1, 2, 3, 4, 5]

# cycle: 无限循环
# for color in itertools.cycle(['red', 'green', 'blue']):
#     print(color)

# combinations & permutations
print(list(itertools.combinations('ABC', 2)))
# [('A', 'B'), ('A', 'C'), ('B', 'C')]

print(list(itertools.permutations('ABC', 2)))
# [('A', 'B'), ('A', 'C'), ('B', 'A'), ('B', 'C'), ('C', 'A'), ('C', 'B')]

# groupby
data = [('a', 1), ('a', 2), ('b', 3), ('b', 4)]
for key, group in itertools.groupby(data, lambda x: x[0]):
    print(f"{key}: {list(group)}")

# islice: 惰性切片
gen = (x*x for x in range(100))
print(list(itertools.islice(gen, 5, 10)))  # [25, 36, 49, 64, 81]

三、实战案例

3.1 案例一：自定义 range 生成器

python

def custom_range(*args):
    """实现类似 range 的生成器"""
    if len(args) == 1:
        start, stop, step = 0, args[0], 1
    elif len(args) == 2:
        start, stop, step = args[0], args[1], 1
    elif len(args) == 3:
        start, stop, step = args
    else:
        raise TypeError(f"Expected 1-3 args, got {len(args)}")
    
    if step == 0:
        raise ValueError("step cannot be 0")
    
    current = start
    while (step > 0 and current < stop) or (step < 0 and current > stop):
        yield current
        current += step

print(list(custom_range(5)))        # [0, 1, 2, 3, 4]
print(list(custom_range(2, 8)))     # [2, 3, 4, 5, 6, 7]
print(list(custom_range(0, 10, 2))) # [0, 2, 4, 6, 8]

四、本章小结

迭代协议：__iter__ + __next__ 驱动 for/in

生成器通过 yield 实现惰性求值

itertools 提供高效的迭代组合工具

五、思考与练习

生成器的 send() 方法如何工作？

如何实现一个支持 peek 的迭代器？