How to prevent iterator getting exhausted?

There’s no way to “reset” a generator. However, you can use itertools.tee to “copy” an iterator.

>>> z = zip(a, b)
>>> zip1, zip2 = itertools.tee(z)
>>> list(zip1)
[(1, 7), (2, 8), (3, 9)]
>>> list(zip2)
[(1, 7), (2, 8), (3, 9)]

This involves caching values, so it only makes sense if you’re iterating through both iterables at about the same rate. (In other words, don’t use it the way I have here!)

Another approach is to pass around the generator function, and call it whenever you want to iterate it.

def gen(x):
    for i in range(x):
        yield i ** 2

def make_two_lists(gen):
    return list(gen()), list(gen())

But now you have to bind the arguments to the generator function when you pass it. You can use lambda for that, but a lot of people find lambda ugly. (Not me though! YMMV.)

>>> make_two_lists(lambda: gen(10))
([0, 1, 4, 9, 16, 25, 36, 49, 64, 81], [0, 1, 4, 9, 16, 25, 36, 49, 64, 81])

I hope it goes without saying that under most circumstances, it’s better just to make a list and copy it.

Also, as a more general way of explaining this behavior, consider this. The point of a generator is to produce a series of values, while maintaining some state between iterations. Now, at times, instead of simply iterating over a generator, you might want to do something like this:

z = zip(a, b)
while some_condition():
    fst = next(z, None)
    snd = next(z, None)
    do_some_things(fst, snd)
    if fst is None and snd is None:

Let’s say this loop may or may not exhaust z. Now we have a generator in an indeterminate state! So it’s important, at this point, that the behavior of a generator is restrained in a well-defined way. Although we don’t know where the generator is in its output, we know that a) all subsequent accesses will produce later values in the series, and b) once it’s “empty”, we’ve gotten all the items in the series exactly once. The more ability we have to manipulate the state of z, the harder it is to reason about it, so it’s best that we avoid situations that break those two promises.

Of course, as Joel Cornett points out below, it is possible to write a generator that accepts messages via the send method; and it would be possible to write a generator that could be reset using send. But note that in that case, all we can do is send a message. We can’t directly manipulate the generator’s state, and so all changes to the state of the generator are well-defined (by the generator itself — assuming it was written correctly!). send is really for implementing coroutines, so I wouldn’t use it for this purpose. Everyday generators almost never do anything with values sent to them — I think for the very reasons I give above.

Leave a Comment