如何测试Python生成器首个元素且不消耗后续迭代？

我有一个生成器gen，具备以下特性：

• 触发其yield操作的开销远高于创建生成器
• 生成的元素占用大量内存
• 有时所有__next__调用都会抛出异常，但创建生成器时无法预知该情况

该生成器并非由我实现。我需要一种方法，让生成器产出首个元素（用try/except捕获异常），且后续遍历生成器时不会直接从第二个元素开始？

我的尝试代码 #

try:
    first = next(gen)
except StopIterator:
    return None
except Exception:
    print("Generator throws exception on a yield")

# looping also over the first element which we yielded already
for thing in (first, *gen):
    do_something_complicated(thing)

几种不够理想的方案 #

• 创建生成器，测试首个元素，再创建新生成器并遍历第二个生成器。（缺点：重复触发yield操作，额外开销大）
• 将整个for循环放入try/except中；但该方案不佳，因为yield抛出的异常过于通用，可能会捕获其他无关异常。
• 产出首个元素并测试，再将首个元素与生成器剩余部分重新组合为新生成器（理想情况下无需将生成器所有元素提取到列表中，避免占用大量内存）。

我给出的示例代码接近方案3，但该代码会在迭代开始前将gen的所有元素提取到元组中，这正是我想要避免的。

最优解决方案 #

你可以用两种方式实现需求，都不会提前加载生成器的所有元素到内存，而是按需迭代：

方法1：使用itertools.chain #

itertools.chain可以将单个元素（包装成可迭代对象）和原生成器组合，遍历的时候逐个产出元素，不会提前消耗整个生成器：

from itertools import chain

try:
    first = next(gen)
except StopIterator:
    # 生成器为空，直接返回
    return None
except Exception as e:
    print(f"Generator throws exception on first yield: {e}")
    # 根据业务需求处理异常，比如返回或终止流程
    return

# 组合首个元素与剩余生成器，按需迭代
for thing in chain([first], gen):
    do_something_complicated(thing)

方法2：自定义生成器函数 #

如果不想依赖itertools，可以自己写一个简单的生成器函数来拼接首个元素和原生成器：

def prepend_item(first_item, generator):
    yield first_item
    yield from generator

try:
    first = next(gen)
except StopIterator:
    return None
except Exception as e:
    print(f"Generator throws exception on first yield: {e}")
    return

for thing in prepend_item(first, gen):
    do_something_complicated(thing)

这两种方案都满足你的核心需求：

1. 仅触发一次生成器的yield流程，避免额外开销
2. 不会提前把生成器所有元素加载到内存，节省内存占用
3. 可以正常遍历包括首个元素在内的所有生成器产出项

内容的提问来源于stack exchange，提问作者Marses