Generators

14_generators

Generators

Generators are custom iterators.

Generators are one way to make a function or expression that can be iterated in a loop or with the next() function.

Generator Functions

In [0]:
def counter_gen():
    count = 0
    while True:
        count += 1
        yield count
In [0]:
for count in counter_gen():
    print(count)
    if count >= 10:
        break

1
2
3
4
5
6
7
8
9
10

Generator functions will not remember where they left off from one call to the next.

In [0]:
for count in counter_gen():
    print(count)
    if count >= 10:
        break

1
2
3
4
5
6
7
8
9
10

But there is a way to make them remember! Memoization. This ability is not unique to generator functions, you can memoize any function.

In [0]:
def smart_counter_gen():
    if not hasattr(smart_counter_gen, "count"):
        smart_counter_gen.count = 0
    while True:
        smart_counter_gen.count += 1
        yield smart_counter_gen.count
In [0]:
for count in smart_counter_gen():
    print(count)
    if count >= 10:
        break

1
2
3
4
5
6
7
8
9
10
In [0]:
for count in smart_counter_gen():
    print(count)
    if count >= 20:
        break

11
12
13
14
15
16
17
18
19
20

What if we wanted a way to reset the counter?

In [0]:
def smart_counter(reset=False):
    if reset or not hasattr(smart_counter, "count"):
        smart_counter.count = 0
    while True:
        smart_counter.count += 1
        yield smart_counter.count
In [0]:
for count in smart_counter():
    print(count)
    if count >= 10:
        break

1
2
3
4
5
6
7
8
9
10
In [0]:
for count in smart_counter(reset=True):
    print(count)
    if count >= 10:
        break

1
2
3
4
5
6
7
8
9
10
In [0]:
for count in smart_counter():
    print(count)
    if count >= 20:
        break

11
12
13
14
15
16
17
18
19
20

Generator Expressions

Generator expressions look a lot like tuple comprehensions, but there is no such thing as a tuple comprehension in Python. Generator expressions can be a little tricky at first because they can only be evaluated once and this can be a subtle problem as there will be no error – it just wont do what you want after the first evaluation.

In [0]:
square_gen = (i*i for i in range(1, 11))

for square in square_gen:
    print(square)

1
4
9
16
25
36
49
64
81
100
In [0]:
for square in square_gen:
    print(square)  # Prints Nothing! The generator is empty.
In [0]:

Generator Expression that works like a Generator Function

To make a generator expression or any iterator that can be used more than once, we must wrap it in a callable, like a lambda. Every time we call this lambda a new generator is created for us. This gives the added ability to parameterize the generator expression.

In [0]:
square_gen = lambda n: (i*i for i in range(n+1))
In [0]:
for square in square_gen(9):
    print(square)

0
1
4
9
16
25
36
49
64
81
In [0]:
for square in square_gen(5):
    print(square)

0
1
4
9
16
25