Iterators and Generators in Python

If you have written some code in Python, something more than the simple “Hello World” program, you have probably used iterable objects. Iterable objects are objects that conform to the “Iteration Protocol” and can hence be used in a loop.

For example:

for i in range(50):
    print(i)

In this example, the range(50) is an iterable object that provides, at each iteration, a different value that is assigned to the “i” variable. Quite easy, but what if we would like to create an iterable object ourselves?

The iteration protocol

Creating an iterable object in Python is as easy as implementing the iteration protocol. Let’s pretend that we want to create an object that would let us iterate over the “Fibonacci sequence”. The Fibonacci sequence is a sequence of integer numbers characterized by the fact that every number after the first two is the sum of the two preceding ones. So the sequence starts with 0 and 1 and then each number that follows is just the sum of the two previous numbers in the sequence. So the third number is 1 (0+1), the fourth is 2 (1+1), the fifth is 3 (1+2), the sixth is 5 (2+3) and so on.

Enough said, let’s the code:

class fibonacci:

    def __init__(self, max=1000000):
        self.a, self.b = 0, 1
        self.max = max

    def __iter__(self):
        # Return the iterable object (self)
        return self

    def next(self):
        # When we need to stop the iteration we just need to raise
        # a StopIteration exception
        if self.a > self.max:
            raise StopIteration

        # save the value that has to be returned
        value_to_be_returned = self.a

        # calculate the next values of the sequence
        self.a, self.b = self.b, self.a + self.b

        return value_to_be_returned

    def __next__(self):
        # For compatibility with Python3
        return self.next()


if __name__ == '__main__':
    MY_FIBONACCI_NUMBERS = fibonacci()
    for fibonacci_number in MY_FIBONACCI_NUMBERS:
        print(fibonacci_number)

As you can see, all we’ve done has been creating a class that implements the iteration protocol. This protocol consists in two methods: the “iter” method that returns the object we would iterate over and the “next” method that is called automatically on each iteration and that returns the value for the current iteration.

Please note that the protocol in Python 2 is a little different and the “__next__” method is called just “next” so it is quite common to use the old Python 2 style method to generate the value and then create the Python 3 style method to simply return the value generated by the former one, so as to have code that can works both with Python 2 and Python 3.

Generators

Generators in Python are just another way of creating iterable objects and are usually used when you need to create iterable object quickly, without the need of creating a class and adopting the iteration protocol. To create a generator you just need to define a function and then use the yield keyword instead of return.

So, the Fibonacci sequence in a generator could be something like this:

def fibonacci(max):
    a, b = 0, 1
    while a < max:
        yield a
        a, b = b, a+b

Yes, so simple! Now, if you want to test it just use your new Fibonacci generator function:

if __name__ == '__main__':
    # Create a generator of fibonacci numbers smaller than 1 million
    fibonacci_generator = fibonacci(1000000)

    # print out all the sequence
    for fibonacci_number in fibonacci_generator:
        print(fibonacci_number)

0
1
1
2
3
5
8
13
21
34
55
89
144
233
377
610
987
1597
2584
4181
6765
10946
17711
28657
46368
75025
121393
196418
317811
514229
832040

Please note that once we have “consumed” the generator, we can’t use it anymore because generators in Python can’t be rewound.

So, if after the code above we tried to print out all the sequence again, we won’t get any values.

    # since the sequence is over, we will not get any value here
    for fibonacci_number in fibonacci_generator:
        print(fibonacci_number)

So, if you need to use the generator again, you have to call the generator function again

    # So, if you need to use the generator again... recreate it!
    fibonacci_generator = fibonacci(1000000)

    # Ok, let's list 'em again
    for fibonacci_number in fibonacci_generator:
        print(fibonacci_number)

Now, if you can, take some time to debug the generator code above and look at how the values are generated and returned. You will find out that the values are generated in a lazy way, just when they need to be generated and then they are returned by the yield statement as it’s hit. Hence, the line after the yield is executed just when it needs to be executed when the next value is requested.

About debugging the code I have to say that one of the best tool to write and debug Python code I know is from Microsoft and it’s Visual Studio Code. It’s really good and available for Windows, macOS and Linux for free.

Playing with iterable objects

Iterable objects give you a lot of possibilities. For example, if you need to create a list from the previous generator you can simply do

    my_fibonacci_list = list(fibonacci(100000))
    print("My fibonacci list: {0}".format(my_fibonacci_list))

My fibonacci list: [0, 1, 1, 2, 3, 5, 8, 13, 21, 34, 55, 89, 144, 233, 377, 610, 987, 1597, 2584, 4181, 6765, 10946, 17711, 28657, 46368, 75025

another way of creating a list from an iterable object is by using list comprehension that allows you to create a list in a very natural way, specifying also which elements to choose for the list. For example, if you need to create a list with only the odd Fibonacci numbers you can do

    fibonacci_odds_list = [x for x in fibonacci(100000) if x%2!=0]
    print("The odds number are: {0}".format(fibonacci_odds_list))

The odds number are: [1, 1, 3, 5, 13, 21, 55, 89, 233, 377, 987, 1597, 4181, 6765, 17711, 28657, 75025]

And you can use them also for all the functions based on iterables, like sum, max, min and so on …

    print("The min number is: {0}".format(min(fibonacci(1000000)))) 
    print("The max number is: {0}".format(max(fibonacci(1000000))))
    print("The sum of is: {0}".format(sum(fibonacci(1000000))))

The min number is: 0
The max number is: 832040
The sum is: 2178308

… or for functional programming functions like map and reduce, but this is another story for a future article. 🙂

Previous Entries Web-Test Automation in Python Next Entries Using virtual environments with Python