Iterators and Iterables in Python: Run Efficient Iterations

What Is an Iterator in Python?

In Python, an iterator is an object that allows you to iterate over collections of data, such as lists, tuples, dictionaries, and sets.

Python iterators implement the iterator design pattern, which allows you to traverse a container and access its elements. The iterator pattern decouples the iteration algorithms from container data structures.

Iterators take responsibility for two main actions:

1. Returning the data from a stream or container one item at a time
2. Keeping track of the current and visited items

In summary, an iterator will yield each item or value from a collection or a stream of data while doing all the internal bookkeeping required to maintain the state of the iteration process.

Python iterators must implement a well-established internal structure known as the iterator protocol. In the following section, you’ll learn the basics of this protocol.

What Is the Python Iterator Protocol?

A Python object is considered an iterator when it implements two special methods collectively known as the iterator protocol. These two methods make Python iterators work. So, if you want to create custom iterator classes, then you must implement the following methods:

The .__iter__() method of an iterator typically returns self, which holds a reference to the current object: the iterator itself. This method is straightforward to write and, most of the time, looks something like this:

def __iter__(self):
    return self

The only responsibility of .__iter__() is to return an iterator object. So, this method will typically just return self, which holds the current instance. Don’t forget that this instance must define a .__next__() method.

The .__next__() method will be a bit more complex depending on what you’re trying to do with your iterator. This method must return the next item from the data stream. It should also raise a StopIteration exception when no more items are available in the data stream. This exception will make the iteration finish. That’s right. Iterators use exceptions for control flow.

When to Use an Iterator in Python?

The most generic use case of a Python iterator is to allow iteration over a stream of data or a container data structure. Python uses iterators under the hood to support every operation that requires iteration, including for loops, comprehensions, iterable unpacking, and more. So, you’re constantly using iterators without being conscious of them.

In your day-to-day programming, iterators come in handy when you need to iterate over a dataset or data stream with an unknown or a huge number of items. This data can come from different sources, such as your local disk, a database, and a network.

In these situations, iterators allow you to process the datasets one item at a time without exhausting the memory resources of your system, which is one of the most attractive features of iterators.

Yielding the Original Data

Okay, now it’s time to learn how to write your own iterators in Python. As a first example, you’ll write a classic iterator called SequenceIterator. It’ll take a sequence data type as an argument and yield its items on demand.

Fire up your favorite code editor or IDE and create the following file:

# sequence_iter.py

class SequenceIterator:
    def __init__(self, sequence):
        self._sequence = sequence
        self._index = 0

    def __iter__(self):
        return self

    def __next__(self):
        if self._index < len(self._sequence):
            item = self._sequence[self._index]
            self._index += 1
            return item
        else:
            raise StopIteration

Your SequenceIterator will take a sequence of values at instantiation time. The class initializer, .__init__(), takes care of creating the appropriate instance attributes, including the input sequence and an ._index attribute. You’ll use this latter attribute as a convenient way to walk through the sequence using indices.

The .__iter__() method does only one thing: returns the current object, self. In this case, self is the iterator itself, which implies it has a .__next__() method.

Finally, you have the .__next__() method. Inside it, you define a conditional statement to check if the current index is less than the number of items in the input sequences. This check allows you to stop the iteration when the data is over, in which case the else clause will raise a StopIteration exception.

In the if clause, you grab the current item from the original input sequence using its index. Then you increment the ._index attribute using an augmented assignment. This action allows you to move forward in the iteration while you keep track of the visited items. The final step is to return the current item.

Here’s how your iterator works when you use it in a for loop:

>>> from sequence_iter import SequenceIterator

>>> for item in SequenceIterator([1, 2, 3, 4]):
...     print(item)
...
1
2
3
4

Great! Your custom iterator works as expected. It takes a sequence as an argument and allows you to iterate over the original input data.

Before diving into another example, you’ll learn how Python’s for loops work internally. The following code simulates the complete process:

>>> sequence = SequenceIterator([1, 2, 3, 4])

>>> # Get an iterator over the data
>>> iterator = sequence.__iter__()
>>> while True:
...     try:
...         # Retrieve the next item
...         item = iterator.__next__()
...     except StopIteration:
...         break
...     else:
...         # The loop's code block goes here...
...         print(item)
...
1
2
3
4

After instantiating SequenceIterator, the code prepares the sequence object for iteration by calling its .__iter__() method. This method returns the actual iterator object. Then, the loop repeatedly calls .__next__() on the iterator to retrieve values from it.

When a call to .__next__() raises the StopIteration exception, you break out of the loop. In this example, the call to print() under the else clause of the try block represents the code block in a normal for loop. As you can see, the for loop construct is a kind of syntactic sugar for a piece of code like the one above.

Transforming the Input Data

Say that you want to write an iterator that takes a sequence of numbers, computes the square value of each number, and yields those values on demand. In this case, you can write the following class:

# square_iter.py

class SquareIterator:
    def __init__(self, sequence):
        self._sequence = sequence
        self._index = 0

    def __iter__(self):
        return self

    def __next__(self):
        if self._index < len(self._sequence):
            square = self._sequence[self._index] ** 2
            self._index += 1
            return square
        else:
            raise StopIteration

The first part of this SquareIterator class is the same as your SequenceIterator class. The .__next__() method is also pretty similar. The only difference is that before returning the current item, the method computes its square value. This computation performs a transformation on each data point.

Here’s how the class will work in practice:

>>> from square_iter import SquareIterator

>>> for square in SquareIterator([1, 2, 3, 4, 5]):
...     print(square)
...
1
4
9
16
25

Using an instance of SquareIterator in a for loop allows you to iterate over the square values of your original input values.

The option of performing data transformation on a Python iterator is a great feature. It can make your code quite efficient in terms of memory consumption. Why? Well, imagine for a moment that iterators didn’t exist. In that case, if you wanted to iterate over the square values of your original data, then you’d need to create a new list to store the computed squares.

This new list would consume memory because it would have to store all the data simultaneously. Only then would you be able to iterate over the square values.

However, if you use an iterator, then your code will only require memory for a single item at a time. The iterator will compute the following items on demand without storing them in memory. In this regard, iterators are lazy objects.

Generating New Data

You can also create custom iterators that generate a stream of new data from a given computation without taking a data stream as input. Consider the following iterator, which produces Fibonacci numbers:

# fib_iter.py

class FibonacciIterator:
    def __init__(self, stop=10):
        self._stop = stop
        self._index = 0
        self._current = 0
        self._next = 1

    def __iter__(self):
        return self

    def __next__(self):
        if self._index < self._stop:
            self._index += 1
            fib_number = self._current
            self._current, self._next = (
                self._next,
                self._current + self._next,
            )
            return fib_number
        else:
            raise StopIteration

This iterator doesn’t take a data stream as input. Instead, it generates each item by performing a computation that yields values from the Fibonacci sequence. You do this computation inside the .__next__() method.

You start this method with a conditional that checks if the current sequence index hasn’t reached the ._stop value, in which case you increment the current index to control the iteration process. Then you compute the Fibonacci number that corresponds to the current index, returning the result to the caller of .__next__().

When ._index grows to the value of ._stop, you raise a StopIteration, which terminates the iteration process. Note that you should provide a stop value when you call the class constructor to create a new instance. The stop argument defaults to 10, meaning the class will generate ten Fibonacci numbers if you create an instance without arguments.

Here’s how you can use this FibonacciIterator class in your code:

>>> from fib_iter import FibonacciIterator

>>> for fib_number in FibonacciIterator():
...     print(fib_number)
...
0
1
1
2
3
5
8
13
21
34

Instead of yielding items from an existing data stream, your FibonacciIterator class computes every new value in real time, yielding values on demand. Note that in this example, you relied on the default number of Fibonacci numbers, which is 10.

Coding Potentially Infinite Iterators

An interesting feature of Python iterators is that they can handle potentially infinite data streams. Yes, you can create iterators that yield values without ever reaching an end! To do this, you just have to skip the StopIteration part.

For example, say that you want to create a new version of your FibonacciIterator class that can produce potentially infinite Fibonacci numbers. To do that, you just need to remove the StopIteraton and the condition that raises it:

# inf_fib.py

class FibonacciInfIterator:
    def __init__(self):
        self._index = 0
        self._current = 0
        self._next = 1

    def __iter__(self):
        return self

    def __next__(