Python - Creating your own Iterator

In our first example, we will create iterable objects, which, when iterated over, will give out cubes of numbers, and these objects will support multiple iterations.

class Cubes:
    def __init__(self, start, stop):
        self.start = start
        self.stop = stop

    def __iter__(self):
        return CubesIterator(self)


class CubesIterator:
    def __init__(self, source):
       self.source = source  
       self.current = source.start

    def __next__(self):
        if self.current > self.source.stop:
            raise StopIteration
        else:
            x = self.current
            self.current += 1
            return x * x * x


x = Cubes(2,8)
for i in x:
    print(i, end=' ')
print('Sum =',sum(x))

Output-

8 27 64 125 216 343 512 Sum = 1295

Let us understand the code for the class Cubes. The initializer method takes two arguments, and inside it, we have created two instance variables start and stop. The __iter__ method should return a fresh iterator object every time it is invoked. So, inside this method we create and return an instance of the class CubesIterator. Now, let us understand the code for the class CubesIterator.

The __init__ method has a parameter named source which will accept the iterable that has to be iterated over. In the __iter__ method of the Cubes class, when we create an instance of this CubesIterator class, we have sent self as an argument, which is the iterable that needs to be iterated over.

Inside the __init__ method, we have written self.source = source, so now this iterator class has access to all the instance variables of the object named source. Next we create an instance variable current that is set equal to source.start. Now, let us see the __next__ method.

If the current is greater than the stop of the source object, then a StopIteration exception will be raised. If it is not, then the cube of the current number will be returned, and the value of the current will be increased by 1. So, this class will create stateful iterator objects that store the current state. The state in this iterator is kept inside the instance variable current. When the method __next__ is called, it produces and returns the result for the current call and it also modifies the state for the next call. So, you can think of an iterator as a value factory; whenever you request the next value from it, it knows how to compute it because it holds the current internal state. It remembers its state between calls.

In our next example, we have created a class named Fibonacci which will create iterable objects that give out Fibonacci numbers up to a certain value. In Fibonacci series, each number is the sum of previous two numbers.

# Fibonacci series: 0 1 1 2 3 5 8 13 21 34 55 89

class Fibonacci:
    def __init__(self, max):
        self.max = max

    def __iter__(self):
        return FiboIterator(self)


class FiboIterator:
    def __init__(self, source):
        self.source = source
        self.a = 0
        self.b = 1

    def __next__(self):
        f = self.a
        if self.a > self.source.max:
            raise StopIteration
        self.a, self.b = self.b, self.a + self.b
        return f


x = Fibonacci(100)
for i in x:
    print(i, end=' ')
print(55 in x, 50 in x)

Output-

0 1 1 2 3 5 8 13 21 34 55 89 True False

We have created an instance object of class Fibonacci and used it in a for loop and in the in operator. The class Fibonacci has two methods __init__ and __iter__. The __init__ method takes an argument and creates an instance variable max which denotes the number up to which we want to generate the Fibonacci numbers. The __iter__ method returns an iterator object of the class FiboIterator. This is the class that maintains the state information.

In the __init__ method of FiboIterator, we have created three instance variables. source is the iterable that has to be iterated over. The state inside the iterator is maintained with the help of instance variables a and b. The instance variable a is initialized to 0 and b to 1.

The __next__ method is responsible for calculating and returning the next term of the series.

In the two examples that we have seen, we created two classes each: the iterable class and the iterator class. The iterator class needs to access the data members of the iterable class, and that is why the iterable class passes a reference of its current object to the initializer of the iterator class. Instead of passing reference to the object, you can simply pass the instance variable that will be needed by the iterator. For example, in FiboIterator you could pass self.max, and in CubesIterator, you could pass self.start and self.stop. It will also work, but it is better to pass the self object instead of passing individual instance variables. When you pass the source object, in the iterator you only have to create instance variables that are responsible for maintaining the state.

We saw two examples where we created objects that supported multiple active iterations. If we do not want our objects to support multiple iterations then the whole thing can be put inside a single class.

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

    def __iter__(self):
        return self

    def __next__(self):
        f = self.a
        if f > self.max:
            raise StopIteration
        self.a, self.b = self.b, self.a + self.b
        return f


x = Fibonacci(100)
for i in x:
    print(i, end = ' ')
print(50 in x, 55 in x)

Output-

0 1 1 2 3 5 8 13 21 34 55 89 False False

In this program, we have only a single class, and its objects will support a single iteration.

The __iter__ method now returns self instead of returning a fresh iterator. The __next__ method is now written in this class itself instead of a separate class. The state-maintaining variables a and b are also a part of this class. Now, we do not need a separate iterator class; this class is its own iterator.

From the output, we can see that object x supports only a single iteration. The in operator did not work because after the for loop, the iterator was exhausted. When the in operator demanded an iterator, it got the same exhausted iterator because now the __iter__ method returns self instead of returning a fresh iterator.

This way, we can write a single class for both the iterator and the iterable. This type of class will create an iterator object that supports only a single iteration.