Python 迭代器与生成器高级用法及实战应用
Python 的迭代器和生成器是语言中最优雅的特性之一。它们不仅提供了简洁的数据遍历方式,更重要的是在处理大规模数据时能够显著降低内存消耗。本文将从原理到实战,全面掌握迭代器与生成器的高级用法。
一、迭代器的核心概念
迭代器是 Python 数据遍历的统一接口。任何实现了 `__iter__()` 和 `__next__()` 方法的对象都是迭代器。`__iter__()` 返回迭代器对象本身,`__next__()` 返回下一个值,当没有更多数据时抛出 `StopIteration` 异常。
让我们通过一个斐波那契数列迭代器来理解其工作原理:
class FibonacciIterator:
"""自定义斐波那契数列迭代器"""
def __init__(self, max_value=None):
self.max_value = max_value
self.a, self.b = 00, 1
def __iter__(self):
self.a, self.b = 00, 1 # 重置状态
return self
def __next__(self):
if self.max_value and self.a > self.max_value:
raise StopIteration
value = self.a
self.a, self.b = self.b, self.a self.b
return value
# 使用示例
fib = FibonacciIterator(100)
for num in fib:
print(num, end=' ')
# 输出: 0 1 1 2 3 5 8 13 21 34 55 89
自定义迭代器的优势在于:惰性计算(只生成需要的值)、无限序列表示(通过省略 max_value)、内存效率高。
二、生成器的基础与进阶
生成器是创建迭代器的简化方式,使用 `yield` 关键字即可。与普通函数不同,生成器在每次 `yield` 时暂停执行,保存当前状态,下次调用时从暂停处继续。
基础生成器
import random
def random_number_generator():
"""无限生成 0-1 之间的随机数"""
while True:
yield random.random()
带参数的生成器
生成器可以接受参数,实现更灵活的数据处理:
def file_line_reader(filename, chunk_size=1024):
"""按块读取文件,逐行生成,内存高效"""
with open(filename, 'r', encoding='utf-8') as f:
buffer = ''
while True:
chunk = f.read(chunk_size)
if not chunk:
if buffer:
yield buffer
break
buffer = chunk
lines = buffer.split('\n')
buffer = lines.pop() # 保留不完整的行
for line in lines:
yield line
这个文件读取器的妙处在于:即使处理 GB 级别的文件,内存占用也仅限于 chunk_size 的大小。
生成器表达式
类似于列表推导式,但使用圆括号,返回生成器而非列表:
def is_prime(n):
"""判断是否为质数"""
if n < 2:
return False
for i in range(2, int(n ** 0.5) 1):
if n % i == 0:
return False
return True
# 列表推导式 - 立即计算,占用内存
primes_list = [n for n in range(2, 1000) if is_prime(n)]
# 生成器表达式 - 惰性计算,节省内存
primes_gen = (n for n in range(2, 1000) if is_prime(n))
当数据量大或只需部分结果时,生成器表达式是更好的选择。
三、生成器链:数据流水线
生成器的真正威力在于可以组合使用,形成数据处理流水线。每个生成器只负责一个简单的转换任务,串联起来实现复杂处理。
def read_data():
"""模拟读取数据"""
data = [' hello ', ' world ', ' python ', ' generator ']
for item in data:
yield item
def strip_whitespace(data_gen):
"""去除空白字符"""
for item in data_gen:
yield item.strip()
def to_uppercase(data_gen):
"""转大写"""
for item in data_gen:
yield item.upper()
def filter_short(data_gen, min_length=5):
"""过滤短字符串"""
for item in data_gen:
if len(item) >= min_length:
yield item
# 构建处理链
gen = filter_short(
to_uppercase(
strip_whitespace(
read_data()
)
),
min_length=5
)
for item in gen:
print(item)
# 输出:
# WORLD
# PYTHON
# GENERATOR
这种模式的优势:清晰的数据流、易于测试每个环节、可以动态替换处理步骤。
四、yield from:委托生成器
Python 3.3 引入的 `yield from` 语法允许一个生成器委托部分工作给其他生成器,避免了嵌套循环的繁琐。
def sub_generator1():
yield from [1, 2, 3]
def sub_generator2():
yield from [4, 5, 6]
def main_generator():
"""使用 yield from 委托给子生成器"""
yield from sub_generator1()
yield from sub_generator2()
yield from [7, 8, 9]
print(list(main_generator()))
# 输出: [1, 2, 3, 4, 5, 6, 7, 8, 9]
`yield from` 不仅传递数据,还双向传递 `send()` 和 `throw()` 的值,实现真正的生成器委托。
五、生成器的双向通信:send() 和 throw()
生成器不仅产出数据,还可以接收外部输入,实现协程式的协作。
def interactive_generator():
"""交互式生成器,接收外部输入"""
total = 0
received = yield total # 初始化,等待接收值
while True:
if received is None:
received = 0
total = received
received = yield total
# 使用示例
gen = interactive_generator()
print(next(gen)) # 输出: 0(启动生成器)
print(gen.send(10)) # 输出: 10
print(gen.send(20)) # 输出: 30
print(gen.send(30)) # 输出: 60
这种模式适用于:状态机、累积计算、交互式协议处理等场景。
六、生成器作为上下文管理器
将生成器与上下文管理器结合,可以确保资源正确释放:
class GeneratorContextManager:
"""将生成器包装为上下文管理器"""
def __init__(self, gen_func):
self.gen_func = gen_func
def __enter__(self):
self.gen = self.gen_func()
next(self.gen) # 启动生成器
return self.gen
def __exit__(self, exc_type, exc_val, exc_tb):
self.gen.close()
return False
# 使用示例
def resource_handler():
resource = open('data.txt', 'r')
try:
yield resource
finally:
resource.close()
with GeneratorContextManager(resource_handler) as f:
data = f.read()
七、实战应用场景
场景1:合并多个有序迭代器
def merge_sorted_iterators(*iterators):
"""合并多个已排序的迭代器,保持有序"""
import heapq
heap = []
# 将第一个元素加入堆
for i, it in enumerate(iterators):
try:
value = next(it)
heap.append((value, i))
except StopIteration:
pass
heapq.heapify(heap)
while heap:
value, i = heapq.heappop(heap)
yield value
try:
next_value = next(iterators[i])
heapq.heappush(heap, (next_value, i))
except StopIteration:
pass
这个函数的时间复杂度是 O(n log k),其中 n 是总元素数,k 是迭代器数量。
场景2:滑动窗口处理
from collections import deque
def sliding_window(data, window_size=3):
"""创建滑动窗口"""
window = deque(maxlen=window_size)
for item in data:
window.append(item)
if len(window) == window_size:
yield list(window)
# 使用示例
data = [1, 2, 3, 4, 5, 6]
for window in sliding_window(data, 3):
print(window)
# 输出:
# [1, 2, 3]
# [2, 3, 4]
# [3, 4, 5]
# [4, 5, 6]
滑动窗口在时间序列分析、移动平均计算、模式匹配等场景非常有用。
场景3:批处理生成器
def batch_processor(items, batch_size=10):
"""将数据分批处理"""
batch = []
for item in items:
batch.append(item)
if len(batch) >= batch_size:
yield batch
batch = []
if batch:
yield batch
这个生成器适用于:批量数据库插入、分块 API 调用、批量文件写入等操作。
场景4:树遍历生成器
class TreeNode:
"""简单的树节点"""
def __init__(self, value, left=None, right=None):
self.value = value
self.left = left
self.right = right
def tree_inorder_traversal(root):
"""中序遍历生成器"""
if root is None:
return
yield from tree_inorder_traversal(root.left)
yield root.value
yield from tree_inorder_traversal(root.right)
# 使用示例
root = TreeNode(4,
TreeNode(2,
TreeNode(1),
TreeNode(3)),
TreeNode(6,
TreeNode(5),
TreeNode(7)))
print(list(tree_inorder_traversal(root)))
# 输出: [1, 2, 3, 4, 5, 6, 7]
生成器实现的树遍历无需维护额外的栈结构,代码更简洁。
场景5:协程式数据处理器
def data_filter(min_value=0, max_value=100):
"""数据过滤器协程"""
while True:
data = yield
if min_value <= data <= max_value:
print(f"通过过滤: {data}")
else:
print(f"被过滤: {data}")
# 使用示例
filter = data_filter(10, 50)
next(filter) # 启动协程
filter.send(25) # 通过过滤: 25
filter.send(75) # 被过滤: 75
这种协程模式可以构建复杂的数据处理管道,每个环节独立且可测试。
场景6:生成器实现限流
def rate_limiter(max_calls=5, time_window=1.0):
"""限流生成器"""
import time
from collections import deque
calls = deque()
while True:
now = time.time()
# 移除过期的调用记录
while calls and now - calls[0] > time_window:
calls.popleft()
if len(calls) < max_calls:
calls.append(now)
yield True
else:
# 等待最早的调用过期
time.sleep(calls[0] time_window - now)
calls.append(time.time())
yield True
# 使用示例
import time
limiter = rate_limiter(max_calls=3, time_window=1.0)
for i in range(10):
next(limiter)
print(f"Call {i 1} at {time.time():.2f}")
限流生成器可以保护 API 端点、数据库查询等资源不被过度消耗。
八、性能对比与最佳实践
内存使用对比
import sys
# 列表存储
numbers_list = [i for i in range(1000000)]
print(f"列表内存: {sys.getsizeof(numbers_list) / 1024 / 1024:.2f} MB")
# 输出: 列表内存: 8.00 MB
# 生成器存储
numbers_gen = (i for i in range(1000000))
print(f"生成器内存: {sys.getsizeof(numbers_gen)} bytes")
# 输出: 生成器内存: 200 bytes
生成器在内存占用上有巨大的优势,特别是处理大数据集时。
最佳实践总结
1. 何时使用迭代器/生成器:
- 处理大型数据集或无限序列
- 需要惰性计算,按需生成数据
- 构建数据处理流水线
- 实现自定义遍历逻辑
2. 避免的反模式:
- 对小数据集使用生成器(复杂性高于收益)
- 重复消费同一生成器(生成器只能遍历一次)
- 在生成器中修改共享状态(可能导致难以调试的问题)
3. 调试技巧:
- 使用 `list(gen)` 快速查看生成器内容(注意内存)
- 添加打印语句追踪 yield 位置
- 使用 `inspect.getgeneratorstate()` 检查生成器状态
4. 与异步编程结合:
- Python 3.6 的 `async for` 可以遍历异步生成器
- `asyncio` 库提供了异步生成器的完整支持
- 在异步 I/O 密集型任务中,异步生成器非常高效
九、总结
迭代器和生成器是 Python 中不可或缺的高级特性。通过本文的学习,我们掌握了:
- 自定义迭代器的实现原理
- 生成器的各种用法:基础、带参数、生成器表达式
- 生成器链和 `yield from` 委托模式
- 生成器的双向通信能力(send/throw)
- 丰富的实战应用场景:有序合并、滑动窗口、批处理、树遍历等
在实际项目中,合理使用迭代器和生成器能够:
- 降低内存消耗,处理更大的数据集
- 简化代码逻辑,提高可读性
- 构建灵活的数据处理流水线
- 实现高效的协程协作
建议在日常编码中多思考:"这个场景是否适合用生成器?",你会发现代码变得更简洁、更高效。
