Python生成器实现分批输出列表的策略与实践

import itertools

def compute_add_generator_batch_problematic(batch_size):
    data = range(5)
    cases = list(itertools.permutations(data, 2)) # 生成所有排列组合

    res = []
    for x, y in cases:
        ans = x + y

        if len(res) != batch_size: # 如果当前批次未满
            res.append(ans)
            continue # 继续填充批次

        # 当批次已满时
        yield res # 输出当前批次
        res = [] # 重置批次列表

# 测试代码
batch_size = 3
print(f"cases={cases}") # 打印所有组合
print("--- 错误的批处理输出 ---")
for res_batch in compute_add_generator_batch_problematic(batch_size):
    print(f"res={res_batch}")

import itertools

def compute_add_generator_batch(batch_size):
    """
    一个生成器函数，用于分批返回计算结果。

    Args:
        batch_size (int): 每个批次中包含的元素数量。
                          必须大于0。
    Yields:
        list: 包含 batch_size 个元素（或在最后批次中少于 batch_size 个）的列表。
    """
    assert batch_size > 0, "batch_size 必须大于0"

    data = range(5)
    # 避免在生成器内部将itertools.permutations的结果转换为list，
    # 这样可以保持惰性求值，但在本示例中数据量小，影响不大。
    # 对于大数据集，直接迭代itertools.permutations对象更优。
    cases_iterator = itertools.permutations(data, 2) 

    batch = [] # 初始化一个空列表来存储当前批次的元素
    for x, y in cases_iterator:
        ans = x + y
        batch.append(ans) # 将计算结果添加到当前批次

        if len(batch) == batch_size: # 如果当前批次已满
            yield batch # 输出完整批次
            batch = [] # 重置批次列表，准备下一个批次

    # 循环结束后，处理可能存在的最后一个不满 batch_size 的批次
    if batch: # 如果 batch 列表非空，说明有剩余元素
        yield batch # 输出剩余元素

# 示例调用与验证
batch_size = 3
report = []
print(f"--- 正确的批处理输出 (batch_size={batch_size}) ---")
for res_batch in compute_add_generator_batch(batch_size):
    report.append(res_batch)
    print(f"res={res_batch}")

print(f"\n收集到的所有批次: {report}")

--- 正确的批处理输出 (batch_size=3) ---
res=[1, 2, 3]
res=[4, 1, 3]
res=[4, 5, 2]
res=[3, 5, 6]
res=[3, 4, 5]
res=[7, 4, 5]
res=[6, 7]

收集到的所有批次: [[1, 2, 3], [4, 1, 3], [4, 5, 2], [3, 5, 6], [3, 4, 5], [7, 4, 5], [6, 7]]