Python生成器中嵌套非生成器函数实现yield的优雅方案

你需要消除生成器中重复的"添加元素-检查批量-产出-重置"逻辑，同时保留生成器的惰性特性（不提前存储所有结果），以下是几种优雅的实现方式：

方案一：封装逻辑为辅助生成器，用yield from产出结果 #

把重复的逻辑封装成内部辅助函数，该函数返回一个生成器（要么为空，要么包含当前批量），通过yield from批量产出结果。既消除了代码重复，又完全保留生成器的惰性：

from typing import Generator

def f() -> Generator[list[int], None, None]:
    result = []
    
    def add_and_check(value: int) -> Generator[list[int], None, None]:
        nonlocal result
        result.append(value)
        if len(result) == 2:
            yield result
            result = []
    
    # 逐个添加元素，通过yield from产出符合条件的批量
    yield from add_and_check(1)
    yield from add_and_check(2)
    yield from add_and_check(3)
    yield from add_and_check(4)

print(list(f()))  # 输出: [[1, 2], [3, 4]]

方案二：使用协程（Send方法）处理动态添加 #

如果你的实际场景中元素是动态产生（而非预先知道所有元素），可以用协程实现批量收集器，通过send()方法添加元素，同时产出批量结果：

from typing import Generator

def batch_collector(batch_size: int) -> Generator[list[int] | None, int, None]:
    batch = []
    while True:
        # 接收外部send的元素
        value = yield
        batch.append(value)
        if len(batch) == batch_size:
            # 产出批量结果，之后回到等待接收状态
            yield batch
            batch = []

def f() -> Generator[list[int], None, None]:
    collector = batch_collector(2)
    next(collector)  # 启动协程到第一个接收点
    
    # 逐个发送元素，并处理产出的批量
    for val in [1, 2, 3, 4]:
        collector.send(val)
        # 尝试获取批量结果（若存在）
        batch = next(collector)
        if batch:
            yield batch

print(list(f()))  # 输出: [[1, 2], [3, 4]]

方案三：面向对象封装批量逻辑 #

如果需要更灵活的状态管理（比如支持中途修改批量大小、处理剩余元素），可以用类封装批量生成逻辑，通过add()方法返回迭代器，再用yield from产出：

from typing import Generator, Iterator

class BatchGenerator:
    def __init__(self, batch_size: int):
        self.batch_size = batch_size
        self.current_batch = []
    
    def add(self, value: int) -> Iterator[list[int]]:
        self.current_batch.append(value)
        if len(self.current_batch) == self.batch_size:
            yield self.current_batch
            self.current_batch = []
    
    def flush(self) -> Iterator[list[int]]:
        # 处理最后剩余的不足批量的元素（若需要）
        if self.current_batch:
            yield self.current_batch

def f() -> Generator[list[int], None, None]:
    bg = BatchGenerator(2)
    yield from bg.add(1)
    yield from bg.add(2)
    yield from bg.add(3)
    yield from bg.add(4)
    # 若有剩余元素，可添加 yield from bg.flush()

print(list(f()))  # 输出: [[1, 2], [3, 4]]

方案四：利用迭代器工具简化（适用于已知元素序列） #

如果你的元素是预先存在的序列，可以直接用itertools结合生成器表达式实现分组，这是最简洁的方式：

from typing import Generator
import itertools

def f() -> Generator[list[int], None, None]:
    items = [1, 2, 3, 4]
    it = iter(items)
    # 每次从迭代器取2个元素组成批量
    yield from (list(itertools.islice(it, 2)) for _ in itertools.repeat(None))

print(list(f()))  # 输出: [[1, 2], [3, 4]]

以上所有方案都遵循了生成器的惰性原则，不会提前将所有结果存储在内存中，完美贴合你的需求。

内容的提问来源于stack exchange，提问作者Vedran Šego