如何在不消耗生成器的情况下测试它

Question

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

• 让它产生收益是相当昂贵的（比创建生成器更昂贵）
• 元素占用相当多的内存
• 有时所有调用__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)

我看到的解决方案不太好：

1. 创建生成器，测试第一个元素，创建一个新生成器，循环第二个元素。
2. 将整个 for 循环放入 try/ except 中；不太好，因为yield抛出的异常非常普遍，它可能会捕获其他东西。
3. 生成第一个元素，对其进行测试，然后从第一个元素和其余元素重组一个新的生成器gen（理想情况下不要将所有gen's 元素提取到列表中，因为这可能会占用大量内存）。

对于 3，这似乎是最好的解决方案，一个近乎完美的例子就是我上面给出的例子，但我相信它只会在gen我们开始迭代之前将 的所有元素提取到一个元组中，这是我想避免的。

Answer 1

我想我有你正在寻找的使用 more_itertools 库的东西：

import more_itertools

if __name__ == "__main__":
    generator = range(100)

    peekable_generator = more_itertools.peekable(generator)

    print(f"peek {peekable_generator.peek()}")
    print(f"next {next(peekable_generator)}")
    print(f"next {next(peekable_generator)}")

output: 
    peek 0
    next 0
    next 1

请参阅此处的文档：https ://more-itertools.readthedocs.io/en/stable/api.html#more_itertools.peekable

如果我没记错的话，查看第一个项目的能力是您需要的关键。