Iterators and Generators

The iteration protocol

A for loop works on lists, strings, dictionaries, files, ranges and more. They have something in common, and it is a small protocol of two ideas:

• an iterable is anything you can loop over. It can produce an iterator when asked, with iter().
• an iterator is the object that does the work. Each call to next() gives the next item, and when there are none left it raises StopIteration.

numbers = [10, 20, 30]
iterator = iter(numbers)

print(next(iterator))
print(next(iterator))
print(next(iterator))

try:
    next(iterator)
except StopIteration:
    print("no more items")

A for loop is a neat way of writing exactly that. Here it is written out by hand:

def my_for(iterable, body):
    iterator = iter(iterable)
    while True:
        try:
            item = next(iterator)
        except StopIteration:
            break
        body(item)

my_for("abc", print)

Writing your own iterator

Any class with two methods works anywhere Python expects something to loop over: __iter__ (returning the iterator itself) and __next__ (returning the next value, or raising StopIteration).

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

    def __iter__(self):
        return self

    def __next__(self):
        if self.current < 0:
            raise StopIteration
        value = self.current
        self.current -= 1
        return value

print(list(Countdown(3)))
for n in Countdown(2):
    print("liftoff in", n)

Iterators are used up

A list can be looped over again and again, because every call to iter() on it gives a fresh iterator. An iterator itself remembers its place and does not rewind:

items = iter([1, 2, 3])
print(sum(items))
print(sum(items))                  # nothing is left

numbers = [1, 2, 3]
print(iter(numbers) is iter(numbers))      # each call makes a new iterator
print(iter(items) is items)                # an iterator is its own iterator

Generators: functions that pause

Writing a class with __iter__ and __next__ takes a lot of code for a simple idea. A generator function does the same thing with far less. Put yield where you would put return, and the function becomes a factory for iterators.

def countdown(n):
    while n >= 0:
        yield n               # hand back a value, and pause here
        n -= 1

gen = countdown(3)
print(type(gen).__name__)
print(next(gen), next(gen))
print(list(gen))              # whatever is left

Calling a generator function does not run it. It returns a generator object. The body only runs when something asks for a value with next() (or a loop does). It runs until the next yield, hands the value over, and freezes, with all its local variables intact, until it is asked again. Watch the order of the prints:

def noisy():
    print("start")
    yield 1
    print("middle")
    yield 2
    print("end")

g = noisy()
print("created")
print(next(g))
print(next(g))
try:
    next(g)
except StopIteration:
    print("finished")

When the function runs off the end, or reaches return, the generator raises StopIteration for you. That is the sign to a loop to stop.

Lazy means on demand

A list holds all its items in memory at once. A generator makes each value only when it is wanted. That has two big benefits. The first is that a generator can be infinite, because it never has to hold everything:

from itertools import islice

def fibonacci():
    a, b = 0, 1
    while True:                # no end: it stops only when the caller stops asking
        yield a
        a, b = b, a + b

print(list(islice(fibonacci(), 12)))      # take just the first twelve

The second is that work can be avoided. This program wants the first square above 50. Compare how many squares each version calculates:

work = {"squares": 0}

def squares_list(n):
    result = []
    for i in range(n):
        work["squares"] += 1
        result.append(i * i)
    return result

def squares_gen(n):
    for i in range(n):
        work["squares"] += 1
        yield i * i

def first_over(values, limit):
    for value in values:
        if value > limit:
            return value

first_over(squares_list(100000), 50)
print("list version worked out", work["squares"], "squares")

work["squares"] = 0
first_over(squares_gen(100000), 50)
print("generator version worked out", work["squares"], "squares")

The list version built all 100 000 squares before the search began. The generator stopped as soon as it found one. The same applies to memory: a generator holds one value at a time, however long the sequence.

Generator expressions, yield from and send

A list comprehension builds a whole list. Change the square brackets to round ones and you have a generator expression, which builds nothing until asked:

import sys

squares_list = [n * n for n in range(5)]
squares_gen = (n * n for n in range(5))
print(squares_list)
print(type(squares_gen).__name__)

total = sum(n * n for n in range(1, 100001))       # a hundred thousand squares, never held at once
print(total)

big_list = list(range(100000))
small = (n for n in range(100000))
print("generator is smaller:", sys.getsizeof(small) < sys.getsizeof(big_list))

When a generator expression is the only argument to a function, the extra brackets can go: sum(n * n for n in ...).

Short-circuiting

any and all stop as soon as the answer is known. With a generator they also stop the work:

def check(n):
    print("checking", n)
    return n > 2

print(any(check(n) for n in [1, 2, 3, 4, 5]))

Items 4 and 5 were never even checked.

yield from

A generator can hand over to another generator with yield from, which passes on every value it produces. It makes recursive generators short. This one flattens lists nested to any depth:

def flatten(items):
    for item in items:
        if isinstance(item, list):
            yield from flatten(item)       # delegate to the inner generator
        else:
            yield item

print(list(flatten([1, [2, [3, 4]], 5])))

send: values going in as well

yield can also receive a value. The caller uses send(), which resumes the generator and makes the yield expression equal to whatever was sent. A generator used like this is called a coroutine. This one keeps a running average:

def running_average():
    total, count, average = 0, 0, None
    while True:
        value = yield average         # hand back the average, wait for the next value
        total += value
        count += 1
        average = total / count

avg = running_average()
next(avg)                              # run to the first yield
print(avg.send(10))
print(avg.send(20))
print(avg.send(60))

Coroutines like this are the idea behind the tasks in the scheduling topic, and behind async and await.

itertools: building blocks

The itertools module is a toolbox of fast, lazy building blocks for working with iterators. Each one takes iterables and returns an iterator, so they can be chained.

from itertools import count, cycle, islice, chain, accumulate, groupby
from itertools import product, permutations, combinations, pairwise, takewhile, zip_longest

print(list(islice(count(10, 5), 4)))
print(list(islice(cycle("AB"), 5)))
print(list(chain([1, 2], "ab", (3,))))
print(list(accumulate([1, 2, 3, 4])))
print([(letter, len(list(group))) for letter, group in groupby("aaabbc")])
print(len(list(product("ab", repeat=3))), len(list(permutations("abcd", 2))), len(list(combinations(range(5), 2))))
print(list(pairwise([1, 4, 9, 16])))
print(list(takewhile(lambda n: n < 20, (n * n for n in count(1)))))
print(list(zip_longest("ab", [1, 2, 3], fillvalue="-")))

Notice what groupby did there: ("a", 3), ("b", 2), ("c", 1) is run length encoding, from the compression topic, in one line. It only groups neighbours, so sort first if you want every equal item together.

product is the lazy form of a nested loop, and counting is a quick way to size a search. A four digit PIN has len(list(product(range(10), repeat=4))) possibilities:

from itertools import product, islice

pins = product(range(10), repeat=4)
print(len(list(product(range(10), repeat=4))), "possible four digit PINs")
print(list(islice(pins, 3)))

Pipelines: processing data of any size

Generators fit together like the stages of a factory line. Each takes an iterable and yields its output, and the next stage pulls from it. Items flow through one at a time. A pipeline can process a file of a hundred gigabytes using almost no memory, because only the item in each stage exists at any moment.

This pipeline reads log lines, turns each into a record, and keeps only the errors. The print in the first stage shows when each line is read:

log = '''2025-03-01 INFO login ok user=amy
2025-03-01 ERROR disk full path=/var
2025-03-01 INFO upload ok user=ben size=40
2025-03-02 ERROR timeout host=db1
2025-03-02 WARN slow query ms=900
2025-03-02 ERROR timeout host=db2'''

def read_lines(text):
    for line in text.split("\n"):
        print("read:   ", line[:24])
        yield line

def parse(lines):
    for line in lines:
        date, level, message = line.split(" ", 2)
        yield {"date": date, "level": level, "message": message}

def only(level, records):
    for record in records:
        if record["level"] == level:
            yield record

pipeline = only("ERROR", parse(read_lines(log)))
print("pipeline built, nothing has run yet")
for record in pipeline:
    print("ERROR:  ", record["message"])

Nothing happens when the pipeline is built. Then, as the loop asks for each error, a line is read, parsed and tested, and the process repeats. The reads and the results are interleaved. A normal function-by-function version would read, and store, every line first.

Because each stage does one small job, each stage can be tested and reused on its own. The last stage can be any consumer: a loop, sum, max, list, or a Counter:

from collections import Counter

log = '''2025-03-01 INFO login ok user=amy
2025-03-01 ERROR disk full path=/var
2025-03-01 INFO upload ok user=ben size=40
2025-03-02 ERROR timeout host=db1
2025-03-02 WARN slow query ms=900
2025-03-02 ERROR timeout host=db2'''

def parse(lines):
    for line in lines:
        date, level, message = line.split(" ", 2)
        yield {"date": date, "level": level, "message": message}

records = parse(log.split("\n"))
by_level = Counter(record["level"] for record in records)
print(by_level)
print("errors:", by_level["ERROR"], "| warnings:", by_level["WARN"])

Files are already lazy

When you write for line in open("big.log"), Python does not read the whole file. A file object is an iterator that reads one line at a time. A pipeline that starts with a file keeps that benefit all the way down.

Key terms you must be able to define

Iterable

An object that can be looped over, because it can produce an iterator when iter() is called on it.

Iterator

An object that returns the next item from a sequence each time next() is called, and raises StopIteration when there are no more.

StopIteration

The exception that signals that an iterator has no more items. A for loop catches it and ends.

Generator

A function containing yield, which returns an iterator that produces its values one at a time and pauses between them.

yield

Hands a value back to the caller and pauses the generator, keeping its state, until the next value is requested.

Lazy evaluation

Working out values only when they are needed, instead of all at once.

Generator expression

A compact form like a list comprehension, written in round brackets, which produces values on demand.

yield from

Delegates to another iterator, passing each of its values on as if they came from this generator.

Coroutine

A generator that can also receive values, using send(), and so can be paused and resumed in both directions.

Pipeline

A chain of generators in which each stage takes values from the one before and passes results on one at a time.

Flashcards

Say the definition out loud before you flip the card. Retrieving it from memory is what builds recall, reading it again is not.

Reach the top grade on this topic

Generator questions ask what is produced, and in what order, and why it is worth doing.

• Say that calling a generator function does not run its body: it creates a generator object
• For tracing, write the output in the order things happen: each next runs to the next yield
• Give two benefits of laziness: low memory use, and no wasted work when the loop stops early
• Say a generator can be used only once, and cannot be indexed or measured with len
• For a pipeline, say each item passes through every stage before the next one is read

Mistakes that cost marks here

def mystery(n):
    for i in range(n):
        print("making", i)
        yield i * 2

g = mystery(3)
print("start")
print(next(g))
print(next(g))

def evens(limit):
    for n in range(0, limit, 2):
        yield n

print(list(evens(9)))