Python Generator：一個(gè)被低估的性能利器
作者：piperliu
本文討論這背后的 Python Generator 。下次當(dāng)你需要處理大量數(shù)據(jù)或?qū)崿F(xiàn)流式處理時(shí)，不要忘了考慮使用 Generator。它可能會(huì)讓你的代碼更優(yōu)雅，性能更好。
調(diào)用 OpenAI 的 API 時(shí)，設(shè)置 stream=True ，接著 for chunk in completion: 我們就可以“流式”地獲取響應(yīng)的內(nèi)容。而非等待遠(yuǎn)程的模型將所有內(nèi)容生成完畢，再返回給我們（這通常要等很久）。
本文討論這背后的 Python Generator 。
從一個(gè)經(jīng)典問題開始
假設(shè)我們要處理一個(gè)超大的日志文件，需要按行讀取并分析。傳統(tǒng)的做法是：
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)這段代碼有什么問題？它會(huì)一次性將所有符合條件的行都讀入內(nèi)存。如果日志文件有 10GB，而符合條件的行有 5GB，那么我們的程序就需要 5GB 的內(nèi)存。
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)看起來很相似，但運(yùn)行機(jī)制完全不同。這個(gè)版本無論日志文件多大，內(nèi)存占用都很小。
Generator 的工作原理
Generator 的核心特點(diǎn)是"懶加載"（lazy evaluation）。當(dāng)我們調(diào)用一個(gè)生成器函數(shù)時(shí)，它并不會(huì)立即執(zhí)行函數(shù)體，而是返回一個(gè)生成器對象。只有在實(shí)際請求下一個(gè)值時(shí)，它才會(huì)執(zhí)行到下一個(gè) yield 語句。
來看一個(gè)更直觀的例子：
def counter():
print("Starting")
i = 0
while True:
print(f"Generating {i}")
yield i
i += 1
# 創(chuàng)建生成器對象
c = counter() # 此時(shí)不會(huì)打印任何內(nèi)容
print("Generator created")
# 獲取前三個(gè)值
print(next(c)) # 打印 "Starting" 和 "Generating 0"，返回 0
print(next(c)) # 打印 "Generating 1"，返回 1
print(next(c)) # 打印 "Generating 2"，返回 2在流式 API 中的應(yīng)用
現(xiàn)在我們理解了為什么流式 API 會(huì)使用 Generator。以 OpenAI 的流式響應(yīng)為例：
def stream_completion(prompt):
# 模擬 API 調(diào)用
response = ["生成", "AI", "回復(fù)", "需要", "時(shí)間"]
for token in response:
yield token
# 使用方式
for chunk in stream_completion("你好"):
print(chunk, end="", flush=True)這樣的設(shè)計(jì)有幾個(gè)好處：
- 節(jié)省內(nèi)存：不需要等待全部內(nèi)容生成完畢
- 實(shí)時(shí)響應(yīng)：用戶可以立即看到部分結(jié)果
- 可中斷：如果用戶不需要更多結(jié)果，可以隨時(shí)停止
高級(jí)用法：Generator 表達(dá)式和的雙向通信 send 的魔法
Generator 還有一些高級(jí)特性。比如 Generator 表達(dá)式：
# 列表推導(dǎo)式
squares = [x*x for x in range(1000000)] # 立即生成所有結(jié)果
# Generator 表達(dá)式
squares = (x*x for x in range(1000000)) # 按需生成GGenerator 不只是能產(chǎn)出值，還能接收值！這就是 send 方法的精妙之處。讓我們通過一個(gè)計(jì)算移動(dòng)平均值的例子，一步步看看 send 是如何工作的：
def averager():
total = 0
count = 0
average = None
while True:
# yield 在這里扮演雙重角色：
# 1. 向外返回 average 值
# 2. 接收外部發(fā)送的 value
value = yield average
if value is None:
break
total += value
count += 1
average = total / count
# 讓我們一步步看發(fā)生了什么
avg = averager() # 創(chuàng)建生成器對象，但函數(shù)體還未開始執(zhí)行
print("第一步：創(chuàng)建生成器")
next(avg) # 啟動(dòng)生成器，運(yùn)行到第一個(gè) yield，返回 None
print("第二步：生成器已啟動(dòng)，等待第一個(gè)值")
print(avg.send(10)) # 1. send(10) 將 10 傳給 value
# 2. 計(jì)算 average = 10/1 = 10.0
# 3. 到達(dá) yield，返回 10.0
# 4. 生成器暫停，等待下一個(gè)值
print(avg.send(20)) # 1. value 獲得值 20
# 2. 計(jì)算 average = 30/2 = 15.0
# 3. yield 返回 15.0
print(avg.send(30)) # 1. value 獲得值 30
# 2. 計(jì)算 average = 60/3 = 20.0
# 3. yield 返回 20.0每次 send 調(diào)用，生成器都會(huì)在 value = yield average 這行代碼處經(jīng)歷一個(gè)完整的"接收-計(jì)算-返回"周期。這種優(yōu)雅的設(shè)計(jì)讓生成器不僅能產(chǎn)出數(shù)據(jù)，還能根據(jù)外部輸入動(dòng)態(tài)調(diào)整其行為。
類型系統(tǒng)中的 Generator：優(yōu)雅的泛型設(shè)計(jì)
在 Python 的類型系統(tǒng)中，Generator 的類型定義也別具匠心。它使用了三個(gè)泛型參數(shù)：
from typing import Generator, Iterator
from typing import TypeVar, Generic
T_co = TypeVar('T_co', covariant=True) # 生成器產(chǎn)出的類型
T_contra = TypeVar('T_contra', contravariant=True) # send 方法接收的類型
V_co = TypeVar('V_co', covariant=True) # return 語句返回的類型
def number_processor() -> Generator[int, str, float]:
# 這個(gè)生成器：
# - yield 產(chǎn)出 int
# - 接收 str 類型的輸入
# - 最終 return float 類型的值
count = 0
total = 0.0
while count < 3:
text = yield count # 產(chǎn)出 int，接收 str
if text:
total += len(text)
count += 1
return total # 返回 float
# 使用示例
proc = number_processor()
next(proc) # 啟動(dòng)生成器，返回 0
print(proc.send("hello")) # 輸出 1
print(proc.send("world")) # 輸出 2
try:
proc.send("!") # 生成器將結(jié)束，拋出 StopIteration
except StopIteration as e:
print(f"最終結(jié)果：{e.value}") # 打印 11.0 (len("hello") + len("world") + len("!"))這個(gè)類型定義展現(xiàn)了 Python 類型系統(tǒng)中一些概念：
- T_co 是協(xié)變的（covariant），表示生成器產(chǎn)出的類型可以是基類
- T_contra 是逆變的（contravariant），表示接收的類型可以是子類
- V_co 也是協(xié)變的，表示返回值類型可以是基類
可以通過具體的例子來解釋協(xié)變和逆變：
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]:
# 這個(gè)生成器:
# - 產(chǎn)出 Dog (協(xié)變位置)
# - 接收 Animal (逆變位置)
dog = yield Dog() # dog 的類型是 Animal
# 協(xié)變（T_co）示例：
# 如果一個(gè)函數(shù)返回 Dog，它也可以用在需要返回 Animal 的地方
generator1: Generator[Animal, Animal, None] = dog_generator() # 沒問題！
# 因?yàn)?Dog 是 Animal 的子類，所以可以用 Dog 替代 Animal
# 逆變（T_contra）示例：
# 如果一個(gè)函數(shù)接收 Animal，它也可以接收 Dog 或 Chihuahua
generator2 = dog_generator()
generator2.send(Chihuahua()) # 沒問題！
# 因?yàn)楹瘮?shù)期望接收 Animal，那接收 Animal 的子類當(dāng)然也可以簡單理解：
- 協(xié)變（covariant）：允許使用更具體的類型
- 如果方法返回 Dog，可以用在需要 Animal 的地方。
- 因?yàn)?Dog 一定是 Animal，所以這樣是安全的。
- 逆變（contravariant）：允許使用更寬泛的類型
- 如果方法接收 Animal，可以傳入 Dog 或 Chihuahua
- 因?yàn)榉椒芴幚硭?Animal，當(dāng)然也能處理具體的 Dog
- Generator[T_co, T_contra, V_co] 中：
- T_co：產(chǎn)出值的類型（協(xié)變），因?yàn)樯善魈峁┲?/li>
- T_contra：send 方法接收的類型（逆變），因?yàn)樯善鹘邮罩?/li>
- V_co：return 語句的返回值類型（協(xié)變），因?yàn)槭翘峁┲?/li>
這種設(shè)計(jì)讓 Generator 類型在靜態(tài)類型檢查時(shí)既保持了類型安全，又提供了足夠的靈活性。
實(shí)戰(zhàn)應(yīng)用：構(gòu)建流式處理管道
讓我們把學(xué)到的知識(shí)組合起來，構(gòu)建一個(gè)優(yōu)雅的流式處理管道：
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]:
# 處理單個(gè) chunk 中的行
# 注意：chunk 可能在行中間截?cái)?，需要處理這種情況
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):
# 當(dāng)前行的未完成部分
partial_line = ''
for chunk in read_chunks(file_path):
# 處理可能被截?cái)嗟男? if partial_line:
chunk = partial_line + chunk
partial_line = ''
# 分割成行，保留最后一個(gè)可能不完整的行
lines = chunk.split('\n')
if not chunk.endswith('\n'):
partial_line = lines[-1]
lines = lines[:-1]
# 處理完整的行
for line in lines:
if line.strip():
# 直接在這里處理，無需存儲(chǔ)所有行
upper_line = line.upper()
if any(k in upper_line for k in {'ERROR', 'WARNING', 'CRITICAL'}):
print(upper_line)
# 處理最后一個(gè)不完整的行（如果有的話）
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)
# 或者，使用更函數(shù)式的寫法
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
# 現(xiàn)在我們真正實(shí)現(xiàn)了流式處理
lines = handle_chunks()
upper_lines = (line.upper() for line in lines)
filtered_lines = filter_keywords(upper_lines)
for line in filtered_lines:
print(line)這個(gè)例子展示了 Generator 在實(shí)際應(yīng)用中的優(yōu)雅之處：
- 每個(gè)函數(shù)職責(zé)單一，易于測試和維護(hù)
- 數(shù)據(jù)流處理清晰，內(nèi)存占用小
- 類型提示清晰，代碼更容易理解
下次當(dāng)你需要處理大量數(shù)據(jù)或?qū)崿F(xiàn)流式處理時(shí)，不要忘了考慮使用 Generator。它可能會(huì)讓你的代碼更優(yōu)雅，性能更好。
責(zé)任編輯：姜華
來源：
Piper蛋窩
