Get the nth item of a generator in Python

Question

Is there a more syntactically concise way of writing the following?

gen = (i for i in xrange(10))
index = 5
for i, v in enumerate(gen):
    if i is index:
           


        

Answer 1

I'd argue against the temptation to treat generators like lists. The simple but naive approach is the simple one-liner:

gen = (i for i in range(10))
list(gen)[3]

But remember, generators aren't like lists. They don't store their intermediate results anywhere, so you can't go backwards. I'll demonstrate the problem with a simple example in the python repl:

>>> gen = (i for i in range(10))
>>> list(gen)[3]
3
>>> list(gen)[3]
Traceback (most recent call last):
  File "", line 1, in 
IndexError: list index out of range

Once you start going through a generator to get the nth value in the sequence, the generator is now in a different state, and attempting to get the nth value again will return you a different result, which is likely to result in a bug in your code.

Let's take a look at another example, based on the code from the question.

One would initially expect the following to print 4 twice.

gen = (i for i in range(10))
index = 4
for i, v in enumerate(gen):
    if i == index:
        answer = v
        break
print(answer)
for i, v in enumerate(gen):
    if i == index:
        answer = v
        break
print(answer)

but type this into the repl and you get:

>>> gen = (i for i in range(10))
>>> index = 4
>>> for i, v in enumerate(gen):
...     if i == index:
...             answer = v
...             break
... 
>>> print(answer)
4
>>> for i, v in enumerate(gen):
...     if i == index:
...             answer = v
...             break
... 
>>> print(answer)
9

Good luck tracing that bug down.

EDIT:

As pointed out, if the generator is infinitely long, you can't even convert it to a list. The expression list(gen) will never finish.

There is a way you could put a lazily evaluated caching wrapper around an infinite generator to make it look like an infinitely long list you could index into at will, but that deserves its own question and answer, and would have major performance implications.