Using a list comprehension, we can directly create a list. However, due to memory constraints, the size of a list is inevitably limited. Moreover, creating a list with a million elements not only occupies a great deal of storage space, but if we only need to access the first few elements, the space occupied by the vast majority of subsequent elements is wasted.
So, if the list elements can be calculated according to a certain algorithm, can we continuously calculate the subsequent elements during the loop? This way, there’s no need to create a complete list, thus saving a lot of space. In Python, this mechanism of calculating while looping is called a generator.
There are several ways to create a generator. The first method is simple: just change the square brackets [] of a list comprehension to parentheses (), and you’ve created a generator:
>>> L = [x * x for x in range(10)]
>>> L
[0, 1, 4, 9, 16, 25, 36, 49, 64, 81]
>>> g = (x * x for x in range(10))
>>> g
<generator object <genexpr> at 0x1022ef630>
The only difference between creating L and g is the outermost [] and (). L is a list, while g is a generator.
We can directly print each element of a list, but how do we print each element of a generator?
If you want to print them one by one, you can use the next() function to get the next return value of the generator:
>>> next(g)
0
>>> next(g)
1
>>> next(g)
4
>>> next(g)
9
>>> next(g)
16
>>> next(g)
25
>>> next(g)
36
>>> next(g)
49
>>> next(g)
64
>>> next(g)
81
>>> next(g)
Traceback (most recent call last):
File "<stdin>", line 1, in <module>
StopIteration
As we mentioned, a generator saves an algorithm. Each call to next(g) calculates the value of the next element of g until the last element is reached. When there are no more elements, a StopIteration error is raised.
Of course, continuously calling next(g) like the above is quite cumbersome. The correct approach is to use a for loop, because a generator is also an iterable object:
>>> g = (x * x for x in range(10))
>>> for n in g:
... print(n)
...
0
1
4
9
16
25
36
49
64
81
Therefore, after creating a generator, we almost never call next() directly. Instead, we iterate over it using a for loop and don’t need to worry about the StopIteration error.
Generators are very powerful. If the calculation algorithm is complex and cannot be implemented using a for loop similar to a list comprehension, it can also be implemented using a function.
For example, the famous Fibonacci sequence: except for the first and second numbers, any number can be obtained by adding the previous two numbers:
1, 1, 2, 3, 5, 8, 13, 21, 34, …
The Fibonacci sequence cannot be written using a list comprehension, but it’s easy to print it using a function:
def fib(max):
n, a, b = 0, 0, 1
while n < max:
print(b)
a, b = b, a + b
n = n + 1
return 'done'
Note the assignment statement:
a, b = b, a + b
This is equivalent to:
t = (b, a + b) # t is a tuple
a = t[0]
b = t[1]
But there’s no need to explicitly write the temporary variable t for the assignment.
The above function can output the first N numbers of the Fibonacci sequence:
>>> fib(6)
1
1
2
3
5
8
'done'
Upon closer inspection, you can see that the fib function actually defines the calculation rule for the Fibonacci sequence. It can calculate any subsequent element starting from the first element. This logic is actually very similar to a generator.
That is to say, the above function is only one step away from being a generator. To turn the fib function into a generator function, you just need to change print(b) to yield b:
def fib(max):
n, a, b = 0, 0, 1
while n < max:
yield b
a, b = b, a + b
n = n + 1
return 'done'
This is another way to define a generator. If a function definition contains the yield keyword, then this function is no longer an ordinary function but a generator function. Calling a generator function returns a generator:
>>> f = fib(6)
>>> f
<generator object fib at 0x104feaaa0>
Here, the most difficult thing to understand is that the execution flow of a generator function is different from that of an ordinary function. An ordinary function executes sequentially and returns when it encounters a return statement or the last line of the function. A function that has become a generator executes each time next() is called, returns when it encounters a yield statement, and resumes execution from the yield statement where it last returned the next time it is executed.
Let’s take a simple example: define a generator function that returns the numbers 1, 3, and 5 in sequence:
def odd():
print('step 1')
yield 1
print('step 2')
yield(3)
print('step 3')
yield(5)
When calling this generator function, you first need to create a generator object, and then continuously obtain the next return value using the next() function:
>>> o = odd()
>>> next(o)
step 1
1
>>> next(o)
step 2
3
>>> next(o)
step 3
5
>>> next(o)
Traceback (most recent call last):
File "<stdin>", line 1, in <module>
StopIteration
As you can see, odd is not an ordinary function but a generator function. During execution, it pauses when it encounters yield and resumes the next time. After executing yield three times, there are no more yield statements to execute, so the fourth call to next(o) results in an error.
Note
Calling a generator function creates a generator object. Calling the generator function multiple times creates multiple independent generators.
Some beginners might find that calling next() like this returns 1 every time:
>>> next(odd())
step 1
1
>>> next(odd())
step 1
1
>>> next(odd())
step 1
1
The reason is that odd() creates a new generator object each time. The above code actually creates three completely independent generators. Calling next() on each of the three generators will of course return the first value each time.
The correct way is to create a single generator object and then continuously call next() on that one object:
>>> g = odd()
>>> next(g)
step 1
1
>>> next(g)
step 2
3
>>> next(g)
step 3
5
Returning to the fib example, we continuously call yield during the loop, which causes continuous pauses. Of course, we need to set a condition for the loop to exit; otherwise, an infinite sequence will be generated.
Similarly, after changing a function into a generator function, we almost never use next() to get the next return value but directly iterate over it using a for loop:
>>> for n in fib(6):
... print(n)
...
1
1
2
3
5
8
However, when calling a generator using a for loop, you’ll notice that you cannot obtain the return value of the generator’s return statement. If you want to get the return value, you must catch the StopIteration error; the return value is contained in the value attribute of StopIteration:
>>> g = fib(6)
>>> while True:
... try:
... x = next(g)
... print('g:', x)
... except StopIteration as e:
... print('Generator return value:', e.value)
... break
...
g: 1
g: 1
g: 2
g: 3
g: 5
g: 8
Generator return value: done
We will explain how to catch errors in detail later in the error handling section.
Generators are very powerful tools. In Python, you can simply convert a list comprehension into a generator, or implement generators with complex logic through functions.
To understand how generators work, they continuously calculate the next element during a for loop and terminate the loop under appropriate conditions. For generators converted from functions, encountering a return statement or executing the last line of the function body serves as the instruction to terminate the generator, and the for loop ends accordingly.
Please note the distinction between ordinary functions and generator functions. Calling an ordinary function directly returns the result:
>>> r = abs(6)
>>> r
6
Calling a generator function actually returns a generator object:
>>> g = fib(6)
>>> g
<generator object fib at 0x1022ef948>