Skip to Content
Course content

170: Writing a Custom Iterator Class

Click on the "Edit" button in the top corner of the screen to edit your slide content.

Think about a ticket dispenser at a deli. You walk up to the machine, pull a slip of paper, and it gives you number 42. The machine doesn't hand you a giant stack of every possible ticket number at once—that would be a waste of paper and a mess in your hands. Instead, it maintains an internal state (the current number) and gives you exactly one ticket every time you interact with it. Once you've pulled a ticket, the machine is already primed to give the next person number 43.

This is exactly how a custom iterator works in Python. Instead of returning a full list of data (which could crash your program if the list has a billion items), an iterator gives you one item at a time, only when you ask for it. It remembers where it left off, just like that deli machine.

The Two-Method Contract

To make a class "iterable," you can't just wing it; you have to follow a specific protocol. I like to think of this as a contract between your class and Python's for loop. If you implement these two methods, Python will treat your object as an iterator:

  • __iter__(): This is the handshake. It tells Python, "Yes, I can be iterated over." It must return the iterator object itself.
  • __next__(): This is the actual "pulling the ticket" action. This method calculates the next value and returns it.

If you've used list or range, you've used these under the hood. But when you're building something complex—like a stream of data from a network socket or a mathematical sequence—you'll want to write your own.

Signaling the End of the Line

A deli machine eventually runs out of tickets. In Python, we don't return None or False to signal the end of a sequence, because those might actually be valid pieces of data you're trying to iterate over. Instead, we raise a StopIteration exception.

It feels weird to raise an exception to signal a normal event, but trust me, this is how for loops work. The loop internally catches StopIteration and quietly exits without crashing your program.

Implementing a Prime Number Sequence

Let's put this into practice. Suppose we want an iterator that generates prime numbers. We don't know how many primes the user will want, so calculating a list beforehand is impossible. We'll build a class that finds the next prime on demand.

class PrimeIterator:
    def __init__(self):
        self.current = 1

    def __iter__(self):
        # We return self because this class implements __next__
        return self

    def __next__(self):
        self.current += 1
        while True:
            if self._is_prime(self.current):
                return self.current
            self.current += 1

    def _is_prime(self, n):
        if n < 2: return False
        for i in range(2, int(n**0.5) + 1):
            if n % i == 0:
                return False
        return True

# Using our custom iterator
primes = PrimeIterator()
for p in primes:
    print(p)
    if p > 20: 
        break # We stop manually, otherwise this would run forever!

Notice that __iter__ simply returns self. I've seen a lot of beginners try to put the logic inside __iter__, but that's a mistake. __iter__ just initializes the process; __next__ does the heavy lifting. In the example above, the state is maintained by self.current, ensuring that every time the for loop asks for the next value, we start searching from where we last stopped.




📋 Practical Task

Build a Circular Buffer Iterator

In many systems, you need to loop through a set of items repeatedly (like rotating through a list of available servers for load balancing). Your task is to create a class called CircularIterator.

Requirements:

  • The class should take a list of items in its __init__ method.
  • It should implement __iter__ and __next__.
  • Instead of raising StopIteration, the __next__ method should wrap back around to the start of the list once it reaches the end.
  • The iterator should be infinite (it never stops).

Testing your code: Create an instance of CircularIterator with the list ['Red', 'Green', 'Blue'] and use a for loop with a break condition to print the first 7 items. You should see: Red, Green, Blue, Red, Green, Blue, Red.

Rating
0 0

There are no comments for now.

to be the first to leave a comment.