One of the joys of Python - Generators

Posted on December 18, 2009

This post will use Fibonacci numbers as a tool to build up to explaining the usefulness of Python generators.

This post will feature both C++ and Python code.

Fibonacci numbers are defined as the sequence: 0, 1, 1, 2, 3, 5, 8, 13, 21, 34, ….

Or in general:

F0 = 0
F1 = 1
Fn = Fn-1 + Fn-2

This can be transferred into a C++ function extremely easily:

size_t Fib(size_t n)
{
    //Fib(0) = 0
    if(n == 0)
        return 0;

    //Fib(1) = 1
    if(n == 1)
        return 1;

    //Fib(N) = Fib(N-2) + Fib(N-1)
    return Fib(n-2) + Fib(n-1);
}

But if you want to print the first 6 Fibonacci numbers, you will be recalculating a lot of the values with the above function.

For example: Fib(3) = Fib(2) + Fib(1), but Fib(2) also recalculates Fib(1). The higher the value you want to calculate, the worse off you will be.

So one may be tempted to re-write the above by keeping track of the state in main.

//Not supported for the first 2 elements of Fib
size_t GetNextFib(size_t &pp, size_t &p)
{
    int result = pp + p;
    pp = p;
    p = result;
    return result;
}

int main(int argc, char *argv[])
{
    size_t pp = 0;
    size_t p = 1;
    std::cout << "0 " << "1 ";
    for(size_t i = 0; i <= 4; ++i)
    {
        size_t fibI = GetNextFib(pp, p);
        std::cout << fibI << " ";
    }
return 0;
}

But this is very ugly, and it complicates our logic in main, it would be better to not have to worry about state in our main function.

We could return a vector of values and use an iterator to iterate over that set of values, but this requires a lot of memory all at once for a large number of return values.

So back to our old approach, what happens if we wanted to do something else besides print the numbers? We’d have to copy and paste the whole block of code in main and change the output statements to whatever else we wanted to do. And if you copy and paste code, then you should be shot. You don’t want to get shot do you?

To solve these problems, and to avoid getting shot, we may re-write this block of code using a callback function. Every time a new Fibonacci number is encountered, we would call the callback function.

void GetFibNumbers(size_t max, void(*FoundNewFibCallback)(size_t))
{
    if(max-- == 0) return;
    FoundNewFibCallback(0);
    if(max-- == 0) return;
    FoundNewFibCallback(1);

    size_t pp = 0;
    size_t p = 1;
    for(;;)
    {
        if(max-- == 0) return;
        int result = pp + p;
        pp = p;
        p = result;
        FoundNewFibCallback(result);
    }
}

void foundNewFib(size_t fibI)
{
    std::cout << fibI << " ";
}

int main(int argc, char *argv[])
{
    GetFibNumbers(6, foundNewFib);
    return 0;
}

This is clearly an improvement, your logic in main is not as cluttered, and you can do anything you want with the Fibonacci numbers, simply define new callbacks.

But this is still not perfect. What if you wanted to only get the first 2 Fibonacci numbers, and then do something, then get some more, then do something else.

Well we could go on like we have been, and we could start adding state again into main, allowing GetFibNumbers to start from an arbitrary point. But this will further bloat our code, and it already looks too big for a simple task like printing Fibonacci numbers.

We could implement a producer and consumer model via a couple threads. But this complicates the code even more.

Instead let’s talk about generators.

Python has a very nice language feature that solves problems like these called generators.

A generator allows you to execute a function, stop at an arbitrary point, and then continue again where you left off. Each time returning a value.

Consider the following code that uses a generator:

def fib():
    pp, p = 0, 1  
    while 1:
        yield pp
        pp, p = p, pp+p

g = fib()
for i in range(6):
    g.next()

Which gives us the results:

0
1
1
2
3
5

The yield statement is used in conjuction with Python generators. It saves the state of the function and returns the yeilded value. The next time you call the next() function on the generator, it will continue where the yield left off.

This is by far more clean than the callback function code. We have cleaner code, smaller code, and not to mention much more functional code (Python allows arbitrarily large integers).

c++ math python
Next page