生成器与协程
00:00
Python生成器与协程详解:yield、send、yield from。
概述
生成器(Generator)是 Python 中使用 yield 语句实现的惰性迭代器,它可以按需产生值,避免一次性加载所有数据到内存。协程(Coroutine)是基于生成器的协作式并发机制,通过 yield 和 send 实现双向通信。理解生成器与协程是掌握 Python 高级编程的关键。
基础概念
生成器函数与生成器对象
包含 yield 语句的函数称为生成器函数,调用它不会执行函数体,而是返回一个生成器对象。生成器对象实现了迭代器协议:
def count_up(n):
"""简单的生成器函数"""
i = 0
while i < n:
yield i
i += 1
# 调用生成器函数返回生成器对象
gen = count_up(5)
print(type(gen)) # <class 'generator'>
# 迭代获取值
for value in gen:
print(value) # 0, 1, 2, 3, 4
yield 的工作原理
yield 暂停函数执行并返回一个值。下次调用 next() 时,从 yield 的下一行继续执行:
def demo():
print("第一步")
yield 1
print("第二步")
yield 2
print("第三步")
return "结束"
gen = demo()
print(next(gen)) # 输出: 第一步 → 返回 1
print(next(gen)) # 输出: 第二步 → 返回 2
# print(next(gen)) # 输出: 第三步 → StopIteration: 结束
快速上手
生成器表达式
# 列表推导式:一次性创建所有元素
squares_list = [x**2 for x in range(1000000)] # 占用大量内存
# 生成器表达式:按需生成元素
squares_gen = (x**2 for x in range(1000000)) # 几乎不占内存
# 使用
total = sum(squares_gen) # 逐个计算,不需要全部存储
读取大文件
def read_large_file(filepath, chunk_size=8192):
"""逐块读取大文件"""
with open(filepath, 'r', encoding='utf-8') as f:
while True:
chunk = f.read(chunk_size)
if not chunk:
break
yield chunk
# 使用:不会一次性加载整个文件到内存
for chunk in read_large_file("large_log.txt"):
process(chunk)
详细用法
send 方法
send 方法向生成器发送值,yield 表达式的返回值就是 send 传入的值:
def accumulator():
"""累加器协程"""
total = 0
while True:
value = yield total # yield 返回 total,接收 send 的值
if value is None:
break
total += value
gen = accumulator()
next(gen) # 启动生成器,返回 0
print(gen.send(10)) # 返回 10
print(gen.send(20)) # 返回 30
print(gen.send(5)) # 返回 35
throw 与 close
def robust_generator():
"""可以处理异常的生成器"""
try:
while True:
value = yield
print(f"收到: {value}")
except GeneratorExit:
print("生成器被关闭")
except ValueError as e:
print(f"收到异常: {e}")
gen = robust_generator()
next(gen) # 启动
gen.send("hello") # 收到: hello
gen.throw(ValueError, "测试异常") # 收到异常: 测试异常
gen.close() # 生成器被关闭
yield from
yield from 将迭代委托给另一个生成器:
def chain(*iterables):
"""将多个可迭代对象串联"""
for iterable in iterables:
yield from iterable
# 使用
list(chain([1, 2], [3, 4], [5])) # [1, 2, 3, 4, 5]
# yield from 也可以转发 send 和 throw
def inner():
result = yield "inner"
print(f"inner 收到: {result}")
return "inner_done"
def outer():
result = yield from inner() # 委托给 inner,并获取其返回值
print(f"inner 返回: {result}")
gen = outer()
print(next(gen)) # inner
print(gen.send("hi")) # inner 收到: hi → inner 返回: inner_done → StopIteration
生成器管道
def read_lines(filepath):
"""读取文件行"""
with open(filepath) as f:
yield from f
def filter_comments(lines):
"""过滤注释行"""
for line in lines:
if not line.strip().startswith('#'):
yield line
def strip_lines(lines):
"""去除空白"""
for line in lines:
yield line.strip()
def parse_csv(lines):
"""解析 CSV 行"""
for line in lines:
if line:
yield line.split(',')
# 组合管道
pipeline = parse_csv(strip_lines(filter_comments(read_lines("data.csv"))))
for row in pipeline:
print(row)
协程实现状态机
def tcp_connection():
"""TCP 连接状态机"""
state = "CLOSED"
while True:
event = yield state
if state == "CLOSED" and event == "connect":
state = "CONNECTING"
elif state == "CONNECTING" and event == "connected":
state = "ESTABLISHED"
elif state == "ESTABLISHED" and event == "close":
state = "CLOSED"
elif state == "ESTABLISHED" and event == "error":
state = "ERROR"
conn = tcp_connection()
next(conn) # 启动,返回 "CLOSED"
print(conn.send("connect")) # "CONNECTING"
print(conn.send("connected")) # "ESTABLISHED"
print(conn.send("close")) # "CLOSED"
常见场景
场景一:分页数据获取
def fetch_all_pages(api, page_size=100):
"""分页获取所有数据"""
page = 1
while True:
response = api.get_page(page, page_size)
yield from response.items
if not response.has_next:
break
page += 1
# 使用
all_items = list(fetch_all_pages(api))
场景二:树遍历
def traverse_tree(node):
"""深度优先遍历树"""
yield node
for child in node.children:
yield from traverse_tree(child)
# 使用
for node in traverse_tree(root):
print(node.name)
场景三:协程实现生产者-消费者
def producer(target, items):
"""生产者"""
for item in items:
target.send(item)
target.close()
def consumer():
"""消费者协程"""
while True:
try:
item = yield
print(f"处理: {item}")
except GeneratorExit:
print("消费者结束")
break
c = consumer()
next(c) # 启动消费者
producer(c, ["任务1", "任务2", "任务3"])
注意事项
- 生成器只能迭代一次,迭代完毕后再次迭代不会产生任何值
- 首次调用 send 前必须先调用 next() 启动生成器(或 send(None))
- yield from 会将子生成器的返回值作为 yield from 表达式的值
- 生成器不是线程安全的,不要在多个线程中同时操作同一个生成器
- Python 3.5+ 的 async/await 是协程的现代实现,基于生成器的协程已不推荐使用
- 生成器在处理大数据时特别有用,但不要过度使用导致代码难以理解
进阶用法
生成器与 itertools 配合
from itertools import islice, chain, groupby
# 分批处理
def batched(iterable, size):
"""将可迭代对象分批"""
it = iter(iterable)
while True:
batch = list(islice(it, size))
if not batch:
break
yield batch
# 使用
for batch in batched(range(100), 10):
print(f"批次: {batch}")
生成器实现无限序列
def fibonacci():
"""无限斐波那契数列"""
a, b = 0, 1
while True:
yield a
a, b = b, a + b
# 取前 10 个
from itertools import islice
first_10 = list(islice(fibonacci(), 10))
print(first_10) # [0, 1, 1, 2, 3, 5, 8, 13, 21, 34]
使用 return 返回值
def search(items, predicate):
"""搜索满足条件的元素,返回是否找到"""
found = False
for item in items:
if predicate(item):
found = True
yield item
return found # 生成器返回值通过 StopIteration.value 获取
gen = search([1, 2, 3, 4, 5], lambda x: x > 3)
print(list(gen)) # [4, 5]
# 获取返回值
gen = search([1, 2, 3], lambda x: x > 5)
try:
next(gen)
except StopIteration as e:
print(f"搜索结果: {e.value}") # False
异步生成器
async def async_read_lines(filepath):
"""异步逐行读取文件"""
async with aiofiles.open(filepath) as f:
async for line in f:
yield line.strip()
# 使用
async for line in async_read_lines("data.txt"):
print(line)