Python Iterators and Generators

Every for loop works the same way underneath. It asks the collection for an iterator, an object that hands out one item at a time, and keeps asking for the next item until there are none left. Understanding this lets you write your own sequences, and generators make that remarkably easy.

Iterables and iterators

An iterable is anything you can loop over: a list, a string, a dictionary, a file, a range. iter() gets an iterator from it, and next() takes the next item. When the items run out, next raises StopIteration, which is the signal a for loop waits for:

colours = ["red", "green", "blue"]
it = iter(colours)
print(next(it))
print(next(it))
print(next(it))
print(next(it, "no more"))

red
green
blue
no more

The second argument of next is returned instead of raising StopIteration. An iterator is used up once it has handed out all its items; to go round again, ask the iterable for a new iterator.

Generator functions

A function that contains yield is a generator function. Calling it does not run its body; it returns a generator, an iterator that runs the body step by step. Each yield hands out one value and pauses the function until the next value is asked for:

def countdown(n):
    print("starting")
    while n > 0:
        yield n
        n -= 1
    print("finished")

for x in countdown(3):
    print(x)

starting
3
2
1
finished

The function's local variables are kept between the steps, so a generator remembers where it was. Writing the same thing as a class with __iter__ and __next__ would take much more code.

Generator expressions and lazy sequences

A generator expression looks like a list comprehension with round brackets. It produces values one at a time instead of building the whole list in memory, which matters for very large or endless sequences:

total = sum(n * n for n in range(1, 1_000_001))
print(total)

def fibonacci():
    a, b = 0, 1
    while True:            # an endless sequence
        yield a
        a, b = b, a + b

from itertools import islice
print(list(islice(fibonacci(), 12)))

333333833333500000
[0, 1, 1, 2, 3, 5, 8, 13, 21, 34, 55, 89]

fibonacci never stops, but that is fine: it only computes the values that are asked for, and islice asks for the first twelve.

Useful iterator tools

The itertools module has building blocks for working with iterators: count counts for ever, cycle repeats a sequence, chain joins sequences, combinations and permutations produce arrangements. Built-ins such as enumerate, zip, map, filter and reversed also return iterators.

Practice exercises

1. Use iter and next to take the first two letters of the string "Python".
2. Write a generator evens(limit) that yields the even numbers up to a limit.
3. Write a generator that yields the lines of a list of strings in upper case, and loop over it.
4. Use itertools.combinations to print every pair of three names.