第 8 章 迭代器与生成器
8.1 可迭代对象 vs 迭代器
这是 Python 中最重要的概念区分之一:
- 可迭代对象(Iterable):实现了
__iter__()方法的对象,可以放在for循环中 - 迭代器(Iterator):实现了
__iter__()和__next__()方法的对象,记住遍历位置
# list 是可迭代对象,但不是迭代器
lst = [1, 2, 3]
print(hasattr(lst, "__iter__")) # True — 可迭代
print(hasattr(lst, "__next__")) # False — 不是迭代器
# 通过 iter() 获取迭代器
it = iter(lst)
print(hasattr(it, "__next__")) # True — 是迭代器
print(next(it)) # 1
print(next(it)) # 2
print(next(it)) # 3
# next(it) # StopIteration 异常 — 耗尽了
for 循环的内部机制:
# for item in lst:
# print(item)
# 等价于:
it = iter(lst)
while True:
try:
item = next(it)
print(item)
except StopIteration:
break
关键区别:
lst = [1, 2, 3]
# 可迭代对象可以多次遍历
for x in lst: print(x, end=" ") # 1 2 3
for x in lst: print(x, end=" ") # 1 2 3 — 再次遍历
# 迭代器只能遍历一次
it = iter(lst)
for x in it: print(x, end=" ") # 1 2 3
for x in it: print(x, end=" ") # (无输出)— 已耗尽
8.2 __iter__() 与 __next__() 协议
自己实现迭代器:
class Countdown:
def __init__(self, start):
self.start = start
def __iter__(self):
return self # 迭代器返回自身
def __next__(self):
if self.start <= 0:
raise StopIteration
self.start -= 1
return self.start + 1
for n in Countdown(5):
print(n, end=" ")
# 5 4 3 2 1
更好的做法——分离可迭代对象和迭代器:
class Range:
def __init__(self, start, end):
self.start = start
self.end = end
def __iter__(self):
return RangeIterator(self.start, self.end)
class RangeIterator:
def __init__(self, current, end):
self.current = current
self.end = end
def __iter__(self):
return self
def __next__(self):
if self.current >= self.end:
raise StopIteration
value = self.current
self.current += 1
return value
r = Range(1, 4)
for x in r: print(x, end=" ") # 1 2 3
for x in r: print(x, end=" ") # 1 2 3 — 可重复遍历!
8.3 生成器函数与 yield
生成器是创建迭代器的简洁方式——用 yield 代替 return:
def countdown(n):
while n > 0:
yield n # 暂停并产出值
n -= 1
# 调用生成器函数返回的是生成器对象,不会立即执行
gen = countdown(5)
print(type(gen)) # <class 'generator'>
# 用 next() 驱动
print(next(gen)) # 5
print(next(gen)) # 4
print(next(gen)) # 3
# 用 for 循环遍历
for n in countdown(5):
print(n, end=" ") # 5 4 3 2 1
yield 的执行机制
def simple_gen():
print("开始")
yield 1
print("第一次 yield 之后")
yield 2
print("第二次 yield 之后")
yield 3
print("结束")
gen = simple_gen() # 不打印任何东西
print(next(gen)) # 打印 "开始",返回 1
print(next(gen)) # 打印 "第一次 yield 之后",返回 2
print(next(gen)) # 打印 "第二次 yield 之后",返回 3
# next(gen) # 打印 "结束",抛出 StopIteration
每次 next() 调用,函数从上次 yield 的位置恢复执行,直到遇到下一个 yield 或函数结束。
生成器的优势:惰性求值
# 列表:一次性把所有数据放入内存
squares_list = [x**2 for x in range(10_000_000)] # 占用大量内存
# 生成器:按需计算,几乎不占内存
squares_gen = (x**2 for x in range(10_000_000))
# 读取大文件
def read_large_file(filepath):
with open(filepath) as f:
for line in f: # 文件对象本身就是迭代器
yield line.strip()
# 无限序列
def natural_numbers():
n = 1
while True:
yield n
n += 1
# 取前 10 个
from itertools import islice
first_10 = list(islice(natural_numbers(), 10))
# [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]
8.4 生成器表达式
生成器表达式的语法与列表推导式几乎相同,只是用圆括号:
# 列表推导式 — 立即计算,存入内存
squares_list = [x**2 for x in range(1000)]
# 生成器表达式 — 惰性求值
squares_gen = (x**2 for x in range(1000))
# 作为函数参数时可以省略外层括号
total = sum(x**2 for x in range(1000))
max_val = max(len(word) for word in words)
# 链式处理
import os
py_files = (
entry.name
for entry in os.scandir(".")
if entry.is_file() and entry.name.endswith(".py")
)
8.5 yield from 委托生成器
yield from 用于将一个生成器的产出委托给另一个可迭代对象:
# 不用 yield from
def chain_manual(*iterables):
for it in iterables:
for item in it:
yield item
# 用 yield from — 更简洁
def chain(*iterables):
for it in iterables:
yield from it
list(chain([1, 2], [3, 4], [5, 6]))
# [1, 2, 3, 4, 5, 6]
遍历嵌套结构:
def flatten(nested):
for item in nested:
if isinstance(item, (list, tuple)):
yield from flatten(item) # 递归委托
else:
yield item
data = [1, [2, 3, [4, 5]], 6, [7, [8, 9]]]
print(list(flatten(data)))
# [1, 2, 3, 4, 5, 6, 7, 8, 9]
yield from 还能传递 send/throw 和获取返回值:
def accumulate():
total = 0
while True:
value = yield total
if value is None:
break
total += value
return total # 生成器的返回值
def main():
result = yield from accumulate()
print(f"总计: {result}")
gen = main()
next(gen) # 启动
gen.send(10) # total = 10
gen.send(20) # total = 30
try:
gen.send(None) # 触发 break
except StopIteration:
pass
# 输出: 总计: 30
8.6 生成器的 send()、throw()、close()
生成器不仅能产出数据,还能接收数据,形成双向通信:
send() — 向生成器发送数据
def echo():
while True:
received = yield # yield 表达式的值就是 send() 的参数
print(f"收到: {received}")
gen = echo()
next(gen) # 必须先 next() 启动(推进到第一个 yield)
gen.send("hello") # 收到: hello
gen.send("world") # 收到: world
# 实用:运行平均值计算器
def running_average():
total = 0
count = 0
average = None
while True:
value = yield average
total += value
count += 1
average = total / count
avg = running_average()
next(avg) # None(启动)
avg.send(10) # 10.0
avg.send(20) # 15.0
avg.send(30) # 20.0
throw() — 向生成器抛入异常
def careful_gen():
try:
while True:
value = yield
print(f"处理: {value}")
except ValueError as e:
print(f"捕获异常: {e}")
yield "error_handled"
gen = careful_gen()
next(gen)
gen.send("data") # 处理: data
result = gen.throw(ValueError, "无效数据") # 捕获异常: 无效数据
print(result) # error_handled
close() — 关闭生成器
def resource_gen():
print("获取资源")
try:
yield "resource"
finally:
print("释放资源") # close() 时会执行 finally
gen = resource_gen()
print(next(gen)) # 获取资源 \n resource
gen.close() # 释放资源
8.7 itertools 模块精讲
itertools 提供了高效的迭代器工具:
无限迭代器
from itertools import count, cycle, repeat
# count — 无限计数
for i in count(start=10, step=2):
if i > 20: break
print(i, end=" ") # 10 12 14 16 18 20
# cycle — 无限循环
colors = cycle(["red", "green", "blue"])
for _, color in zip(range(7), colors):
print(color, end=" ") # red green blue red green blue red
# repeat — 重复
list(repeat("hello", 3)) # ['hello', 'hello', 'hello']
终止迭代器
from itertools import (
chain, islice, takewhile, dropwhile,
accumulate, groupby, starmap, compress, filterfalse
)
# chain — 串联多个可迭代对象
list(chain([1, 2], [3, 4], [5, 6]))
# [1, 2, 3, 4, 5, 6]
# islice — 切片迭代器(不支持负索引)
list(islice(count(), 5, 10)) # [5, 6, 7, 8, 9]
list(islice(count(), 0, 10, 2)) # [0, 2, 4, 6, 8]
# takewhile — 条件为真时持续取值
list(takewhile(lambda x: x < 5, [1, 3, 5, 2, 4]))
# [1, 3] — 遇到 5 时停止,即使后面还有 < 5 的值
# dropwhile — 跳过条件为真的前缀
list(dropwhile(lambda x: x < 5, [1, 3, 5, 2, 4]))
# [5, 2, 4]
# accumulate — 累积
list(accumulate([1, 2, 3, 4, 5]))
# [1, 3, 6, 10, 15] — 前缀和
import operator
list(accumulate([1, 2, 3, 4, 5], operator.mul))
# [1, 2, 6, 24, 120] — 前缀积
# groupby — 分组(数据需要预先排序)
from itertools import groupby
data = [("A", 1), ("A", 2), ("B", 3), ("B", 4), ("A", 5)]
data.sort(key=lambda x: x[0]) # 先排序
for key, group in groupby(data, key=lambda x: x[0]):
print(f"{key}: {list(group)}")
# A: [('A', 1), ('A', 2), ('A', 5)]
# B: [('B', 3), ('B', 4)]
# compress — 根据选择器过滤
list(compress("ABCDEF", [1, 0, 1, 0, 1, 1]))
# ['A', 'C', 'E', 'F']
# starmap — 解包参数后 map
list(starmap(pow, [(2, 3), (3, 2), (10, 3)]))
# [8, 9, 1000]
排列组合
from itertools import product, permutations, combinations, combinations_with_replacement
# product — 笛卡尔积
list(product("AB", "12"))
# [('A', '1'), ('A', '2'), ('B', '1'), ('B', '2')]
list(product(range(2), repeat=3))
# [(0,0,0), (0,0,1), (0,1,0), (0,1,1), (1,0,0), (1,0,1), (1,1,0), (1,1,1)]
# permutations — 排列
list(permutations("ABC", 2))
# [('A','B'), ('A','C'), ('B','A'), ('B','C'), ('C','A'), ('C','B')]
# combinations — 组合
list(combinations("ABCD", 2))
# [('A','B'), ('A','C'), ('A','D'), ('B','C'), ('B','D'), ('C','D')]
# combinations_with_replacement — 允许重复的组合
list(combinations_with_replacement("ABC", 2))
# [('A','A'), ('A','B'), ('A','C'), ('B','B'), ('B','C'), ('C','C')]
实用模式
# 分块读取
def chunked(iterable, n):
it = iter(iterable)
while chunk := list(islice(it, n)):
yield chunk
list(chunked(range(10), 3))
# [[0, 1, 2], [3, 4, 5], [6, 7, 8], [9]]
# 滑动窗口(Python 3.12+)
from itertools import pairwise
list(pairwise([1, 2, 3, 4, 5]))
# [(1, 2), (2, 3), (3, 4), (4, 5)]
# 通用滑动窗口
from collections import deque
def sliding_window(iterable, n):
it = iter(iterable)
window = deque(islice(it, n), maxlen=n)
if len(window) == n:
yield tuple(window)
for item in it:
window.append(item)
yield tuple(window)
list(sliding_window([1, 2, 3, 4, 5], 3))
# [(1, 2, 3), (2, 3, 4), (3, 4, 5)]
本章小结:迭代器和生成器是 Python 的核心概念。生成器用
yield实现惰性求值,节省内存又简洁优雅。yield from简化了生成器的委托。itertools提供了丰富的迭代器工具,是处理序列数据的利器。理解这些概念后,你会发现 Python 中到处都是迭代器协议的应用。