Python Custom Iterator

1. Create a Custom Iterator With the Help of Class Methods __iter__ and __next__ in Python
2. Create a Custom Iterator With the Help of a Generator in Python

In this article, we will learn what iterators are and how to create a custom iterator with the help of class methods called __iter__ and __next__. We will also learn how to create a custom iterator with the help of a generator in Python.

Create a Custom Iterator With the Help of Class Methods __iter__ and __next__ in Python

Now when you say iterator, as the name suggests, it will be used for iterations.

If you come from languages like C, C++, and Java, we have a common concept: ++. One more concept is known when we talk about iterations: the for loop.

Let’s say we have a list of 5 elements. Using a loop, we can iterate from the first element to the last element.

This way, we have a list with some values, and if we want to print them one by one using the index numbers.

If we pass our index 0, it will give us the first element, 2. If we use index value 1, it will provide us with the second element, 5.

Numbers = [2, 5, 10, 3, 1]

print(Numbers[0])
print(Numbers[1])

Output:

2
5

The other way of printing the values is using a loop. We can iterate on a list and print all the values.

Numbers = [2, 5, 10, 3, 1]

for i in Numbers:
    print(i)

Output:

2
5
10
3
1

Now we have one more way: by using an iterator. In general, an iterator works behind the scene of the for loop.

We will use an iter() function to understand an iterator, which will convert our list into an iterator. This iterator will not give us all the values; it will give just one value at a time.

Numbers = [2, 5, 10, 3, 1]

ITR = iter(Numbers)
print(ITR)

We can see it is printing the object of an iterator.

<list_iterator object at 0x000001FEDCFF7310>

But, if you want the value, you can set ITR.__next__(), which is an inbuilt method. It will give us the first value during the first call of the __next__() method on a particular object; this is the same way we used the index value, but the advantage is we do not have to use the index value.

umbers = [2, 5, 10, 3, 1]

ITR = iter(Numbers)
print(ITR.__next__())

Output:

2

It will give the next value when we execute ITR.__next__() in the next execution line. Behind this scene, the iterator will have multiple values, so when we call the __next__() method, it will pick up one value.

And again, when we call __next__(), it knows the last value of i, which means it will preserve the state of the last value. That is the beauty of iterator; when we call the function again, it will preserve the old value.

We can create our iterator using a class. Using the class, we can print the top 10 values individually.

To achieve that, we will have a class called TopValues, and in this class, we will specify a counter.

And for that, we will use the __init__ function, where we will define our counter variable self.N=1; obviously, the counter will be initialized from 1. We need two important methods to create our custom iterator.

The first one is the iter() method, which will give us the object of the iterator, and then the next() method will provide the next value or object. In the __next__ method, instead of returning self.N, we use a variable called VALUE in which we will assign self.N, and in the next line, we increment self.N with one.

So that it will be incremented in each next iteration and then return the VALUE not self.N because self.N is incremented for the next iteration.

class TopValues:
    def __init__(self):
        self.N = 1

    def __iter__(self):
        return self

    def __next__(self):
        VALUE = self.N
        self.N += 1
        return VALUE

Our iterator is now ready so that we can create an object of the TopValues class. Let’s use a loop because usually when you have an iterator, we can use loops.

T_Val = TopValues()

for i in T_Val:
    print(i)

Now when we execute this code, we get thousands of values.

1
2
3
4
5
....
1000

Let’s try to understand what is happening. To check if the iterator is working, we will use __next__().

print(T_Val.__next__())

This line will print the first iteration value, 1, but what goes wrong with the loop? The problem is that the loop will go from start to end; we assumed the end would be 10, but we have not mentioned where we should stop.

When we use a loop, it will call the next() function, which is how it works. The for loops internally use the next() function; hence, it will call this next() function again and again.

We must apply a condition inside the __next__ method.

We also have to set the else block; inside this block, we will raise an exception. Otherwise, after printing 10 values, the loop will print None values.

class TopValues:
    def __init__(self):
        self.N = 1

    def __iter__(self):
        return self

    def __next__(self):
        if self.N <= 10:
            VALUE = self.N
            self.N += 1
            return VALUE
        else:
            raise StopIteration


T_Val = TopValues()

for i in T_Val:
    print(i)

Output:

1
2
3
4
5
6
7
8
9
10

Create a Custom Iterator With the Help of a Generator in Python

Let’s now try to create a generator function that does the same thing. Generators are a lot simpler since they take care of the __iter__() and __next__() methods for us.

If we try to write a generator function called Our_Gen(), we will pass the parameter inside this function.

Since we are looping over the list, we are yielding one item at a time. And, when there are no more items to loop over, it will automatically take care of that and raise an exception for stop iteration.

Now we are looping over the Our_Gen() function, so it should print out one item at a time. If we run that, we can see that we get one at a time.

def Our_Gen(n_list):
    for i in n_list:
        yield i


Func = Our_Gen([2, 5, 10, 3, 1])

for i in Func:
    print(i)

Output:

2
5
10
3
1

To ensure that our next() function is still working. After executing, we get five elements we get a stop iteration exception.

def Our_Gen(n_list):
    for i in n_list:
        yield i


Func = Our_Gen([2, 5, 10, 3, 1])

print(Func.__next__())
print(Func.__next__())
print(Func.__next__())
print(Func.__next__())
print(Func.__next__())
print(Func.__next__())

Output:

2
5
10
3
1
Traceback (most recent call last):
  File "c:\Users\Dell\Desktop\demo\demo.py", line 56, in <module>
    print(Func.__next__())
StopIteration

The generators are easier to write than classes. But depending on your use case, you might need to know how to do the __iter__() method and the __next__() method in a class.

You can read more about creating a custom iterator from here.