Python生成器批量输出：高效处理数据的实现与常见陷阱

import itertools

# 传统函数，一次性返回所有结果
def compute_add_sync():
    data = range(5)
    cases = list(itertools.permutations(data, 2))
    print(f"{cases=}")
    result = []
    for x, y in cases:
        ans = x + y
        result.append(ans)
    return result

report_sync = compute_add_sync()
print(f"{report_sync=}")

# 转换为生成器，每次生成一个值
def compute_add_generator_single():
    data = range(5)
    cases = list(itertools.permutations(data, 2))
    print(f"{cases=}")
    for x, y in cases:
        ans = x + y
        yield ans

report_gen_single = []
for res in compute_add_generator_single():
    report_gen_single.append(res)
print(f"{report_gen_single=}")

# 错误的批量生成器实现
def compute_add_generator_batch_flawed(batch_size):
    data = range(5)
    cases = list(itertools.permutations(data, 2))
    print(f"{cases=}")

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

        if len(res_batch) != batch_size: # 当批次未满时，添加元素
            res_batch.append(ans)
            continue # 继续循环，不执行下面的yield

        # 当批次已满时，yield批次，然后重置批次
        yield res_batch
        res_batch = [] # 重置批次列表

    # 错误：循环结束后，如果res_batch中还有剩余元素，它们将被遗漏
    # 并且如果批次大小刚好等于总元素数量的倍数，也可能遗漏最后的空批次检查

print("\n--- 错误批量生成器输出 ---")
batch_size_flawed = 3
for res in compute_add_generator_batch_flawed(batch_size_flawed):
    print(f"{res=}")

import itertools

def compute_add_generator_batch_correct(batch_size):
    # 确保批次大小有效
    assert batch_size > 0, "batch_size 必须大于 0"

    data = range(5)
    # 这里的 itertools.permutations 也可以直接作为生成器使用，避免一次性生成所有cases
    # 但为了与原始问题保持一致，这里先生成列表
    all_cases = list(itertools.permutations(data, 2)) 

    current_batch = []
    for x, y in all_cases:
        ans = x + y
        current_batch.append(ans) # 始终将元素添加到当前批次

        if len(current_batch) == batch_size: # 当批次达到指定大小
            yield current_batch # 产出完整批次
            current_batch = [] # 重置批次列表，准备下一个批次

    # 循环结束后，处理可能存在的不足一个批次的剩余元素
    if current_batch: # 如果 current_batch 不为空
        yield current_batch # 产出剩余批次

print("\n--- 正确批量生成器输出 ---")
report_batches = []
batch_size_correct = 3
for res_batch in compute_add_generator_batch_correct(batch_size_correct):
    report_batches.append(res_batch)
    print(f"{res_batch=}")

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

res_batch=[1, 2, 3]
res_batch=[4, 1, 3]
res_batch=[4, 5, 2]
res_batch=[3, 5, 6]
res_batch=[3, 4, 5]
res_batch=[7, 4, 5]
res_batch=[6, 7]

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