Using next() on generator function

Question

I have this generator function:-

>>> def gen():
...     for i in range(3):
...             yield i*i

Now when i am calling next() on gen(); it is giving the first element each time.

>>> next(gen())
0
>>> next(gen())
0

But when i use that in a for loop; it works as expected:-

>>> for i in gen():
...     print(i)
... 
0
1
4

Can someone explain the cause of this effect; and the concept that i am missing here??

Answer 1

Each time you call the function it returns a generator object. And each time you call the next on it you'll get the first item which is 0 * 0 because you're not calling the next on same object (each time a new one). But in second case you're looping over one generator object and it will continue to consuming the generator until it hit the StopIteration .

For a better demonstration you can create two iterator objects from your generator function and loop over them simultaneously:

In [17]: g = gen()

In [18]: k = gen()

In [19]: for i, j in zip(g, k):
    ...:     print(i, j)
    ...:     
0 0
1 1
4 4

Answer 2

Your function returns a new generator each time you call gen() .

I think the simplest way to understand is to contrast what you are doing:

>>> next(gen())
0
>>> next(gen())
0

with this:

>>> my_gen = gen()
>>> next(my_gen)
0
>>> next(my_gen)
1
>>> next(my_gen)
4
>>> next(my_gen)Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
StopIteration

In this latter case, I am getting new values from the same generator.