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

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.

Answer 2

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