9.3 Custom Iterators in Python

NicePython

— 5 min read

In Python, a custom iterator is an object you create that implements the iteration protocol by defining two methods: __iter__() and __next__(). This allows you to control how and when the next item is produced during iteration. Custom iterators are useful when you want to build a special sequence, manage complex data flows, or perform specific actions during iteration.

In this section, we will explore how to create custom iterators, how they work, and how to implement the iteration protocol in your own classes. We’ll also look at examples that demonstrate different use cases for custom iterators.

9.3.1 The Iteration Protocol

To create a custom iterator in Python, you need to define a class that implements two key methods:

  1. __iter__(): This method should return the iterator object itself and is used to retrieve an iterator from an iterable.
  2. __next__(): This method is used to return the next element in the sequence. It raises a StopIteration exception when there are no more elements to return.

An object that implements both of these methods is considered an iterator.

9.3.2 Creating a Simple Custom Iterator

Let's start with a basic example of a custom iterator that returns numbers from a starting point up to an end point.

Example: Simple Range Iterator:

class MyRange:
    def __init__(self, start, end):
        self.current = start
        self.end = end

    def __iter__(self):
        return self

    def __next__(self):
        if self.current >= self.end:
            raise StopIteration  # Stop iteration when the end is reached
        else:
            self.current += 1
            return self.current - 1  # Return the current number

# Using the custom iterator
my_range = MyRange(1, 5)
for number in my_range:
    print(number)

Output:

1
2
3
4

In this example:

  • The MyRange class implements the __iter__() and __next__() methods, making it a custom iterator.
  • The __next__() method returns the next number in the range from start to end - 1. When it reaches end, it raises a StopIteration exception to signal that the iteration is complete.
  • The for loop automatically handles the iteration using the __next__() method.

9.3.3 The __iter__() Method

The __iter__() method returns the iterator object itself. This method is called at the start of iteration when the for loop or another iteration mechanism is used. It allows an object to be treated as an iterable, even if it's also an iterator.

Example:

class MyRange:
    def __init__(self, start, end):
        self.current = start
        self.end = end

    def __iter__(self):
        return self  # The object itself is the iterator

In this example:

  • The __iter__() method returns the object itself, which implements the iteration logic in __next__(). This is standard for custom iterators.

9.3.4 The __next__() Method

The __next__() method defines how to retrieve the next element in the iteration. It should raise a StopIteration exception when there are no more elements to return, which signals the end of the iteration.

Example:

class MyRange:
    def __init__(self, start, end):
        self.current = start
        self.end = end

    def __next__(self):
        if self.current >= self.end:
            raise StopIteration  # Stop when end is reached
        else:
            self.current += 1
            return self.current - 1  # Return the next element

In this example:

  • The __next__() method increments the current value and returns it. Once the end of the range is reached, it raises the StopIteration exception to signal that iteration is finished.

9.3.5 Example: Custom Iterator for Fibonacci Sequence

Let’s create a custom iterator that generates numbers in the Fibonacci sequence. The Fibonacci sequence is a series of numbers where each number is the sum of the two preceding ones, starting from 0 and 1.

Example: Fibonacci Iterator:

class Fibonacci:
    def __init__(self, max_count):
        self.prev = 0
        self.current = 1
        self.count = 0
        self.max_count = max_count

    def __iter__(self):
        return self

    def __next__(self):
        if self.count >= self.max_count:
            raise StopIteration  # Stop after max_count numbers
        self.count += 1
        fib_number = self.prev
        self.prev, self.current = self.current, self.prev + self.current
        return fib_number

# Using the Fibonacci iterator
fib_iter = Fibonacci(10)  # Generate first 10 Fibonacci numbers
for number in fib_iter:
    print(number)

Output:

0
1
1
2
3
5
8
13
21
34

In this example:

  • The Fibonacci class is a custom iterator that generates the Fibonacci sequence up to max_count numbers.
  • The __next__() method calculates the next Fibonacci number by adding the previous two numbers in the sequence.
  • The iteration stops when the count exceeds max_count, and a StopIteration exception is raised.

9.3.6 Infinite Iterators

You can also create infinite iterators, which don’t raise StopIteration and continue generating values indefinitely. These are useful when you want to iterate over an endless stream of data.

Example: Infinite Counter Iterator:

class InfiniteCounter:
    def __init__(self, start=0):
        self.current = start

    def __iter__(self):
        return self

    def __next__(self):
        self.current += 1
        return self.current

# Using the infinite iterator
counter = InfiniteCounter()
for number in counter:
    if number > 5:
        break  # Stop manually after 5 iterations
    print(number)

Output:

1
2
3
4
5

In this example:

  • The InfiniteCounter class generates an endless sequence of numbers, starting from a specified value.
  • Since this iterator doesn’t raise StopIteration, it would run indefinitely if not for the manual stopping condition (if number > 5: break).

9.3.7 Combining Iterators with Built-In Functions

Python provides several built-in functions that work well with iterators:

  • iter(): Converts an iterable to an iterator.
  • next(): Retrieves the next item from an iterator.
  • enumerate(): Returns an iterator that produces pairs of an index and the corresponding item.
  • zip(): Combines multiple iterators into a single iterator of tuples.

Example: Using zip() and enumerate() with Custom Iterators:

class MyRange:
    def __init__(self, start, end):
        self.current = start
        self.end = end

    def __iter__(self):
        return self

    def __next__(self):
        if self.current >= self.end:
            raise StopIteration
        self.current += 1
        return self.current - 1

# Using zip() to combine two iterators
my_range1 = MyRange(1, 4)
my_range2 = MyRange(10, 13)

for x, y in zip(my_range1, my_range2):
    print(x, y)

# Using enumerate() to track indices
my_range = MyRange(5, 8)
for index, value in enumerate(my_range):
    print(f"Index {index}: Value {value}")

Output:

1 10
2 11
3 12
Index 0: Value 5
Index 1: Value 6
Index 2: Value 7

In this example:

  • zip() combines two custom iterators (my_range1 and my_range2) into a single iterator that returns pairs of elements.
  • enumerate() adds an index to each element produced by the custom iterator my_range.

9.3.8 Best Practices for Custom Iterators

  1. Use Iterators When Needed: Only create custom iterators when they are necessary for handling complex sequences or special iteration logic. Built-in iterators for lists, tuples, etc., are usually sufficient for simple tasks.
  2. Handle StopIteration Gracefully: Always raise StopIteration to signal the end of iteration when appropriate.
  3. Use __iter__() and __next__() Together: Ensure your class implements both __iter__() and __next__() methods to follow the iteration protocol.
  4. Avoid Side Effects: Custom iterators should avoid introducing unintended side effects during iteration, such as modifying the internal state in ways that could cause errors in subsequent iterations.

9.3.9 Summary

  • A custom iterator in Python is an object that implements the iteration protocol by defining the __iter__() and __next__() methods.
  • The __iter__() method returns the iterator object itself, while the __next__() method defines how to retrieve the next element in the iteration.
  • Custom iterators allow you to generate complex sequences or manage data flows during iteration.
  • You can create finite iterators that raise StopIteration when finished or infinite iterators that continue producing elements indefinitely.
  • Built-in Python functions like zip(), enumerate(), and next() work seamlessly with custom iterators, providing additional flexibility in iteration.

By mastering custom iterators, you can create efficient, flexible data structures and algorithms that are tailored to your specific needs and optimize iteration performance in your Python programs.

Read more