Python Generators

Generators

• List comprehensions can create lists directly, however, list capacity is limited by memory
• Creating a list with a million elements takes up a lot of storage space
• If you only need to access the first few elements, the space used by the majority of the elements will be wasted
• Instead of creating complete lists, elements can be calculated during the loop, which saves a lot of space
• This mechanism is called a generator in Python

Creating Generators

• Many ways to create generators
• Change the square brackets, [], to parentheses, (), in a list comprehension
• The difference between creating L and g lies only in the outermost [] and ()
• L is a list
• g is a generator
• Print each element of a list directly
• Next() function obtains 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

• The generator stores an algorithm: each call to next(g) calculates the next value of g until the last element is calculated
• StopIteration error raised if there are no more elements
• Use a for loop
• Generators are also iterable objects

>>> g = (x * x for x in range(10))
>>> for n in g:
... print(n)
...
0
1
4
9
16
25
36
49
64
81

• Generators are never called using next()
• instead they are iterated over by for loops so there is no need to care about the StopIteration error
• Generators are very powerful
• Use of functions to implement more complex algorithms that list comprehensions cannot achieve

Fibonacci Sequence Example

• In a Fibonacci sequence, each number (except the first and second) can be obtained by adding the two preceding numbers

1, 1, 2, 3, 5, 8, 13, 21, 34

• Fibonacci sequence cannot be written using list comprehension, but it is very easy to print it out 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'

• Assignment statement example: a, b = b, a + b
• Equivalent to:

t = (b, a + b) #t is a tuple
a = t [0]
b = t [1]

• No need to explicitly write out the temporary variable t to assign a value

>>> fib(6)
1
1
2
3
5
'done'

• The fib function defines the calculation rules of the Fibonacci sequence
• To convert the fib function into a generator, replace print(b) with 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'

• A function that contains the yield keyword is no longer an ordinary function, but a generator:

>>> f = fib(6)
>>> f
<generator object fib at 0x104feaaa0>

• Most difficult to understand is that the execution flow of a generator is different from that of a function:
• The function is executed sequentially and returns when it encounters a return statement or the last line of the function statement
• A function that turns into a generator is executed each time next() is called
• When it encounters a yield statement, it returns and continues execution from the last returned yield statement
• Can define a generator to return 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 a generator, first create a generator object, and then use the next() function to continuously obtain the next return value

>>> 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

• odd is not an ordinary function, but a generator
• The execution is interrupted upon encountering yield
• The program then resumes during the next call
• The error is thrown after three yield calls and a fourth call of next()
• In the Fibonacci example, the loop is continuously interrupted when yield is called
• Conditions have to be set to exit the loop, otherwise, an infinite sequence is made
• Using for loops to iterate, is more appropriate than directly using next() to obtain the next return value from generator functions

>>> for n in fib(6):
... print(n)
...
1
1
2
3
5
8

• The return statement's return value cannot be obtained when calling a generator using a 'for' loop
• To obtain the return value, it is necessary to capture the StopIteration error and the return value contained in the StopIteration's value

>>> 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

Additional Information

• Error capture will be explained in detail in the error handling section later on.

Summary

• Generators are powerful tools in Python that can easily turn list comprehensions into generators or implement complex logic through functions
• When a generator is created from a function, the return statement or the last line of the function body is the instruction to end the generator and the for loop also ends
• Normal function calls return results directly
• Function calls return a function object
• To differentiate between ordinary functions and generator functions, ordinary function calls return results directly:

>>> r = abs(6)
>>> r
6

• generator function calls return a generator object:

>>> g = fib(6)
>>> g
<generator object fib at 0x1022ef948>