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.