Python Generator：一個被低估的性能利器

作者：piperliu 2024-12-13 08:02:10

本文討論這背后的 Python Generator 。下次當你需要處理大量數據或實現流式處理時，不要忘了考慮使用 Generator。它可能會讓你的代碼更優雅，性能更好。

調用 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)

這段代碼有什么問題？它會一次性將所有符合條件的行都讀入內存。如果日志文件有 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)

看起來很相似，但運行機制完全不同。這個版本無論日志文件多大，內存占用都很小。

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"，返回 2

在流式 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)

這樣的設計有幾個好處：

節省內存：不需要等待全部內容生成完畢
實時響應：用戶可以立即看到部分結果
可中斷：如果用戶不需要更多結果，可以隨時停止

高級用法：Generator 表達式和的雙向通信 send 的魔法

Generator 還有一些高級特性。比如 Generator 表達式：

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

# Generator 表達式
squares = (x*x for x in range(1000000))  # 按需生成

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.0

每次 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("!"))

這個類型定義展現了 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 的子類當然也可以

簡單理解：

協變（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)

這個例子展示了 Generator 在實際應用中的優雅之處：

每個函數職責單一，易于測試和維護
數據流處理清晰，內存占用小
類型提示清晰，代碼更容易理解

下次當你需要處理大量數據或實現流式處理時，不要忘了考慮使用 Generator。它可能會讓你的代碼更優雅，性能更好。

責任編輯：姜華來源： Piper蛋窩