13. Iterators and generators: lazy pipelines

13.1 The iterator protocol

for x in obj works on anything iterable. Under the hood:

nums = [10, 20]
it = iter(nums)        # calls nums.__iter__() -> an iterator
print(next(it))        # 10   (calls it.__next__())
print(next(it))        # 20
next(it)               # StopIteration: the for loop catches this and stops

An iterable has __iter__ (lists, dicts, files, strings). An iterator also has __next__ and remembers its position. Iterators are single-use: once exhausted, they stay exhausted.

squares = (n * n for n in range(3))    # a generator expression is an iterator
print(list(squares))                   # [0, 1, 4]
print(list(squares))                   # []  <- already consumed

That second empty list is a common bug when a generator is passed to two consumers. If you need the values twice, store them in a list.

13.2 Generators: functions that pause

A function containing yield returns a generator. Each next() runs the body until the next yield, hands out the value, and freezes the function's local state until the next call:

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

g = countdown(2)       # nothing runs yet
print(next(g))         # start, then 2
print(next(g))         # 1
print(list(g))         # done, then []

Generators give you lazy evaluation: values are produced one at a time, on demand. That means constant memory for unbounded or huge inputs.

13.3 Real-world example: a streaming log pipeline

Chain generators like Unix pipes. Each stage pulls one item at a time from the previous one, so a 50 GB log flows through in a few kilobytes of memory:

import gzip
import json
from collections.abc import Iterable, Iterator
from pathlib import Path

def read_lines(paths: Iterable[Path]) -> Iterator[str]:
    for p in paths:
        opener = gzip.open if p.suffix == ".gz" else open
        with opener(p, "rt", encoding="utf-8", errors="replace") as f:
            yield from f                       # delegate to the file iterator

def parse_json(lines: Iterable[str]) -> Iterator[dict]:
    for line in lines:
        try:
            yield json.loads(line)
        except json.JSONDecodeError:
            continue

def only_errors(events: Iterable[dict]) -> Iterator[dict]:
    return (e for e in events if e.get("level") == "error")

# Build a sample input, then run the pipeline.
Path("app.jsonl").write_text(
    '{"level": "info", "msg": "up"}\n{"level": "error", "msg": "db timeout"}\nnot json\n'
)
for e in only_errors(parse_json(read_lines([Path("app.jsonl")]))):
    print(e["msg"])                            # db timeout

This is the same pattern as cat *.gz | zcat | jq | grep in OS Ch 28, and the same pattern PyTorch's IterableDataset and Hugging Face's streaming datasets use to feed terabytes of training text (Ch 60).

13.4 itertools and friends

import itertools as it

print(list(it.islice(it.count(10, 5), 4)))        # [10, 15, 20, 25]: slice an infinite iterator
print(list(it.chain([1, 2], (3,), "ab")))          # [1, 2, 3, 'a', 'b']
print(list(it.batched(range(7), 3)))               # [(0, 1, 2), (3, 4, 5), (6,)]  (3.12+)
print(list(it.pairwise([1, 4, 9, 16])))            # [(1, 4), (4, 9), (9, 16)]
print(list(it.accumulate([3, 1, 4], initial=0)))   # [0, 3, 4, 8]: running totals
for key, group in it.groupby(sorted(["ant", "bee", "asp"]), key=lambda w: w[0]):
    print(key, list(group))                        # groupby needs sorted input
print(list(it.product("ab", [0, 1])))              # cartesian product: grid search
print(list(it.combinations([1, 2, 3], 2)))
print(any(x > 2 for x in [1, 3]), all([]), max([], default=0))

itertools.batched is what you use to send embeddings or database rows in groups of 64. product is a hyperparameter grid.

13.5 Generators that receive values, and cleanup

yield is an expression: gen.send(value) resumes the generator and makes yield return value. This two-way channel is how coroutines started in Python, and why async def coroutines (Ch 29) are built on the same machinery. You will rarely write send yourself, but you will rely on the cleanup guarantee: when a generator is closed or garbage-collected, GeneratorExit is raised at the paused yield, so a finally or with inside the generator runs.

def reader():
    print("open")
    try:
        yield 1
        yield 2
    finally:
        print("close")    # runs even if the consumer stops early

for x in reader():
    print(x)
    break                 # open, 1, close