Iterators are a fundamental concept in Python, allowing you to traverse through all the elements in a collection, such as lists or tuples. While Python provides built-in iterators, there are times when you might need to create your own custom iterator to handle more complex data structures or operations. In this article, we will explore how to build a basic iterator in Python from scratch.

What is an Iterator?

An iterator in Python is an object that implements two methods: __iter__() and __next__(). The __iter__() method returns the iterator object itself and is called once at the beginning of an iteration. The __next__() method returns the next value from the sequence and raises a StopIteration exception when there are no more items to return.

Define the Iterator Class:

• __init__: The constructor initializes the maximum value (max_n) and the current value (current), which starts at 0.
• __iter__: This method returns the iterator object itself. It is called once at the beginning of the iteration.
• __next__: This method calculates the square of the current value, increments the current value, and returns the result. If the current value exceeds max_n, it raises a StopIteration exception to signal the end of the iteration.

Why Use Custom Iterators?

Custom iterators allow you to:

• Encapsulate complex iteration logic.
• Maintain state between iterations.
• Generate values on the fly (lazy evaluation), which can be more memory-efficient for large datasets.

Example 1: Building a Basic Custom Iterator

Let’s build a basic iterator that generates a sequence of square numbers. We’ll start by creating a class that implements the iterator protocol.

Initialization: The SquareIterator class is initialized with a maximum value max_n, which determines how many squares will be generated. The current attribute keeps track of the current number being squared.

Iteration: The __iter__ method is straightforward; it returns the iterator object itself (self). This method is necessary to make our object an iterator.

Generating Values: The __next__ method is where the core logic resides. It checks if the current value has exceeded the maximum (max_n). If it has, it raises a StopIteration exception, which stops the iteration. Otherwise, it calculates the square of the current value, increments current, and returns the squared value.

class SquareIterator:
    def __init__(self, max_n):
        self.max_n = max_n
        self.current = 0

    def __iter__(self):
        return self

    def __next__(self):
        if self.current > self.max_n:
            raise StopIteration
        result = self.current ** 2
        self.current += 1
        return result

# Using the custom iterator
square_iter = SquareIterator(5)

for num in square_iter:
    print(num)

Output:

0
1
4
9
16
25

Example 2: Building a Fibonacci Iterator

We’ll create an iterator that generates Fibonacci numbers up to a specified limit.

1. Class Definition: FibonacciIterator is defined with an __init__ method that initializes the starting values of the Fibonacci sequence (self.a and self.b) and the limit up to which the numbers should be generated.
2. Iterator Method: The __iter__ method returns the iterator object itself.
3. Next Method: The __next__ method generates the next number in the Fibonacci sequence. If the current number exceeds the specified limit, it raises a StopIteration exception.
4. Using the Iterator: The iterator can be used with a for loop or the next() function to generate and print Fibonacci numbers up to the specified limit.

class FibonacciIterator:
    def __init__(self, limit):
        self.limit = limit
        self.a = 0
        self.b = 1
    
    def __iter__(self):
        return self
    
    def __next__(self):
        if self.a > self.limit:
            raise StopIteration
        current = self.a
        self.a, self.b = self.b, self.a + self.b
        return current

# Create an instance of FibonacciIterator
fib_iterator = FibonacciIterator(100)

# Using a for loop to iterate
print("Using for loop:")
for num in fib_iterator:
    print(num)

# Using next() to iterate
print("\nUsing next() function:")
fib_iterator = FibonacciIterator(100)
print(next(fib_iterator)) 
print(next(fib_iterator))  
print(next(fib_iterator))  
print(next(fib_iterator))  
print(next(fib_iterator))  
print(next(fib_iterator)) 
print(next(fib_iterator))  
print(next(fib_iterator))  
print(next(fib_iterator))  
print(next(fib_iterator))  
print(next(fib_iterator)) 
print(next(fib_iterator))  

Output:

Using for loop:
0
1
1
2
3
5
8
13
21
34
55
89

Using next() function:
0
1
1
2
3
5
8
13
21
34
55
89

Conclusion

Building a custom iterator in Python is a powerful way to create more complex and memory-efficient data traversal mechanisms. By implementing the __iter__ and __next__ methods, you can control the iteration logic and manage state between iterations. The example provided here demonstrates how to create a simple iterator that generates square numbers, but the principles can be applied to more complex scenarios as well.