Python函数设计：何时用return，何时用yield？

在Python开发中，我们常常面临一个选择：定义一个函数时，究竟使用 return 返回完整结果，还是使用 yield 逐个产出值？这两者看似只是语法差异，实则背后是两种截然不同的编程思路和内存模型。

一个最直观的例子就是编写一个筛选偶数的函数。如果你想得到一个完整的列表，最直接的写法是：

def find_even(nums):
    result = []
    for n in nums:
        if n % 2 == 0:
            result.append(n)
    return result

这个函数逻辑清晰：遍历输入，符合条件的放入列表，最后将整个列表 return 给调用者。调用者拿到的是一个“已经计算完毕”的集合。然而，当数据量极大时，例如需要处理一个几千万行的日志文件，这种一次性在内存中构建完整列表的方式可能会带来巨大的内存压力。

这时，就该 yield 登场了。同样的逻辑，用 yield 可以这样实现：

def iter_even(nums):
    for n in nums:
        if n % 2 == 0:
            yield n

请注意，这个函数的关键区别在于：当你调用 iter_even(range(10)) 时，函数体并不会立即执行完毕。它只是返回一个生成器对象：

g = iter_even(range(10))
print(g)           # <generator object iter_even at 0x...>

for x in g:
    print(x)       # 此时才一边循环，一边遇到 yield 就吐出一个值

可以这样简单理解两者的区别：

• return：一次性交付。函数把所有工作做完，将最终成果（如列表、字典）整个“拍”在桌子上。
• yield：流式生产。函数变成一个“流水线”，调用者需要时（例如在 for 循环中）才生产一个结果，实现了“边算边给”。

函数中的 yield：一个决定性的转变

这里有一个至关重要的细节：只要函数体内出现了 yield 关键字，无论它藏在多深的条件判断或循环里，这个函数就会被 Python 视为生成器函数。调用它时，返回的永远是生成器对象，而不是普通返回值。

看一个容易让人困惑的例子：

def weird(flag):
    if flag:
        return 1
    else:
        yield 2

很多初学者会认为，传入 True 会返回整数 1，传入 False 会返回一个生成器。但实际上，无论传入什么，weird() 返回的永远是生成器对象。当 flag=True 时，如果你尝试 next(weird(True))，它会直接引发 StopIteration 异常，而 return 的值 1 会作为异常的值被携带，而不是常规的返回值。

因此，一个明确的实践原则是：要么将整个函数设计为生成器（使用 yield），要么就完全不要用它。避免设计那种“半生成器、半普通函数”的混合结构。

return 在生成器中的角色：优雅地终止

那么在生成器函数里，还能使用 return 吗？当然可以，而且这是一种常见的模式，用于在满足某个条件时提前终止生成过程。

def upto(n):
    i = 0
    while True:
        if i > n:
            return  # 这里相当于引发一个干净的 StopIteration
        yield i
        i += 1

从外部调用者的视角看：

for x in upto(3):
    print(x)   # 输出：0, 1, 2, 3，然后循环自然结束

这里的 return 扮演了“终止信号”的角色。它也可以携带一个值，但这个值通常不会在普通的 for 循环中被捕获：

def gen_sum(n):
    s = 0
    for i in range(n):
        s += i
        yield i
    return s  # 注意：这个 s 在 for 循环里是拿不到的

如果你想获取生成器函数中 return 的值，需要手动捕获 StopIteration 异常并访问其 value 属性。不过，在大多数业务场景中，我们更关注 yield 产出的序列，很少需要这个终止值。

哪些场景下 yield 是更优解？

在解决特定类型问题时，yield 能极大提升程序的效率和优雅度。一个经典场景是处理大文件：

# 危险写法：一次性读取所有行
def read_lines(path):
    with open(path, 'r', encoding='utf-8') as f:
        return f.readlines()  # 大文件可能导致内存耗尽

# 推荐写法：流式逐行读取
def iter_lines(path):
    with open(path, 'r', encoding='utf-8') as f:
        for line in f:
            yield line.rstrip('\n')

两种写法调用方式类似，但内存占用天差地别：

for line in iter_lines("access.log"):
    # 每次内存中只保留一行数据
    process(line)

类似地，在算法与数据处理中，以下场景都很适合使用 yield：

• 遍历大型数据集寻找特定元素：无需构建中间列表。
• 树或图的遍历（如 DFS、BFS）：天然适合生成访问路径或节点。
• 分页从数据库拉取数据：每次 yield 一批记录，避免一次性加载过载。

何时应坚持使用 return？

选择 return 的理由通常也很明确：

• 结果集很小：比如只有几十或几百个元素，构建列表的内存开销可忽略不计。
• 需要随机访问：后续逻辑需要频繁通过索引（如 result[3]）或切片访问结果，列表等容器更合适。
• 追求简单直观：函数目的就是完成计算并返回一个明确结果，不希望引入生成器的概念来增加理解成本。

在这种情况下，像开头的 find_even 那样直接返回列表，代码可读性更高，对协作者也更友好。

实践中需要避开的“坑”

1. 生成器只能迭代一次：这是最常见的疏忽。

   g = iter_even(range(100))
   list1 = list(g)  # 第一次消耗生成器
   list2 = list(g)  # 这次 list2 是空的！因为 g 已耗尽

   如果需要对结果进行多次处理，要么重新调用函数创建新的生成器，要么先用 list() 转换一次并保存结果。

2. 在递归结构中使用 yield：编写递归生成器时，语法需要特别注意。例如，二叉树的中序遍历，错误的写法是：

   def inorder(root):
       if not root:
           return
       inorder(root.left)  # 这里返回 None，不会 yield 任何值
       yield root.val
       inorder(root.right) # 同上

   正确的写法是使用 yield from，它将委托给子生成器：

   def inorder(root):
       if not root:
           return
       yield from inorder(root.left)
       yield root.val
       yield from inorder(root.right)

   yield from 是处理这类嵌套生成器场景的利器，能让代码简洁清晰。

总结：如何选择？

决策的关键在于明确你的需求是“结果集合”还是“计算过程”。

• 需要结果集合：数据已完全就绪，且后续操作可能需要随机访问或复用，使用 return 返回列表、元组等容器。
• 需要计算过程：数据量大，适合流式、逐个处理，且内存敏感，使用 yield 定义一个生成器。

在Python开发中清晰地理解这两者的差异，能帮助你写出更高效、更符合场景需求的代码。希望这次的梳理能帮你理清思路。如果你在LeetCode刷题或实际项目中遇到类似困惑，不妨试着将一些函数改写成生成器版本，亲身体验一下其中的差异，理解会更加深刻。欢迎在云栈社区分享你的实践心得与更多关于函数设计的见解。