代码收藏家技术教程 2025-02-15

Python yield 详解

Python `yield` 详解

yield 是 Python 中用于定义生成器函数的关键字。与返回整个结果列表的普通函数不同，生成器函数通过 yield 逐个产生值，允许你在迭代过程中逐步获取数据，而不是一次性将所有数据加载到内存中。这种方式对于处理大型数据集或创建无限序列非常有用，因为它可以显著减少内存占用并提高性能。

`yield` 的基本用法

一个简单的生成器函数示例如下：

def simple_generator():
    yield 1
    yield 2
    yield 3

# 使用 for 循环遍历生成器
for value in simple_generator():
    print(value)

上述代码会依次打印出 1, 2, 3。每次遇到 yield 语句时，函数会暂停执行，并返回 yield 后面的值。当再次调用该函数（例如通过迭代）时，它会从上次暂停的地方继续执行，直到遇到下一个 yield 或者函数结束。

生成器对象

当你调用一个包含 yield 的函数时，实际上返回的是一个生成器对象，而不是立即执行函数体内的代码。你可以通过 next() 函数来手动获取生成器中的下一个值：

gen = simple_generator()
print(next(gen))  # 输出: 1
print(next(gen))  # 输出: 2
print(next(gen))  # 输出: 3
# print(next(gen))  # 这行会引发 StopIteration 异常

一旦生成器的所有值都被消费完毕，再次调用 next() 将会抛出 StopIteration 异常，表示没有更多的值可以产生了。

避免 `StopIteration` 异常

为了避免 StopIteration 异常，可以在调用 next() 时提供一个默认值，或者使用 try-except 块来捕获异常。以下是两种常见的做法：

提供默认值

如果你不希望程序在生成器耗尽时抛出异常，可以为 next() 提供一个默认值。如果生成器没有更多值可返回，next() 将返回这个默认值，而不是抛出异常。

gen = simple_generator()
print(next(gen, None))  # 输出: 1
print(next(gen, None))  # 输出: 2
print(next(gen, None))  # 输出: 3
print(next(gen, None))  # 输出: None

使用 `try-except` 捕获异常

另一种方法是使用 try-except 块来捕获 StopIteration 异常，从而优雅地处理生成器耗尽的情况。

gen = simple_generator()
while True:
    try:
        value = next(gen)
        print(value)
    except StopIteration:
        print("Generator has been exhausted.")
        break

这两种方法都可以有效避免 StopIteration 异常，选择哪种取决于你的具体需求和代码风格。

生成器表达式

类似于列表推导式，Python 也支持生成器表达式，它提供了一种简洁的方式来创建生成器。生成器表达式的语法与列表推导式相似，只是使用圆括号 () 而不是方括号 []：

gen_expr = (x * x for x in range(5))
for num in gen_expr:
    print(num)  # 依次输出: 0, 1, 4, 9, 16

生成器的优点

节省内存：生成器不会一次性创建所有的值，而是在需要的时候才生成，因此非常适合处理大规模数据。

惰性求值：只有在迭代时才会计算下一个值，这使得程序可以在任何时候停止或继续，而不必等待所有数据都准备好。

流式处理：生成器非常适合处理流式数据，如文件读取、网络请求等场景。

`yield` 与 `return` 的区别

return 用于从函数中返回一个值，并终止函数的执行。而 yield 只是暂停函数的执行，并保存当前状态，以便后续可以从中断处恢复。

如果生成器函数中有 return 语句（Python 3.3+），它可以用来传递一个最终值给生成器的使用者，但这个值会被封装在一个 StopIteration 异常中。此外，return 也可以用来提前终止生成器，不再产生新的值。

发送值给生成器

除了产生值，生成器还可以接收外部发送的值。这可以通过 send() 方法实现，它允许你向生成器传递数据，并将其作为 yield 表达式的值：

def echo():
    while True:
        received = yield
        print(f"Received: {received}")

gen = echo()
next(gen)  # 必须先启动生成器
gen.send("Hello")  # 输出: Received: Hello
gen.send("World")  # 输出: Received: World

注意，必须先调用 next() 来启动生成器，否则 send() 会抛出 TypeError。

抛出异常和关闭生成器

你可以使用 throw() 方法向生成器抛出异常，或者使用 close() 方法显式地关闭生成器。一旦生成器被关闭，任何进一步的调用都会导致 StopIteration 或 GeneratorExit 异常：

def my_generator():
    try:
        yield 1
        yield 2
        yield 3
    except GeneratorExit:
        print("Generator was closed")

gen = my_generator()
print(next(gen))  # 输出: 1
gen.close()       # 输出: Generator was closed
# next(gen)      # 这行会引发 StopIteration 异常