Iterators

5.1 Iterable and Iterator

As you already know, iterators are objects that implement the iterator protocol, allowing you to sequentially retrieve elements from a collection. Iterators are widely used in Python to loop through elements of sequences like lists, tuples, and strings.

Let's take a look at how iterators work and how to use them.

Iterable Object (Iterable)

For an object to be looped over with a for loop, it must be iterable – Iterable. This means that our object needs to implement the method __iter__(), which returns an iterator object.

Iterator Object (Iterator)

This is a special object that has a function __next__() that returns the next element of the sequence. When the elements are finished, the __next__() method raises the StopIteration exception as a signal to stop iteration.

An iterator should also implement the __iter__() method, which returns the iterator itself.

Example Using Built-in Python Functions

In this example, the list numbers is an iterable object. We get an iterator using the iter() function and use the next() function to iterate through elements until the StopIteration exception is raised.

# Iterable object
numbers = [1, 2, 3, 4, 5]
            
# Get an iterator from the iterable object
iterator = iter(numbers)
            
# Use the iterator to loop through elements
try:
    while True:
        number = next(iterator)
        print(number)
except StopIteration:
    pass
        

This is exactly what happens when you write code like:

# Iterable object
numbers = [1, 2, 3, 4, 5]
            
for number in numbers:
    print(number)
        

5.2 The Essence of an Iterator

An iterator is an object that helps us sequentially go through a group of elements. There can be many different implementations. Let's write our own class that implements all the requirements for an iterator.

Step 1. First, let's create our class

Let it sequentially return numbers from start to end

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

Step 2. Support the __iter__ function

Now we need to add the __iter__ function to it, which will return the iterator object that will call the __next()__ function. We will return a reference to our own object – this is allowed.

class MyIterator:
    def __init__(self, start, end):
        self.current = start
        self.end = end
        
    def __iter__(self):
        return self
        

Step 3. Support the __next__ function

Now we need to add the __next__ function to our iterator object, which will return the next element of our list. We will simply use the current variable:

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

Step 4. Stop the iterator

If the iterator has already returned all the values it planned to, it should raise the StopIteration exception. Let's tweak our last function a bit:

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

Great. Now we can use our iterator. Here's an example of all our code:

class MyIterator:
    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
        current = self.current
        self.current += 1
        return current
        
# Create an instance of the custom iterator
my_iter = MyIterator(1, 5)
        
# Use the iterator to loop through elements
for num in my_iter:
    print(num)

5.3 The Proper Iterator

What's wrong with the iterator from the previous example? Yes, it's an iterator, it works, but it's too primitive. You can't use it to iterate over the same collection of elements simultaneously with different iterators.

A better way would be to write code that doesn't return a reference to itself in the __iter__ method, but instead returns a separate object that would properly yield all the elements.

Example:

class MyIterable:
    def __init__(self, data):
        self.data = data
    
    def __iter__(self):
        return MyIterator(self.data)
    
class MyIterator:
    def __init__(self, data):
        self.data = data
        self.index = 0
    
    def __iter__(self):
        return self
    
    def __next__(self):
        if self.index >= len(self.data):
            raise StopIteration
        item = self.data[self.index]
        self.index += 1
        return item
    
# Usage
my_iterable = MyIterable([1, 2, 3, 4])
for item in my_iterable:
    print(item)
    

In this example, we have two classes — the first one takes a collection to be iterated over, and the second is the iterator itself, which returns the elements of the collection in the next() method. It's pretty simple, but this is exactly how you should add iterators to your classes.