python - 与 itertools.dropwhile 相反(如何在 N 次迭代后停止生成器)

Question

是否有一种简单的方法可以在 N 循环后停止迭代器？当然我可以写这样的东西:

for i, val in enumerate(gen()):
    if i > N: break

但我想写一些类似的东西

for val in stop_after(gen(), N):
    ...

我尝试使用 itertools.dropwhile 但它似乎做了相反的事情。当然，我可以用逆逻辑重写 itertools.dropwhile ，但我想知道是否已经实现了一些东西。

Answer 1

使用islice:

for val in itertools.islice(gen(), N):
    ....

假设您的示例是:

for i, val in enumerate(gen()):
    if i > N: break