<< Go back to Posts

Python Generators

A Cheat Sheet.


Introduction

A generator is a function that when called, runs from its last position to the next yield x, outputs x, and then stop until it is called again.

A generator can have a finite number of execution, for instance when using a for loop or when reading a file, or an infinite number when using a while loop.

Generators are convenient as data on which they iterate do not need to be stored. For instance:

  • for i in [1, 2, 3, 4, 5, 6]
  • for i in range(1, 7)

In the first case, the elements from 1 to 6 are stored in a list, so 6 elements are stored in the memory.

In the second case, the elements are not stored. We call the range generator, that outputs at each round the next value until reaching 6

The syntax for quick and easy generator is very similar to list generation:

  • [i for i in range(4, 20)] generate a list
  • (i for i in range(4, 20)) generate a generator.

You use it already

Compare:

dic = {"a": 32, "b": 34, "c": 50}

for c in list(dic):
    if c == "b":
        del(dic[c])

And

dic = {"a": 32, "b": 34, "c": 50}

for c in dic.keys(): # equivalent to "for c in dic:"
    if c == "b":
        del(dic[c])

The first example runs correctly, and you end up with {"a": 32, "c": 50}. The second throws an error.

The first example works because you have dic and you ask first to evaluate your generator using list(dic) (which is equivalent to [k for k in dic]). Because when the for loop runs, the dic keys are already been extracted, there is no problem when discarding one element of the dic.

In the second case, this is a generator. A generator may or may not accept to handle modification during its execution. Here, it does not.

Examples

Simple (finite) generator

def my_generator(x):
    print("This is called only once")
    for i in range(x, 12):
        yield i

gen = my_generator(3)

print(next(gen))

>>> This is called only once
>>> 3

next(gen)

>>> 4

for i in gen:
    print(i)

>>> 5
>>> 6
>>> 7
>>> 8
>>> 9
>>> 10
>>> 11

When the generator is out of elements:

next(gen)

---------------------------------------------------------------------------

StopIteration                             Traceback (most recent call last)

Cell In[6], line 1
----> 1 next(gen)


StopIteration:

A Simple (infinite) Generator

Works as finite generators. The only difference is that you never reach the end.

def my_generator_infinite(x):
    while True:
        x += 1
        yield x


gen = my_generator_infinite(2)
for i in gen:
    print(i)

>>> 3 # +1 applied BEFORE the `yield x`
>>> 4
>>> 5
...

Special Operations

There are three operations that can be done on the generator:

  • gen.send(): update a value within the generator
  • gen.close(): stop the generator, useful for infinite generators
  • gen.throw(): throw an error, which can be handled or not by the generator

Send

def my_generator_1(x):
    while True:
        i = yield x
        print("My update from the world", i)
        if i is not None:
            x = i

        x += 1

gen = my_generator_1(5)

gen.send(6)

---------------------------------------------------------------------------

TypeError                                 Traceback (most recent call last)

Cell In[9], line 1
----> 1 gen.send(6)


TypeError: can't send non-None value to a just-started generator

next(gen)

>>> 5


next(gen)

>>>  My update from the world None
>>> 6


gen.send(10) # Send used here !

>>> My update from the world 10
>>> 11

Move from somewhere to \(10\) and apply the +1

What happened ?

The send(10) replaces i by 10 AND call the generator for one loop:

- `x = i`
- `x += 1 # 11`
- New round of the loop
- `i = yield x`: the generator stop here and return `x` which was `11`

Next, you can continue using the generator:

next(gen)

>>>    My update from the world None
>>>    12

Here, i in the print() is None, as there were no send() message, so i was not updated.

When calling next(gen), it is similar to gen.send(None).

Close

Stop the generator.

gen.close()

gen

>>> <generator object my_generator_1 at 0x7f4eec0e0f20>

next(gen)

---------------------------------------------------------------------------

StopIteration                             Traceback (most recent call last)

Cell In[16], line 1
----> 1 next(gen)


StopIteration:

Usefulness of stop ? I do not know. To stop a loop without using a break ?

def my_generator_infinite(x):
    while True:
        x += 1
        yield x

Usage without close():

gen = my_generator_infinite(10)
for i in gen:
    print(i)
    if i > 15:
        break

print("=== break ===")

for i in gen:
    print(i)
    if i > 20:
        break

>>>    11
>>>    12
>>>    13
>>>    14
>>>    15
>>>    16
>>>    === break ===
>>>    17
>>>    18
>>>    19
>>>    20
>>>    21

Usage with close()

gen = my_generator_infinite(10)
for i in gen:
    print(i)
    if i > 15:
        gen.close()

print("=== break ===")

# We do not enter the loop, because the generator does not return anything.
for i in gen:
    print(i)
    if i > 20:
        break

>>>    11
>>>    12
>>>    13
>>>    14
>>>    15
>>>    16
>>>    === break ===

Throw

For errors …

Without error handling

gen = my_generator_infinite(10)

for i in gen:
    print(i)
    if i > 15:
        gen.throw(ValueError("Please stop the loop"))

>>>    11
>>>    12
>>>    13
>>>    14
>>>    15
>>>    16

---------------------------------------------------------------------------

ValueError                                Traceback (most recent call last)

Cell In[20], line 6
      4 print(i)
      5 if i > 15:
----> 6     gen.throw(ValueError("Please stop the loop"))


Cell In[17], line 4, in my_generator_infinite(x)
      2 while True:
      3     x += 1
----> 4     yield x


ValueError: Please stop the loop

This has the effect of stopping the generator

next(gen)

---------------------------------------------------------------------------

StopIteration                             Traceback (most recent call last)

Cell In[21], line 1
----> 1 next(gen)


StopIteration:

With error handling

def my_generator_error_handling(x):
    while True:
        x += 1
        try:
            yield x
        except ValueError as e:
            print("Error message:", e)
            pass

gen = my_generator_error_handling(10)

for i in gen:
    print(i)
    if i > 15:
        gen.throw(ValueError("Please stop the loop"))

    if i > 20:
        # To stop
        gen.close()

>>>    11
>>>    12
>>>    13
>>>    14
>>>    15
>>>    16
>>>    Error message: Please stop the loop
>>>    18
>>>    Error message: Please stop the loop
>>>    20
>>>    Error message: Please stop the loop
>>>    22
>>>    Error message: Please stop the loop

Useful Ressources

Check Real Python - Introduction to python generators for an in-depth introduction


>> You can subscribe to my mailing list here for a monthly update. <<