Custom Iterators

Iterators are one of the foundational concepts in Python, enabling you to traverse through data structures like lists, tuples, and dictionaries effortlessly. But what if you want to create your own object that behaves like these built-in iterable types? This is where custom iterators come into play.

By mastering custom iterators, you gain fine-grained control over how your objects are traversed in a for loop or any context expecting an iterable. This lesson will guide you through the inner workings of iterators in Python, step-by-step, and show you how to craft your own iterator classes from scratch.

Understanding the Iterator Protocol

In Python, the iteration mechanism relies on something called the iterator protocol. This protocol defines two essential methods:

• __iter__(self): Returns the iterator object itself. Called once when an iteration starts.
• __next__(self): Returns the next item in the sequence. Raises StopIteration when no more items are available.

When you use a for loop, Python internally calls iter() on the object, which in turn calls the object's __iter__() method to get an iterator. Then, it repeatedly calls next() (which calls __next__()) on that iterator until StopIteration is raised.

Creating a Simple Custom Iterator

Let's build a simple iterator that returns numbers from 1 up to a given limit. This example will clarify how __iter__() and __next__() work together.

class NumberIterator:
    def __init__(self, limit):
        self.limit = limit
        self.current = 1

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

    def __next__(self):
        if self.current <= self.limit:
            number = self.current
            self.current += 1
            return number
        else:
            # No more data to return, stop iteration
            raise StopIteration

# Using the custom iterator
numbers = NumberIterator(5)
for num in numbers:
    print(num)

In this example, the NumberIterator class keeps track of the current number and the upper limit. The __next__() method returns the current number and increments it until it reaches the limit, then raises StopIteration to signal the end.

Why Implement Both __iter__() and __next__()?

Every iterator needs to implement these methods. But what if you want to create an iterable container that can provide multiple independent iterators?

Typically, you separate the iterable and the iterator into two classes:

• The iterable class implements __iter__(), which returns a new iterator object each time.
• The iterator class implements __next__() and __iter__().

This approach allows multiple loops over the same iterable to work independently without interfering with each other's state.

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

    def __iter__(self):
        return self

    def __next__(self):
        if self.current > 0:
            val = self.current
            self.current -= 1
            return val
        else:
            raise StopIteration

class CountDown:
    def __init__(self, start):
        self.start = start

    def __iter__(self):
        # Return a new CountDownIterator each time
        return CountDownIterator(self.start)

# Demonstrate independent iterators
cd = CountDown(3)

for x in cd:
    print(x, end=' ')
print()

for y in cd:
    print(y, end=' ')

Notice that each for loop uses a fresh iterator starting from 3, so they don't interfere with each other.

Common Use Cases for Custom Iterators

Custom iterators are particularly useful when:

• You want to traverse a custom data structure (e.g., trees, graphs, linked lists).
• You need lazy evaluation, generating items on the fly rather than storing all at once.
• You want to encapsulate complex iteration logic cleanly.
• You want to implement infinite sequences or streams.

For example, iterators are great to generate Fibonacci numbers or prime numbers without precomputing a large list.

Example: Fibonacci Iterator

Let's create an iterator that generates Fibonacci numbers up to a certain count:

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

    def __iter__(self):
        return self

    def __next__(self):
        if self.index < self.count:
            self.index += 1
            fib = self.a
            self.a, self.b = self.b, self.a + self.b
            return fib
        else:
            raise StopIteration

# Use the Fibonacci iterator
fib_seq = Fibonacci(7)
for num in fib_seq:
    print(num, end=' ')

This iterator calculates Fibonacci numbers on-demand, making it very memory efficient compared to storing the entire sequence beforehand.

What Happens Under the Hood When You Use for?

When you write for item in obj:, Python internally translates it roughly as:

iterator = iter(obj)
while True:
    try:
        item = next(iterator)
    except StopIteration:
        break
    # process item

This shows why implementing these two methods correctly allows your custom objects to integrate seamlessly with Python's iteration constructs.

Built-in Iterator Types vs Custom Iterators

Python comes with many built-in iterators, such as those for lists, tuples, dictionaries, and files. Here's a quick comparison between built-in and custom iterators:

Practical Tips When Writing Iterators

• Always raise StopIteration to end iteration: This is the signal Python looks for to stop looping.
• Keep iterator state inside the object: Store indices or pointers as instance variables.
• Return self in __iter__() for iterator classes: This allows the object to be used directly in loops.
• Make separate iterator classes if you want multiple simultaneous iterations: This avoids shared state issues.
• Use generators as a simpler alternative for many cases: Sometimes yield can produce iterators with less boilerplate.

Summary

Custom iterators provide a powerful way to control how your objects produce sequences of values. By implementing __iter__() and __next__(), you enable your classes to integrate seamlessly with Python's iteration tools, including for loops, comprehensions, and functions that consume iterables.

Key takeaways:

• The iterator protocol requires __iter__() and __next__().
• __iter__() returns an iterator object (usually self).
• __next__() returns the next item or raises StopIteration.
• Separating iterable and iterator classes allows multiple independent iterations.
• Custom iterators enable lazy, memory-efficient, and complex iteration logic.

With this knowledge, you can now craft your own iterators to fit any custom data traversal needs.