深入理解Python生成器中StopIteration异常的捕获机制

def test(vid, size):
    while True:
        try:
            # part 是一个生成器表达式
            part = (next(vid) for _ in range(size))
            yield part
        except StopIteration:
            # 期望在此捕获，但实际上不会发生
            break

res = test((i for i in range(100)), 30)
for i in res:
    for j in i: # 异常在此处发生
        print(j, end=" ")
    print()

---------------------------------------------------------------------------
StopIteration                             Traceback (most recent call last)
Cell In[54], line 4, in (.0)
      3 try:
----> 4     part = (next(vid) for _ in range(size))
      5     yield part

StopIteration: 

The above exception was the direct cause of the following exception:

RuntimeError                              Traceback (most recent call last)
Cell In[54], line 11
      9 res = test((i for i in range(100)), 30)
     10 for i in res:
---> 11     for j in i:
     12         print(j, end=" ")
     13         print()

RuntimeError: generator raised StopIteration

def test2():
    try:
        def foo():
            raise StopIteration
        return foo # foo函数在此处并未被调用
    except StopIteration: # 此处不会捕获到异常
        pass

outer_foo = test2()
outer_foo() # <--- StopIteration 在此处被引发

# 这种形式下，StopIteration可以在内部被捕获
for _ in range(size):
    yield next(vid) # <-- StopIteration可以在这里被捕获

def create_batches(source_generator, batch_size):
    """
    将一个源生成器分割成多个子生成器，每个子生成器产生指定大小的批次。
    当源生成器耗尽时，优雅地终止。

    Args:
        source_generator: 原始的生成器或可迭代对象。
        batch_size: 每个批次（子生成器）的元素数量。

    Yields:
        一个子生成器，每次迭代产生一个批次的元素。
    """
    done = False # 标志，指示源生成器是否已完全耗尽

    def batch_generator_inner():
        """
        内部生成器，负责从源生成器中获取单个批次的元素。
        它会在内部捕获StopIteration，并更新外部的done标志。
        """
        nonlocal done # 声明使用外部作用域的done变量
        # print("--- new batch ---") # 调试信息
        for i in range(batch_size):
            # print(f"batch {i+1} / {batch_size}") # 调试信息
            try:
                yield next(source_generator)
            except StopIteration:
                # 捕获到StopIteration，表示源生成器已耗尽
                # print("StopIteration caught, and we are done") # 调试信息
                done = True # 设置标志，通知外部循环停止
                break # 退出当前批次的生成

    # 只要源生成器未完全耗尽，就不断生成新的批次生成器
    while not done:
        yield batch_generator_inner()

# 示例用法
print("--- 示例1：源生成器有余数 ---")
source_data = (i for i in range(10)) # 0到9共10个元素
batch_size = 3
batches = create_batches(source_data, batch_size)

for batch_idx, batch in enumerate(batches):
    print(f"\n处理批次 {batch_idx + 1}:")
    for elem in batch:
        print(f"  元素: {elem}")

print("\n--- 示例2：源生成器刚好整除 ---")
source_data_exact = (i for i in range(9)) # 0到8共9个元素
batch_size_exact = 3
batches_exact = create_batches(source_data_exact, batch_size_exact)

for batch_idx, batch in enumerate(batches_exact):
    print(f"\n处理批次 {batch_idx + 1}:")
    for elem in batch:
        print(f"  元素: {elem}")

--- 示例1：源生成器有余数 ---

处理批次 1:
  元素: 0
  元素: 1
  元素: 2

处理批次 2:
  元素: 3
  元素: 4
  元素: 5

处理批次 3:
  元素: 6
  元素: 7
  元素: 8

处理批次 4:
  元素: 9

--- 示例2：源生成器刚好整除 ---

处理批次 1:
  元素: 0
  元素: 1
  元素: 2

处理批次 2:
  元素: 3
  元素: 4
  元素: 5

处理批次 3:
  元素: 6
  元素: 7
  元素: 8