Iterators & Generators

Every for loop in Python runs on the iterator protocol. Generators — functions that yield values one at a time — make it easy to create your own iterators, process data lazily, and handle streams far larger than memory.

Iterables and iterators

An iterable is anything you can loop over (lists, strings, dictionaries, files). Calling iter() on it gives an iterator, which hands out one item per call to next() and raises StopIteration when it runs out:

colors = ["red", "green"]
it = iter(colors)

print(next(it))
print(next(it))
print(next(it, "no more items"))

red
green
no more items

A for loop does exactly this behind the scenes: call iter(), then next() repeatedly until StopIteration.

Generator functions

A function containing yield becomes a generator function. Calling it doesn't run the body; it returns a generator, which runs the body step by step, pausing at each yield:

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

gen = countdown(3)
print(type(gen).__name__)
print(next(gen))
print(next(gen))
print(list(gen))

generator
starting
3
2
done
[1]

"starting" only prints when the first value is requested, and the function's local variables are preserved between yields.

Lazy evaluation

Generators produce values on demand, so they use almost no memory regardless of how many values there are — and can even be infinite:

import itertools

def squares():
    n = 1
    while True:
        yield n * n
        n += 1

first_ten = list(itertools.islice(squares(), 10))
print(first_ten)

big = (n * n for n in range(10_000_000))
print(next(big), next(big))

[1, 4, 9, 16, 25, 36, 49, 64, 81, 100]
0 1

The generator expression on the last lines would describe ten million squares, but only the two requested are ever computed.

Pipelines

Chaining generators creates efficient data pipelines, where each item flows through every step before the next item is read:

log = """INFO start
ERROR disk full
INFO retry
ERROR timeout
WARN slow"""

def read_lines(text):
    for line in text.splitlines():
        yield line

def errors_only(lines):
    for line in lines:
        if line.startswith("ERROR"):
            yield line

def messages(lines):
    for line in lines:
        yield line.split(" ", 1)[1]

for msg in messages(errors_only(read_lines(log))):
    print(msg)

disk full
timeout

With a real log file, read_lines would read one line at a time — the whole pipeline would handle gigabytes with tiny memory use.

yield from

def walk(tree):
    yield tree["name"]
    for child in tree.get("children", []):
        yield from walk(child)

site = {"name": "home", "children": [{"name": "blog", "children": [{"name": "post-1"}]}, {"name": "about"}]}
print(list(walk(site)))

['home', 'blog', 'post-1', 'about']

yield from passes along every value from another iterable — ideal for recursive structures.

Generators are single-use

gen = (n for n in range(3))
print(list(gen))
print(list(gen))

[0, 1, 2]
[]

Once consumed, a generator is empty. If you need the values twice, store them in a list or create a new generator.

Writing an iterator class

For full control, a class can implement __iter__ and __next__. Often, though, making __iter__ a generator is simpler:

class Countdown:
    def __init__(self, start):
        self.start = start

    def __iter__(self):
        n = self.start
        while n > 0:
            yield n
            n -= 1

timer = Countdown(3)
print(list(timer), list(timer))

[3, 2, 1] [3, 2, 1]

Because __iter__ creates a fresh generator each time, a Countdown object can be looped over repeatedly.

Built-ins that work with iterators

names = ["ana", "ben", "cleo"]
scores = [91, 78, 85]

print(list(enumerate(names, 1)))
print(list(zip(names, scores)))
print(list(map(str.title, names)))
print(sum(s for s in scores if s > 80))
print(sorted(zip(scores, names), reverse=True)[0])

[(1, 'ana'), (2, 'ben'), (3, 'cleo')]
[('ana', 91), ('ben', 78), ('cleo', 85)]
['Ana', 'Ben', 'Cleo']
176
(91, 'ana')

enumerate, zip, map, filter, reversed and generator expressions all return lazy iterators. The itertools module (lesson 51) provides many more.

Try it yourself