Python Generator：一个被低估的性能利器-51CTO.COM

调用 OpenAI 的 API 时，设置 stream=True ，接着 for chunk in completion: 我们就可以“流式”地获取响应的内容。而非等待远程的模型将所有内容生成完毕，再返回给我们（这通常要等很久）。

本文讨论这背后的 Python Generator 。

从一个经典问题开始

假设我们要处理一个超大的日志文件，需要按行读取并分析。传统的做法是：

def read_log_file(filename):
    result = []
    with open(filename) as f:
        for line in f:
            if "ERROR" in line:
                result.append(line)
    return result

# 使用方式
errors = read_log_file("huge.log")
for error in errors:
    process_error(error)1.
2.
3.
4.
5.
6.
7.
8.
9.
10.
11.
12.

这段代码有什么问题？它会一次性将所有符合条件的行都读入内存。如果日志文件有 10GB，而符合条件的行有 5GB，那么我们的程序就需要 5GB 的内存。

Generator 版本

我们用 Generator 改写一下：

def read_log_file(filename):
    with open(filename) as f:
        for line in f:
            if "ERROR" in line:
                yield line

# 使用方式
for error in read_log_file("huge.log"):
    process_error(error)1.
2.
3.
4.
5.
6.
7.
8.
9.

看起来很相似，但运行机制完全不同。这个版本无论日志文件多大，内存占用都很小。

Generator 的工作原理

Generator 的核心特点是"懒加载"（lazy evaluation）。当我们调用一个生成器函数时，它并不会立即执行函数体，而是返回一个生成器对象。只有在实际请求下一个值时，它才会执行到下一个 yield 语句。

来看一个更直观的例子：

def counter():
    print("Starting")
    i = 0
    while True:
        print(f"Generating {i}")
        yield i
        i += 1

# 创建生成器对象
c = counter()  # 此时不会打印任何内容
print("Generator created")

# 获取前三个值
print(next(c))  # 打印 "Starting" 和 "Generating 0"，返回 0
print(next(c))  # 打印 "Generating 1"，返回 1
print(next(c))  # 打印 "Generating 2"，返回 21.
2.
3.
4.
5.
6.
7.
8.
9.
10.
11.
12.
13.
14.
15.
16.

在流式 API 中的应用

现在我们理解了为什么流式 API 会使用 Generator。以 OpenAI 的流式响应为例：

def stream_completion(prompt):
    # 模拟 API 调用
    response = ["生成", "AI", "回复", "需要", "时间"]
    for token in response:
        yield token

# 使用方式
for chunk in stream_completion("你好"):
    print(chunk, end="", flush=True)1.
2.
3.
4.
5.
6.
7.
8.
9.

这样的设计有几个好处：

节省内存：不需要等待全部内容生成完毕
实时响应：用户可以立即看到部分结果
可中断：如果用户不需要更多结果，可以随时停止

高级用法：Generator 表达式和的双向通信 send 的魔法

Generator 还有一些高级特性。比如 Generator 表达式：

# 列表推导式
squares = [x*x for x in range(1000000)]  # 立即生成所有结果

# Generator 表达式
squares = (x*x for x in range(1000000))  # 按需生成1.
2.
3.
4.
5.

GGenerator 不只是能产出值，还能接收值！这就是 send 方法的精妙之处。让我们通过一个计算移动平均值的例子，一步步看看 send 是如何工作的：

def averager():
    total = 0
    count = 0
    average = None
    while True:
        # yield 在这里扮演双重角色：
        # 1. 向外返回 average 值
        # 2. 接收外部发送的 value
        value = yield average
        if value is None:
            break
        total += value
        count += 1
        average = total / count

# 让我们一步步看发生了什么
avg = averager()          # 创建生成器对象，但函数体还未开始执行
print("第一步：创建生成器")

next(avg)                 # 启动生成器，运行到第一个 yield，返回 None
print("第二步：生成器已启动，等待第一个值")

print(avg.send(10))      # 1. send(10) 将 10 传给 value
                         # 2. 计算 average = 10/1 = 10.0
                         # 3. 到达 yield，返回 10.0
                         # 4. 生成器暂停，等待下一个值

print(avg.send(20))      # 1. value 获得值 20
                         # 2. 计算 average = 30/2 = 15.0
                         # 3. yield 返回 15.0

print(avg.send(30))      # 1. value 获得值 30
                         # 2. 计算 average = 60/3 = 20.0
                         # 3. yield 返回 20.01.
2.
3.
4.
5.
6.
7.
8.
9.
10.
11.
12.
13.
14.
15.
16.
17.
18.
19.
20.
21.
22.
23.
24.
25.
26.
27.
28.
29.
30.
31.
32.
33.
34.

每次 send 调用，生成器都会在 value = yield average 这行代码处经历一个完整的"接收-计算-返回"周期。这种优雅的设计让生成器不仅能产出数据，还能根据外部输入动态调整其行为。

类型系统中的 Generator：优雅的泛型设计

在 Python 的类型系统中，Generator 的类型定义也别具匠心。它使用了三个泛型参数：

from typing import Generator, Iterator
from typing import TypeVar, Generic

T_co = TypeVar('T_co', covariant=True)      # 生成器产出的类型
T_contra = TypeVar('T_contra', contravariant=True)  # send 方法接收的类型
V_co = TypeVar('V_co', covariant=True)      # return 语句返回的类型

def number_processor() -> Generator[int, str, float]:
    # 这个生成器：
    # - yield 产出 int
    # - 接收 str 类型的输入
    # - 最终 return float 类型的值
    count = 0
    total = 0.0
    while count < 3:
        text = yield count  # 产出 int，接收 str
        if text:
            total += len(text)
        count += 1
    return total  # 返回 float

# 使用示例
proc = number_processor()
next(proc)  # 启动生成器，返回 0
print(proc.send("hello"))  # 输出 1
print(proc.send("world"))  # 输出 2
try:
    proc.send("!")  # 生成器将结束，抛出 StopIteration
except StopIteration as e:
    print(f"最终结果：{e.value}")  # 打印 11.0 (len("hello") + len("world") + len("!"))1.
2.
3.
4.
5.
6.
7.
8.
9.
10.
11.
12.
13.
14.
15.
16.
17.
18.
19.
20.
21.
22.
23.
24.
25.
26.
27.
28.
29.
30.

这个类型定义展现了 Python 类型系统中一些概念：

T_co 是协变的（covariant），表示生成器产出的类型可以是基类
T_contra 是逆变的（contravariant），表示接收的类型可以是子类
V_co 也是协变的，表示返回值类型可以是基类

可以通过具体的例子来解释协变和逆变：

from typing import Generator, TypeVar

# 定义一些类来演示
class Animal:
    pass

class Dog(Animal):
    pass

class Chihuahua(Dog):
    pass

# 定义类型变量
T_co = TypeVar('T_co', covariant=True)
T_contra = TypeVar('T_contra', contravariant=True)

def dog_generator() -> Generator[Dog, Animal, None]:
    # 这个生成器:
    # - 产出 Dog (协变位置)
    # - 接收 Animal (逆变位置)
    dog = yield Dog()  # dog 的类型是 Animal
    
# 协变（T_co）示例：
# 如果一个函数返回 Dog，它也可以用在需要返回 Animal 的地方
generator1: Generator[Animal, Animal, None] = dog_generator()  # 没问题！
# 因为 Dog 是 Animal 的子类，所以可以用 Dog 替代 Animal

# 逆变（T_contra）示例：
# 如果一个函数接收 Animal，它也可以接收 Dog 或 Chihuahua
generator2 = dog_generator()
generator2.send(Chihuahua())  # 没问题！
# 因为函数期望接收 Animal，那接收 Animal 的子类当然也可以1.
2.
3.
4.
5.
6.
7.
8.
9.
10.
11.
12.
13.
14.
15.
16.
17.
18.
19.
20.
21.
22.
23.
24.
25.
26.
27.
28.
29.
30.
31.
32.

简单理解：

协变（covariant）：允许使用更具体的类型

如果方法返回 Dog，可以用在需要 Animal 的地方。
因为 Dog 一定是 Animal，所以这样是安全的。

逆变（contravariant）：允许使用更宽泛的类型

如果方法接收 Animal，可以传入 Dog 或 Chihuahua
因为方法能处理所有 Animal，当然也能处理具体的 Dog

Generator[T_co, T_contra, V_co] 中：

T_co：产出值的类型（协变），因为生成器提供值
T_contra：send 方法接收的类型（逆变），因为生成器接收值
V_co：return 语句的返回值类型（协变），因为是提供值

这种设计让 Generator 类型在静态类型检查时既保持了类型安全，又提供了足够的灵活性。

实战应用：构建流式处理管道

让我们把学到的知识组合起来，构建一个优雅的流式处理管道：

from typing import Generator, Iterator
from itertools import chain

def read_chunks(file_path: str) -> Generator[str, None, None]:
    with open(file_path) as f:
        while chunk := f.read(1024):
            yield chunk

def process_chunk(chunk: str) -> Generator[str, None, None]:
    # 处理单个 chunk 中的行
    # 注意：chunk 可能在行中间截断，需要处理这种情况
    lines = chunk.split('\n')
    for line in lines:
        if line.strip():
            yield line.upper()

def filter_keywords(lines: Iterator[str]) -> Generator[str, None, None]:
    keywords = {'ERROR', 'WARNING', 'CRITICAL'}
    for line in lines:
        if any(k in line for k in keywords):
            yield line

# 正确的流式处理版本
def process_log_file(file_path: str):
    # 当前行的未完成部分
    partial_line = ''
    
    for chunk in read_chunks(file_path):
        # 处理可能被截断的行
        if partial_line:
            chunk = partial_line + chunk
            partial_line = ''
            
        # 分割成行，保留最后一个可能不完整的行
        lines = chunk.split('\n')
        if not chunk.endswith('\n'):
            partial_line = lines[-1]
            lines = lines[:-1]
            
        # 处理完整的行
        for line in lines:
            if line.strip():
                # 直接在这里处理，无需存储所有行
                upper_line = line.upper()
                if any(k in upper_line for k in {'ERROR', 'WARNING', 'CRITICAL'}):
                    print(upper_line)
    
    # 处理最后一个不完整的行（如果有的话）
    if partial_line and partial_line.strip():
        upper_line = partial_line.upper()
        if any(k in upper_line for k in {'ERROR', 'WARNING', 'CRITICAL'}):
            print(upper_line)

# 或者，使用更函数式的写法
def process_log_file_functional(file_path: str):
    def handle_chunks() -> Generator[str, None, None]:
        partial_line = ''
        for chunk in read_chunks(file_path):
            if partial_line:
                chunk = partial_line + chunk
                partial_line = ''
            
            lines = chunk.split('\n')
            if not chunk.endswith('\n'):
                partial_line = lines[-1]
                lines = lines[:-1]
            
            yield from (line for line in lines if line.strip())
        
        if partial_line and partial_line.strip():
            yield partial_line

    # 现在我们真正实现了流式处理
    lines = handle_chunks()
    upper_lines = (line.upper() for line in lines)
    filtered_lines = filter_keywords(upper_lines)
    
    for line in filtered_lines:
        print(line)1.
2.
3.
4.
5.
6.
7.
8.
9.
10.
11.
12.
13.
14.
15.
16.
17.
18.
19.
20.
21.
22.
23.
24.
25.
26.
27.
28.
29.
30.
31.
32.
33.
34.
35.
36.
37.
38.
39.
40.
41.
42.
43.
44.
45.
46.
47.
48.
49.
50.
51.
52.
53.
54.
55.
56.
57.
58.
59.
60.
61.
62.
63.
64.
65.
66.
67.
68.
69.
70.
71.
72.
73.
74.
75.
76.
77.
78.
79.

这个例子展示了 Generator 在实际应用中的优雅之处：

每个函数职责单一，易于测试和维护
数据流处理清晰，内存占用小
类型提示清晰，代码更容易理解

下次当你需要处理大量数据或实现流式处理时，不要忘了考虑使用 Generator。它可能会让你的代码更优雅，性能更好。