Iterators and generators

Lesson 30 of 62, about 3 min

Anything you can put in a for loop is an iterable. Under the hood, Python asks it for an iterator and then repeatedly asks the iterator for the next value. Knowing this protocol lets you build your own lazy sequences with generators.

iter() and next()

python

letters = ["a", "b", "c"]
it = iter(letters)
print(next(it))
print(next(it))
print(next(it))
try:
    next(it)
except StopIteration:
    print("exhausted")

A for loop does exactly this, and catches StopIteration for you.

Writing an iterator class

An iterator has __iter__ (returns itself) and __next__ (returns the next value or raises StopIteration).

python

class Countdown:
    def __init__(self, start):
        self.n = start
    def __iter__(self):
        return self
    def __next__(self):
        if self.n <= 0:
            raise StopIteration
        self.n -= 1
        return self.n + 1

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

That is a lot of boilerplate. Generators do the same with far less code.

Generator functions

A function that contains yield returns a generator. Each yield hands out one value and pauses; the function resumes from that point on the next request.

python

def countdown(start):
    while start > 0:
        yield start
        start -= 1

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

print(list(countdown(5)))

Generators are lazy

Values are produced only when asked. This means a generator can represent an infinite sequence, or process a huge file without loading it all into memory.

python

def naturals():
    n = 1
    while True:
        yield n
        n += 1

gen = naturals()
print(next(gen), next(gen), next(gen))

def first_n(iterable, n):
    for i, item in enumerate(iterable):
        if i >= n:
            return
        yield item

print(list(first_n(naturals(), 5)))

Generators are single use

Once exhausted, a generator stays empty. Create a new one if you need to loop again.

python

g = (x * x for x in range(3))
print(list(g))
print(list(g))

Generator expressions

Like a list comprehension but with parentheses. Nothing is built up front.

python

squares = (n * n for n in range(1, 6))
print(sum(squares))

# passing straight into a function needs no extra parentheses
print(max(len(w) for w in ["hi", "hello", "hey"]))

Use a generator expression instead of a list comprehension when you only need to iterate once, especially over large ranges.

yield from

Delegates to another iterable.

python

def chain(*iterables):
    for it in iterables:
        yield from it

print(list(chain([1, 2], "ab", range(3))))

The itertools module

The standard library has a toolbox of iterator helpers.

python

import itertools

print(list(itertools.islice(itertools.count(10, 5), 4)))
print(list(itertools.chain([1, 2], [3])))
print(list(itertools.combinations("abc", 2)))
print(list(itertools.product([0, 1], repeat=2)))

Practice

Write a generator evens(limit) that yields even numbers up to limit.
Write a generator that yields the Fibonacci sequence forever, and print the first 10 values with itertools.islice.
Rewrite Countdown as a generator function in three lines.