Infinite Python Generator

RRR week is ending and Finals week at Berkeley is about to start!

Currently, I’ve been studying for the 3 courses that I’m currently taking (CS61A, EE16A, and Econ 2). (As a side note, UGBA10 doesn’t have a final exam so that’s nice.)

Recently, my friend Chandana who also takes the CS61A course with me recently asked me about this question about generators in an optional worksheet for our class:

Write a generator function textbfgen inf that returns a generator which yields all the numbers in the provided list one by one in an infinite loop. 

>>> t = gen_inf([3, 4, 5])
>>> next(t)
3
>>> next(t)
4
>>> next(t)
5
>>> next(t)
3
>>> next(t)
4
def gen_inf(lst):

The solution is:

def gen_inf(lst):
	while True:
		for elem in lst:
			yield elem

Chandana asked me why this solution worked since there wasn’t an explanation in the solutions.

At first glance, I thought that the solution code would simply raise a StopIteration exception at the end of the list. I was very confused why it continued to go on infinitely. How does the code know to go back to the first element after it reaches the end?

To first understand what is the reason why this solution code successfully creates an infinite generator, I removed the while loop and tested it to see if that was the reason why this is able to work infinitely:

>>> def gen_inf(lst):
...     for elem in lst:
...             yield elem
... 
>>> t = gen_inf([3, 4, 5])
>>> next(t)
3
>>> next(t)
4
>>> next(t)
5
>>> next(t)
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
StopIteration

As expected, the StopIteration exception is thrown at the end of the list. That means the while True part of the solutoin code is the reason behind why the generator is able to continue infinitely.

From this new information, I knew that the reason why the solution code worked was because of the “while True” line. I just didn’t understand how the “while True” line makes the generator work infinitely. So, I inserted a print statement right below the “while True” line:

def gen_inf(lst): 
	while True:
		print(lst)
		for elem in lst: 
			yield elem

When I tested it, I got the following:

>>> def gen_inf(lst):
...     while True:
...             print(lst)
...             for elem in lst:
...                     yield elem
... 
>>> t = gen_inf([3, 4, 5])
>>> next(t)
[3, 4, 5]
3
>>> next(t)
4
>>> next(t)
5
>>> next(t)
[3, 4, 5]
3
>>> next(t)
4
>>> next(t)
5
>>> next(t)
[3, 4, 5]
3
>>> next(t)
4
>>> next(t)
5

Interesting! The list gets printed right before every time the first element of the list is yielded, but doesn’t get printed for the other yield statements. So, every time we are at the first element of the list that is passed in, the list passed in is printed.

In our example of gen_inf([3, 4, 5]), the first iteration of the while True loop yields 3, then yields 4, and then yields 5.

The second iteration of the while True loop i yields 3, then yields 4, and then yields 5. This continues on infinitely.

Just to emphasize how the code works, here’s another way of thinking how the code works with our example of gen_inf([3, 4, 5]):

When it yields 3, it’s still in the first iteration of the while loop. When it yields 4, it’s still in the first iteration of the while loop. When it yields 5, it’s still in the first iteration of the while loop.

Then it’s done with the for loop. So, it starts with the next iteration of the while loop.

So, for the next next() statement, we are in the second iteration of the while statement. and we start a new for iteration with the same list

Each iteration in the while loop ends when the for loop inside of it ends.

So, the reason why the solution code never reaches the StopIteration exception is because the inner for loop ends right when the last element is yielded (since the for… in loop ends at the last element of a list)! Right after, the next iteration of the while loop starts and the whole cycle restarts where the each element of the list is yielded.