Generators

By Martin McBride, 2018-03-18
Tags: generator iterator
Categories: functional programming


A generator is a convenient way to create your own iterators.

Generator functions

A generator function looks similar to a normal function, except that instead of a return statement it has a yield statement. This simple difference tells Python that it is a generator function.

Here is an example:

def count_down():
    yield(5)
    yield(4)
    yield(3)
    yield(2)
    yield(1)

You can use count_down anywhere you might use an iterator, for example in a for loop (it does a similar job to a range):

for i in count_down():
    print(i)

This will print:

5
4
3
2
1

How a generator works

A generator function above returns an iterator. What the iterator does is controlled by the generator body - each yield statement creates another element for the iterator. When the generator body exits, that is the end of the iteration.

The for loop first calls count_down to get the iterator. Then it calls next on the iterator to get the next value. The iterator starts executing the function, and pretty quickly hits yield(5). At this point, the iterator returns a next value of 5.

The for loop prints 5, then loops round again. It calls next to get the next value from the iterator. The interator remembers where it was and carries on executing from the next line of the function, this time hitting yield(4). The loop prints 4 than loops round again. This cycle continues until the iterator gets to the end of the function. When the function exits, the iterator runs out of values and the for loop ends.

A triangle numbers generator

The triangle numbers are 1, 3, 6, 10... :

1
3 = 1 + 2
6 = 1 + 2 + 3
10 = 1 + 2 + 3 + 4

Here is a generator for triangle numbers:

def triangles():
    i = 1
    n = 1
    while True:
        yield(n)
        i += 1
        n += i

And here is how to call it:

for i in triangles():
    print(i)
    if i > 1000:
        break

This illustrates two important points about generators:

  • You can use a loop in a generator ... of course!
  • You can make an infinite generator.

The triangles generator will keep creating numbers forever. Fortunately, generators are lazy (that is, they generate values on demand), otherwise the code would never work. If it had to create all its values before it returned, it would never finish.

Since we don't want our program to run forever, we put a break statement in our for loop to stop once we hit a triangle number that is bigger than 1000.

Generators can work on other iterables

We can create a generator that works on one or more input iterables. As an example here is a a simple version of map implmented as a generator:

def fake_map(fn, s):
    for x in s:
        yield(fn(x))

it = fake_map(lambda x: x*2, [1, 2, 3, 4])
print(list(it))

fake_map takes a function and a sequence. It applies the function to each element in the sequence and yields it.

The test code uses a lambda function that multiplies the value by 2, and applies it to the list. This gives:

[2, 4, 6, 8]

Another example

In this example we will double the length of the input sequence by duplicating each element. So:

[1, 2, 3, 4] becomes [1, 1, 2, 2, 3, 3, 4, 4]

Here is the code:

def duplicate(s):
    for x in s:
        yield(x)
        yield(x)

See also

If you found this article useful, you might be interested in the book NumPy Recipes or other books by the same author.

Join the PythonInformer Newsletter

Sign up using this form to receive an email when new content is added:

Popular tags

2d arrays abstract data type alignment and angle animation arc array arrays bar chart bar style behavioural pattern bezier curve built-in function callable object chain circle classes clipping close closure cmyk colour combinations comparison operator comprehension context context manager conversion count creational pattern data science data types decorator design pattern device space dictionary drawing duck typing efficiency ellipse else encryption enumerate fill filter font font style for loop formula function function composition function plot functools game development generativepy tutorial generator geometry gif global variable gradient greyscale higher order function hsl html image image processing imagesurface immutable object in operator index inner function input installing iter iterable iterator itertools join l system lambda function latex len lerp line line plot line style linear gradient linspace list list comprehension logical operator lru_cache magic method mandelbrot mandelbrot set map marker style matplotlib monad mutability named parameter numeric python numpy object open operator optimisation optional parameter or pandas partial application path pattern permutations pie chart pil pillow polygon pong positional parameter print product programming paradigms programming techniques pure function python standard library radial gradient range recipes rectangle recursion reduce regular polygon repeat rgb rotation roundrect scaling scatter plot scipy sector segment sequence setup shape singleton slice slicing sound spirograph sprite square str stream string stroke structural pattern subpath symmetric encryption template tex text text metrics tinkerbell fractal transform translation transparency triangle truthy value tuple turtle unpacking user space vectorisation webserver website while loop zip zip_longest