Iterators

A for loop goes through a list, a string, a dictionary or a file. You can make your own classes work with for too. It's useful for a class whose job is to hold a collection of things, such as a playlist of songs or a team of players.

First, a look at what a for loop does. It asks the collection for an iterator, an object that hands out the items one at a time, and then keeps asking the iterator for the next item. The built-in functions iter and next do the same steps by hand:

>>> songs = iter(["Highway Song", "Slow River"])
>>> next(songs)
'Highway Song'
>>> next(songs)
'Slow River'
>>> next(songs)
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
StopIteration

When the items run out, the iterator raises the exception StopIteration. A for loop catches it by itself, quietly, and that's how the loop knows it's done.

A class of your own doesn't do any of this yet:

class Playlist:
    def __init__(self):
        self._songs = ["Highway Song", "Slow River"]


for song in Playlist():
    print(song)

Traceback (most recent call last):
  File "main.py", line 6, in <module>
    for song in Playlist():
                ~~~~~~~~^^
TypeError: 'Playlist' object is not iterable

iter and next

To make a class iterable, give it two methods:

• __iter__ runs when a loop starts. It sets up the counting from the start, and returns the iterator. The simplest way is to let the object be its own iterator and return self.
• __next__ returns the next item, and raises StopIteration when there are none left.

class Playlist:
    def __init__(self):
        self._songs = []

    def add(self, song):
        self._songs.append(song)

    def __iter__(self):
        self._index = 0
        return self

    def __next__(self):
        if self._index >= len(self._songs):
            raise StopIteration
        song = self._songs[self._index]
        self._index += 1
        return song


mix = Playlist()
mix.add("Highway Song")
mix.add("Slow River")
mix.add("River Bend")
for song in mix:
    print(song)
print(list(mix))
print("Slow River" in mix)

Highway Song
Slow River
River Bend
['Highway Song', 'Slow River', 'River Bend']
True

_index remembers how far the iteration has got. __iter__ sets it back to 0, so every new loop starts from the first song. Everything else that goes through items one by one works now too: list(mix) makes a list of the songs, and in goes through them looking for a match.

The items __next__ returns don't have to be stored in a list. An iterator can work them out as it goes, for example the next number in a series.