Python Generators Tutorial

What is a generator?

When creating iterators there is a lot of overhead. We have to implement a class with the __iter__() and __next__() methods, keep track of states, raise StopIterations etc.

A generator is just a simple way of creating an iterator. All the overhead is managed by the generator.

How to create a generator

To create a generator we must define a function with at least one yield statement. Both return and yield will return some value from a function.

The difference between yield and return is that a return statement terminates the function, while a yield statement pauses the function.

When the function is paused, it saves all its states and continues from there on any successive calls.

A generator:

• Contains one or more yield statements.
• Returns an iterator
• Does not start execution immediately.
• The iter() and next() methods are automatically implemented.
• When the function yields (pauses), it will transfer control to the caller of the function.
• Local scoped variables (inside the function) and their states, are remembered between function calls.
• The StopIteration is automatically raised when the function terminates.

def gen():
    print("First item")
    yield 1

    print("Second Item")
    yield 2

    print("Third Item")
    yield 3

ex = gen()

# iterate through yields by calling
# the same function the amount of
# times there are yields
next(ex)
next(ex)
next(ex)

In the example above our generator has three yield statements. Each time we call next() on the function it will iterate through the yield statements. We print before each yield statement to demonstrate that the state is saved between calls.

In the example above we can’t see that actual values that were returned with yield. That’s because we only iterated the function, we didn’t print its yield.

def gen():
    print("First item")
    yield 1

    print("Second Item")
    yield 2

    print("Third Item")
    yield 3

ex = gen()

yield_1 = next(ex)
yield_2 = next(ex)
yield_3 = next(ex)

print("\nActual yield values:")
print(yield_1)
print(yield_2)
print(yield_3)

In the example above, we assign each iteration through the yields as a different variable, which allows us to work with the actual values that were yielded.

How to loop through a generator

As we’ve seen in the iterators tutorial lesson , looping through a generator with a for loop will automatically use an iteration and __next__() function, and automatically stop when StopIteration is raised.

def gen():
    print("First item")
    yield 1

    print("Second Item")
    yield 2

    print("Third Item")
    yield 3

ex = gen()

for item in ex:
    print("Actual value yielded:", item)

The example above is of less practical use. Ideally, if there is to be a loop, we want it inside the generator.

def range_up_to(num):
    for item in range(num):
        yield item

seq = range_up_to(5)

print(next(seq))
print(next(seq))
print(next(seq))
print(next(seq))
print(next(seq))

The example above is more practical. We don’t have to specify the loop when we use the generator, we just iterate over its yields.

Let’s consider a generator for a game. We can loop through casting spells as long as the player has mana. If we cast a spell, it pauses all further casting of that spell until its off cooldown.

We can then move on to the next iteration, where the generator remembers how much mana we have left after already having cast the previous spell.

Why use generators

Why should we use generators instead of iterators:

1. Generators are easy to implement. They are clear and concise compared to Iterator classes.
2. Generators are efficient. They will only generate sequence items one at a time in memory, instead of the whole sequence.
3. Generators are potentially infinite streams. Because generators only store one item at a time in memory they can be infinite.
4. Generators can be pipelined. We can easily create a chain of processing elements where the output of each element is the input to the next.

Summary: Points to remember

• Generators create iterators for us and handle all the overhead.
• Generators are functions with at least one yield statement. Yield pauses the function, instead of terminating it.
• Using a for loop to loop through a generator will use an iteration and the __next__() function.
• To stop a potential infinite loop we break out of the loop by raising a StopIteration error.
• Generators are efficient and easy to implement.