Search code examples
pythongeneratoryield-keyword

What Happens when a Generator Runs out of Values to Yield?


To illustrate the question, suppose we have this simple generator:

def firstn(n):
    num = 0
    while num < n:
        yield num
        num += 1

for i in firstn(10):
    print i

This will print the digits 0 through 9. But what if we have:

def firstn(n):
    num = 0
    while num < 5 < n:
        yield num
        num += 1

for i in firstn(10):
    print i

(The change is in the while statement.) Then it prints only digits 0 through 4. Once num >= 5, then the generator no longer yields values.

What I'm curious about is what goes on under the hood: I used PythonTutor to step through the code, and the impression I'm under is that once the while statement is no longer True, the function implicitly returns None, which the for loop somehow detects, and then also breaks. I used the next built-in to inspect this more closely:

>>> def firstn(n):
...     num = 0
...     while num < 5 < n:
...         yield num
...         num += 1
... 
>>> 
>>> mygen = firstn(100)
>>> next(mygen)
0
>>> next(mygen)
1
>>> next(mygen)
2
>>> next(mygen)
3
>>> next(mygen)
4
>>> next(mygen)
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
StopIteration

Which supports my theory. My big question: how does StopIteration work, and does this mean that calling a generator with a large value can be equivalent to calling it with its smallest terminating value? In our example, for i in firstn(5) and for i in firstn(9999999999999) should be equivalent, right?


Solution

  • This isn't very mysterious. When a generator runs out of values to yield, it raises a StopIteration exception. You just need to understand how a for-loop works in Python. Essentially, it is equivalent to the following code:

    iterator = iter(collection)
    while True:
        try:
            x = next(iterator)
            # do something
        except StopIteration as e:
            break
    

    The above is equivalent to:

    for x in collection:
        # do something