Emmanuel Akinfulubi
Free courses Python, from zero

Iteration all 3 written

The protocol behind every for loop, generators that produce values lazily, and the itertools that do it correctly for you.

33 The iteration protocol, and what for really does

A for loop is not a language primitive. It is three method calls in a loop, and knowing which three explains why some things can be iterated twice and others cannot — which is the difference that causes the most confusion in everything built on top.

Two words that are not synonyms

An iterable is something you can get an iterator from: it has __iter__. An iterator is the thing that actually produces values one at a time: it has __next__. A list is iterable; the thing iter(list) hands back is an iterator.

iterable, iterator, and the test that tells them apart

ran at build time 
numbers = [10, 20, 30]

it = iter(numbers)                 # calls numbers.__iter__()
print("the list:    ", type(numbers).__name__)
print("its iterator:", type(it).__name__)

print()
print("next:", next(it))
print("next:", next(it))
print("next:", next(it))

try:
    next(it)
except StopIteration:
    print("StopIteration -- the iterator is spent")

# The defining difference: an iterator is its own iterator, a list is not.
print()
print("iter(it) is it        :", iter(it) is it)
print("iter(numbers) is numbers:", iter(numbers) is numbers)

test("iter() returns a different object for a list", iter(numbers) is numbers, False)
test("but an iterator returns itself", iter(it) is it, True)
test("a list has __iter__ but no __next__",
     (hasattr(numbers, "__iter__"), hasattr(numbers, "__next__")), (True, False))
test("an iterator has both", (hasattr(it, "__iter__"), hasattr(it, "__next__")), (True, True))

What it printed, when the build ran it

the list:     list
its iterator: list_iterator

next: 10
next: 20
next: 30
StopIteration -- the iterator is spent

iter(it) is it        : True
iter(numbers) is numbers: False
ok   iter() returns a different object for a list
ok   but an iterator returns itself
ok   a list has __iter__ but no __next__
ok   an iterator has both

A for loop, written out by hand

iter, next, StopIteration — that is the whole protocol

ran at build time 
numbers = [1, 2, 3]

# What the for loop does, spelled out.
collected = []
iterator = iter(numbers)           # 1. get an iterator
while True:
    try:
        value = next(iterator)     # 2. ask for the next value
    except StopIteration:
        break                      # 3. stop when it says it is done
    collected.append(value)

print("by hand:", collected)

# The same thing, written the way anyone would.
print("with for:", [value for value in numbers])

test("the manual version produces the same values", collected, [1, 2, 3])
test("which is exactly what for does", collected, list(numbers))

What it printed, when the build ran it

by hand: [1, 2, 3]
with for: [1, 2, 3]
ok   the manual version produces the same values
ok   which is exactly what for does

This is why StopIteration is in the traceback catalogue as "not a bug". Every for loop you have ever written raised and caught one on its final iteration. You only see it when you call next() yourself and there is no loop to catch it for you.

Why a list can be looped twice and an iterator cannot

The consequence of the distinction

ran at build time 
numbers = [1, 2, 3]

# A list: each loop asks for a FRESH iterator, so both passes see everything.
first = [n for n in numbers]
second = [n for n in numbers]
print("list, pass 1:", first)
print("list, pass 2:", second)

# An iterator: there is only one, and it is consumed.
it = iter(numbers)
first_pass = [n for n in it]
second_pass = [n for n in it]
print()
print("iterator, pass 1:", first_pass)
print("iterator, pass 2:", second_pass, "<- empty, and no error")

test("a list can be iterated repeatedly", first, second)
test("an iterator gives everything on the first pass", first_pass, [1, 2, 3])
test("and nothing on the second", second_pass, [])

What it printed, when the build ran it

list, pass 1: [1, 2, 3]
list, pass 2: [1, 2, 3]

iterator, pass 1: [1, 2, 3]
iterator, pass 2: [] <- empty, and no error
ok   a list can be iterated repeatedly
ok   an iterator gives everything on the first pass
ok   and nothing on the second

The second pass returns empty rather than raising, which is what makes this dangerous. A function handed an exhausted iterator sees an empty sequence and reports zero results — a plausible-looking wrong answer. If a value might be an iterator and you need it twice, call list() on it once and pay the memory deliberately.

A function that consumes its argument twice

ran at build time 
def summarise(values):
    """Looks harmless. Consumes 'values' twice."""
    count = sum(1 for _ in values)
    total = sum(values)
    return count, total

# With a list: correct.
print("given a list:    ", summarise([1, 2, 3]))

# With an iterator: the first sum consumed it, so the second sees nothing.
print("given an iterator:", summarise(iter([1, 2, 3])), "<- total is 0")

# The fix: materialise once, at the boundary.
def summarise_safely(values):
    values = list(values)
    return sum(1 for _ in values), sum(values)

print("with the fix:     ", summarise_safely(iter([1, 2, 3])))

test("a list gives the right answer", summarise([1, 2, 3]), (3, 6))
test("an iterator silently gives a wrong one", summarise(iter([1, 2, 3])), (3, 0))
test("materialising fixes it", summarise_safely(iter([1, 2, 3])), (3, 6))

What it printed, when the build ran it

given a list:     (3, 6)
given an iterator: (3, 0) <- total is 0
with the fix:      (3, 6)
ok   a list gives the right answer
ok   an iterator silently gives a wrong one
ok   materialising fixes it

Making your own class iterable

Implementing __iter__ is enough. The simplest correct approach is to return an iterator over something that already is one — which is what lesson 29's Playlist did.

One method, and the whole ecosystem works

ran at build time 
class Countdown:
    def __init__(self, start):
        self.start = start

    def __iter__(self):
        # A fresh iterator every time, so the object can be looped repeatedly.
        return iter(range(self.start, 0, -1))

c = Countdown(3)
print("pass 1:", list(c))
print("pass 2:", list(c), "<- still works")

# And everything written against the protocol now works.
print()
print("sum:   ", sum(c))
print("max:   ", max(c))
print("sorted:", sorted(c))
print("in:    ", 2 in c)
print("unpack:", [*c])

test("it can be iterated more than once", list(c), list(c))
test("sum works", sum(Countdown(3)), 6)
test("membership works", 2 in Countdown(3), True)
test("and unpacking works", [*Countdown(2)], [2, 1])

What it printed, when the build ran it

pass 1: [3, 2, 1]
pass 2: [3, 2, 1] <- still works

sum:    6
max:    3
sorted: [1, 2, 3]
in:     True
unpack: [3, 2, 1]
ok   it can be iterated more than once
ok   sum works
ok   membership works
ok   and unpacking works

The trap: __iter__ returning self

Writing the iterator state directly on the object and returning self compiles, passes a quick test, and is wrong: the object is now an iterator, so it can only be consumed once.

An object that is its own iterator can only be used once

ran at build time 
class BrokenCountdown:
    def __init__(self, start):
        self.current = start

    def __iter__(self):
        return self                 # the object IS the iterator -- one shot only

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

b = BrokenCountdown(3)
print("pass 1:", list(b))
print("pass 2:", list(b), "<- silently empty")

# Worse, because it breaks in ways that look unrelated.
b2 = BrokenCountdown(3)
print()
print("max() consumed it:", max(b2))
print("so sum() now sees:", sum(b2))

test("the first pass works", list(BrokenCountdown(3)), [3, 2, 1])
test("the second is empty", (lambda c: (list(c), list(c))[1])(BrokenCountdown(3)), [])
test("and any two consumers interfere", sum(b2), 0)

What it printed, when the build ran it

pass 1: [3, 2, 1]
pass 2: [] <- silently empty

max() consumed it: 3
so sum() now sees: 0
ok   the first pass works
ok   the second is empty
ok   and any two consumers interfere

The rule. If your object represents a collection, make __iter__ return a new iterator each time and do not give the object a __next__. Only return self when the object genuinely represents a single-use stream — a file, a network response, a cursor — where being consumable once is the honest description.

Collection and iterator as separate objects

ran at build time 
# When a separate iterator class IS the right answer: state per iteration.
class Fibonacci:
    """A repeatable iterable over the first n Fibonacci numbers."""

    def __init__(self, count):
        self.count = count

    def __iter__(self):
        return FibonacciIterator(self.count)

class FibonacciIterator:
    def __init__(self, count):
        self.remaining = count
        self.a, self.b = 0, 1

    def __iter__(self):
        return self                 # correct here: this IS the iterator

    def __next__(self):
        if self.remaining <= 0:
            raise StopIteration
        self.remaining -= 1
        value, self.a, self.b = self.a, self.b, self.a + self.b
        return value

fib = Fibonacci(8)
print("pass 1:", list(fib))
print("pass 2:", list(fib), "<- repeatable, because __iter__ built a new iterator")

test("the sequence is right", list(Fibonacci(8)), [0, 1, 1, 2, 3, 5, 8, 13])
test("and it is repeatable", list(fib), list(fib))

What it printed, when the build ran it

pass 1: [0, 1, 1, 2, 3, 5, 8, 13]
pass 2: [0, 1, 1, 2, 3, 5, 8, 13] <- repeatable, because __iter__ built a new iterator
ok   the sequence is right
ok   and it is repeatable

The two-argument iter, and next with a default

next(it, default) and iter(callable, sentinel)

ran at build time 
# next() with a default returns it instead of raising.
it = iter([1])
print("first:", next(it, "nothing left"))
print("again:", next(it, "nothing left"))

# iter(callable, sentinel) calls the function until it returns the sentinel.
# This is how you read a stream in chunks without a while-True-and-break.
chunks = []
data = iter([b"abc", b"def", b"", b"ghi"])

def read_chunk():
    return next(data)

for chunk in iter(read_chunk, b""):      # stops at the empty bytes
    chunks.append(chunk)

print()
print("chunks read before the sentinel:", chunks)

test("next with a default does not raise", next(iter([]), "empty"), "empty")
test("the sentinel form stops at the marker", chunks, [b"abc", b"def"])

What it printed, when the build ran it

first: 1
again: nothing left

chunks read before the sentinel: [b'abc', b'def']
ok   next with a default does not raise
ok   the sentinel form stops at the marker

And the error you get when something is not iterable at all, which names the missing method rather than being vague about it.

An integer is not iterable

raises TypeError 
total = 0
for digit in 12345:
    total += digit

The traceback it really produced, when the build ran it

Traceback (most recent call last):
  File "example.py", line 2, in <module>
    for digit in 12345:
                 ^^^^^
TypeError: 'int' object is not iterable

What to take from this lesson

  • An iterable has __iter__; an iterator has __next__ and returns itself from __iter__.
  • A for loop is iter, then next until StopIteration.
  • Iterators are consumed once, and a second pass gives an empty result rather than an error.
  • Make __iter__ return a fresh iterator for a collection; return self only for a genuine one-shot stream.
  • next(it, default) avoids the exception; iter(callable, sentinel) loops until a marker value.

34 Generators, yield, and laziness that pays

A generator is a function that produces values one at a time instead of returning once. Writing one is a matter of using yield instead of return — and what you get back is an iterator, with all the properties lesson 33 established.

A function with yield is a different kind of function

Calling it runs nothing; next() runs it up to the next yield

ran at build time 
def counted():
    print("   -- body starts")
    yield 1
    print("   -- resumed after the first yield")
    yield 2
    print("   -- resumed after the second")

# Calling it runs NOTHING. It builds a generator object.
gen = counted()
print("after the call:", type(gen).__name__)

print()
print("first next():")
print("   got", next(gen))
print("second next():")
print("   got", next(gen))

print()
try:
    next(gen)
except StopIteration:
    print("StopIteration -- the function reached its end")

test("calling a generator function returns a generator", type(counted()).__name__, "generator")
test("and it is an iterator", hasattr(counted(), "__next__"), True)
test("collecting it gives the yielded values", list(counted()), [1, 2])

What it printed, when the build ran it

after the call: generator

first next():
   -- body starts
   got 1
second next():
   -- resumed after the first yield
   got 2

   -- resumed after the second
StopIteration -- the function reached its end
ok   calling a generator function returns a generator
ok   and it is an iterator
   -- body starts
   -- resumed after the first yield
   -- resumed after the second
ok   collecting it gives the yielded values

That is the whole mechanism. yield hands a value out and suspends the function, keeping its local variables and its position. The next next() resumes exactly there. A generator is a function you can pause.

Memory: the practical argument

A million values, or a few hundred bytes

ran at build time 
import sys

def squares_generator(n):
    for i in range(n):
        yield i * i

def squares_list(n):
    return [i * i for i in range(n)]

gen = squares_generator(1_000_000)
lst = squares_list(1_000_000)

print("generator object:", f"{sys.getsizeof(gen):>9,} bytes")
print("list of a million:", f"{sys.getsizeof(lst):>9,} bytes")
print()
print("Both describe the same million values. One is storing them.")

# And the generator can be consumed exactly like the list.
print("first five from each:", list(squares_generator(5)), squares_list(5))

test("the generator is orders of magnitude smaller", sys.getsizeof(gen) * 1000 < sys.getsizeof(lst), True)
test("and yields the same values", list(squares_generator(5)), squares_list(5))

What it printed, when the build ran it

generator object:       208 bytes
list of a million: 8,448,728 bytes

Both describe the same million values. One is storing them.
first five from each: [0, 1, 4, 9, 16] [0, 1, 4, 9, 16]
ok   the generator is orders of magnitude smaller
ok   and yields the same values

Infinite sequences, which a list cannot represent at all

An unbounded sequence, bounded at the point of use

ran at build time 
from itertools import islice

def naturals():
    n = 0
    while True:            # no termination -- and that is fine
        yield n
        n += 1

def fibonacci():
    a, b = 0, 1
    while True:
        yield a
        a, b = b, a + b

print("first 8 naturals: ", list(islice(naturals(), 8)))
print("first 8 fibonacci:", list(islice(fibonacci(), 8)))

# Taking from an infinite generator until a condition holds.
def first_above(threshold):
    for value in fibonacci():
        if value > threshold:
            return value

print()
print("first fibonacci above 1000:", first_above(1000))

test("islice bounds an unbounded generator", list(islice(naturals(), 3)), [0, 1, 2])
test("fibonacci is correct", list(islice(fibonacci(), 8)), [0, 1, 1, 2, 3, 5, 8, 13])
test("and searching an infinite sequence terminates", first_above(1000), 1597)

What it printed, when the build ran it

first 8 naturals:  [0, 1, 2, 3, 4, 5, 6, 7]
first 8 fibonacci: [0, 1, 1, 2, 3, 5, 8, 13]

first fibonacci above 1000: 1597
ok   islice bounds an unbounded generator
ok   fibonacci is correct
ok   and searching an infinite sequence terminates

Never call list() on an unbounded generator. It will consume memory until the process is killed. Bound it first with islice, a break, or a condition — the laziness is what makes the infinite sequence possible, and materialising it throws that away.

Laziness is visible: work happens only when asked

Nothing is computed until something consumes it

ran at build time 
log = []

def tracked(values):
    for v in values:
        log.append(f"produced {v}")
        yield v

# Building the generator does no work.
gen = tracked([1, 2, 3])
log_after_creation = list(log)
print("after creating it, log is:", log_after_creation)

# Each value is produced on demand.
print("asking for one:", next(gen))
print("log is now:", log)

# any() stops as soon as it has an answer, so the third is never produced.
log.clear()
found = any(v == 2 for v in tracked([1, 2, 3]))
print()
print("any(v == 2):", found)
print("log shows it stopped early:", log)

test("creating a generator produced nothing at all", log_after_creation, [])
test("any() short-circuits, so the third value was never produced", log, ["produced 1", "produced 2"])

What it printed, when the build ran it

after creating it, log is: []
asking for one: 1
log is now: ['produced 1']

any(v == 2): True
log shows it stopped early: ['produced 1', 'produced 2']
ok   creating a generator produced nothing at all
ok   any() short-circuits, so the third value was never produced

Pipelines: the pattern generators are for

Each stage takes an iterable and yields a transformed one. Nothing is materialised in between, so a multi-stage pipeline over a huge file uses the memory of a single row.

Four stages, one row in memory at a time

ran at build time 
from pathlib import Path

Path("sales.csv").write_text(
    "region,amount\n"
    "north,100\n"
    "\n"                       # a blank line, as real files have
    "south,250\n"
    "# a comment line\n"
    "east,invalid\n"
    "west,400\n",
    encoding="utf-8",
)

def read_lines(path):
    with open(path, encoding="utf-8") as f:
        for line in f:
            yield line.rstrip("\n")

def drop_blanks_and_comments(lines):
    for line in lines:
        if line and not line.startswith("#"):
            yield line

def skip_header(lines):
    iterator = iter(lines)
    next(iterator, None)          # discard the first
    yield from iterator

def parse(lines):
    for line in lines:
        region, _, amount = line.partition(",")
        try:
            yield region, int(amount)
        except ValueError:
            continue              # skip unparseable rows

# The pipeline is assembled but nothing has been read yet.
pipeline = parse(skip_header(drop_blanks_and_comments(read_lines("sales.csv"))))
print("assembled, nothing read yet:", type(pipeline).__name__)

rows = list(pipeline)
print("rows:", rows)
print("total:", sum(amount for _, amount in rows))

test("the pipeline is a generator until consumed", type(pipeline).__name__, "generator")
test("blank lines, comments, the header and bad rows were all dropped",
     [r for r, _ in rows], ["north", "south", "west"])
test("and the amounts parsed", sum(a for _, a in rows), 750)

What it printed, when the build ran it

assembled, nothing read yet: generator
rows: [('north', 100), ('south', 250), ('west', 400)]
total: 750
ok   the pipeline is a generator until consumed
ok   blank lines, comments, the header and bad rows were all dropped
ok   and the amounts parsed

yield from: delegating to another iterable

yield from, and the return value it forwards

ran at build time 
def inner():
    yield "a"
    yield "b"
    return "inner finished"        # a generator CAN return a value

def manual():
    for value in inner():          # the long way
        yield value

def delegated():
    result = yield from inner()    # shorter, and captures the return value
    yield f"inner said {result!r}"

print("manual:   ", list(manual()))
print("delegated:", list(delegated()))

# Flattening a nested structure, which is where yield from earns its keep.
def flatten(nested):
    for item in nested:
        if isinstance(item, list):
            yield from flatten(item)     # recurse
        else:
            yield item

print()
print("flattened:", list(flatten([1, [2, [3, [4, 5]], 6], 7])))

test("yield from produces the same values as the loop",
     list(manual()), ["a", "b"])
test("and additionally captures the return value",
     list(delegated())[-1], "inner said 'inner finished'")
test("recursive flattening works", list(flatten([1, [2, [3]], 4])), [1, 2, 3, 4])

What it printed, when the build ran it

manual:    ['a', 'b']
delegated: ['a', 'b', "inner said 'inner finished'"]

flattened: [1, 2, 3, 4, 5, 6, 7]
ok   yield from produces the same values as the loop
ok   and additionally captures the return value
ok   recursive flattening works

A generator's return value does not appear among the yielded items. It is attached to the StopIteration, which is where yield from reads it from.

Where a generator's return value goes

ran at build time 
def with_return():
    yield 1
    yield 2
    return "the summary"

# Iterating never shows the return value.
print("list():", list(with_return()))

# It is on the StopIteration.
gen = with_return()
next(gen)
next(gen)
captured = None
try:
    next(gen)
except StopIteration as exc:
    captured = exc.value
    print("StopIteration.value:", repr(captured))

test("the return value is not yielded", list(with_return()), [1, 2])
test("it travels on the StopIteration instead", captured, "the summary")

What it printed, when the build ran it

list(): [1, 2]
StopIteration.value: 'the summary'
ok   the return value is not yielded
ok   it travels on the StopIteration instead

Generators are one-shot, which follows from lesson 33

The same one-shot behaviour, in the place you will meet it

ran at build time 
def three():
    yield from [1, 2, 3]

gen = three()
print("pass 1:", list(gen))
print("pass 2:", list(gen), "<- exhausted, as any iterator would be")

# Which means this common mistake gives a wrong answer, not an error.
gen = three()
print()
print("max:", max(gen))
print("sum of the same generator:", sum(gen), "<- nothing left")

# If you need it twice, call the function again -- it is cheap.
print()
print("calling it again:", max(three()), sum(three()))

test("a generator is exhausted after one pass",
     (lambda g: (list(g), list(g))[1])(three()), [])
test("so two consumers interfere", sum(gen), 0)
test("calling the function again is the fix", (max(three()), sum(three())), (3, 6))

What it printed, when the build ran it

pass 1: [1, 2, 3]
pass 2: [] <- exhausted, as any iterator would be

max: 3
sum of the same generator: 0 <- nothing left

calling it again: 3 6
ok   a generator is exhausted after one pass
ok   so two consumers interfere
ok   calling the function again is the fix

Return the generator function, not the generator. If a caller may need the sequence more than once, hand them something they can iterate repeatedly — a list, or a class with __iter__ as in lesson 33. A function returning a bare generator is an interface that only works once, and nothing in its signature says so.

Generators and with: cleanup is not guaranteed

A generator holding a resource open across a yield only closes it if the generator is driven to completion or explicitly closed. Abandon it half-way and the cleanup waits for garbage collection.

finally in a generator runs on completion or close()

ran at build time 
from pathlib import Path

Path("data.txt").write_text("a\nb\nc\n", encoding="utf-8")

state = {"closed": None}

def read_all(path):
    f = open(path, encoding="utf-8")
    try:
        for line in f:
            yield line.strip()
    finally:
        state["closed"] = True

# Consumed fully: the finally runs.
state["closed"] = False
lines = list(read_all("data.txt"))
closed_after_full = state["closed"]
print("fully consumed:", lines, "closed:", closed_after_full)

# Abandoned after one value: the finally has NOT run yet.
state["closed"] = False
gen = read_all("data.txt")
first = next(gen)
closed_while_open = state["closed"]
print("one value:", first, " closed:", closed_while_open, "<- still open")

# close() forces it.
gen.close()
closed_after_close = state["closed"]
print("after close():", closed_after_close)

test("full consumption runs the finally", closed_after_full, True)
test("abandoning it half-way does not", closed_while_open, False)
test("and close() runs it", closed_after_close, True)

What it printed, when the build ran it

fully consumed: ['a', 'b', 'c'] closed: True
one value: a  closed: False <- still open
after close(): True
ok   full consumption runs the finally
ok   abandoning it half-way does not
ok   and close() runs it

send, throw and close

A generator is bidirectional: yield is an expression, and send supplies its value. This is how coroutines were built before async/await existed, and it is worth recognising rather than reaching for.

send, throw and close

ran at build time 
def accumulator():
    total = 0
    while True:
        try:
            received = yield total        # yield out, and receive in
        except ValueError:
            print("   (told to reset)")
            total = 0
            continue
        total += received

acc = accumulator()
primed = next(acc)                           # run up to the first yield
print("priming with next():", primed)
after_ten = acc.send(10)
print("send 10 ->", after_ten)
after_five = acc.send(5)
print("send  5 ->", after_five)

# throw() raises inside the generator, at the yield.
print("throwing ValueError:")
after_reset = acc.throw(ValueError)
print("after reset ->", after_reset)

after_seven = acc.send(7)
print("send  7 ->", after_seven)
acc.close()

test("priming yields the initial total", primed, 0)
test("each send adds to the running total", (after_ten, after_five), (10, 15))
test("throw() reached the except inside the generator and reset it", after_reset, 0)
test("and it kept accumulating afterwards", after_seven, 7)

What it printed, when the build ran it

priming with next(): 0
send 10 -> 10
send  5 -> 15
throwing ValueError:
   (told to reset)
after reset -> 0
send  7 -> 7
ok   priming yields the initial total
ok   each send adds to the running total
ok   throw() reached the except inside the generator and reset it
ok   and it kept accumulating afterwards

A generator must be primed before it can be sent a value, because there is no suspended yield to receive it yet.

Sending to a generator that has not started

raises TypeError 
def accumulator():
    total = 0
    while True:
        total += yield total

acc = accumulator()
acc.send(10)

The traceback it really produced, when the build ran it

Traceback (most recent call last):
  File "example.py", line 7, in <module>
    acc.send(10)
    ~~~~~~~~^^^^
TypeError: can't send non-None value to a just-started generator

Most code should not use send. It makes a generator into a state machine whose behaviour depends on the order of calls, which is hard to read and harder to test. If you want concurrency, lesson 43's async/await is built for it; if you want accumulating state, a class with a method is clearer.

Generator expressions, once more

Two ways to write the same generator

ran at build time 
import sys

# A generator expression is a generator function without the def.
squares_expr = (n * n for n in range(5))

def squares_func():
    for n in range(5):
        yield n * n

print("expression:", list(squares_expr))
print("function:  ", list(squares_func()))
print("same type: ", type((n for n in [])).__name__ == type(squares_func()).__name__)

# When the generator is the only argument, the brackets can be dropped.
print()
print("sum of squares:", sum(n * n for n in range(5)))
print("any over a file-sized sequence:", any(n > 100 for n in range(1000)))

test("both produce the same values", list(squares_func()), [0, 1, 4, 9, 16])
test("a generator expression is a generator",
     type((n for n in [])).__name__, "generator")
test("and a bare one can be passed as the only argument", sum(n for n in range(4)), 6)

What it printed, when the build ran it

expression: [0, 1, 4, 9, 16]
function:   [0, 1, 4, 9, 16]
same type:  True

sum of squares: 30
any over a file-sized sequence: True
ok   both produce the same values
ok   a generator expression is a generator
ok   and a bare one can be passed as the only argument

Which to write. A generator expression when the transformation fits on one line — the same rule as lesson 12's comprehensions. A generator function when there is setup, multiple yield points, a try/finally, or anything worth a name and a docstring.

What to take from this lesson

  • A function containing yield returns a generator; calling it runs no code.
  • yield suspends the function, keeping its state, and next() resumes it.
  • Generators hold nothing, so they can represent unbounded sequences — never list() one.
  • They are one-shot. Call the function again, or return something repeatable.
  • Chain them into pipelines to process large data a row at a time.
  • yield from delegates and forwards the inner return value; that value arrives on StopIteration, not among the yields.

35 itertools, and the plumbing it writes for you

Almost every loop you write over a sequence has been written before, correctly, in the standard library. itertools is that collection — all lazy, all working on any iterable, and all avoiding a specific hand-rolled loop that is easy to get subtly wrong.

Joining and slicing

chain and islice

ran at build time 
from itertools import chain, islice

first = [1, 2]
second = (3, 4)
third = {5}

# chain treats several iterables as one, without building a combined list.
print("chain:", list(chain(first, second, third)))

# chain.from_iterable flattens one level.
nested = [[1, 2], [3], [4, 5]]
print("from_iterable:", list(chain.from_iterable(nested)))

# islice slices any iterable, including one you cannot index.
gen = (n * n for n in range(100))
print()
print("islice(gen, 5):     ", list(islice(gen, 5)))
print("islice(range, 2, 7):", list(islice(range(100), 2, 7)))
print("islice with a step: ", list(islice(range(20), 0, 20, 5)))

test("chain concatenates different iterable types", list(chain([1], (2,), {3})), [1, 2, 3])
test("from_iterable flattens one level", list(chain.from_iterable([[1, 2], [3]])), [1, 2, 3])
test("islice works on a generator, which cannot be sliced with []",
     list(islice((n for n in range(10)), 3)), [0, 1, 2])
test("and takes start, stop and step", list(islice(range(20), 0, 20, 5)), [0, 5, 10, 15])

What it printed, when the build ran it

chain: [1, 2, 3, 4, 5]
from_iterable: [1, 2, 3, 4, 5]

islice(gen, 5):      [0, 1, 4, 9, 16]
islice(range, 2, 7): [2, 3, 4, 5, 6]
islice with a step:  [0, 5, 10, 15]
ok   chain concatenates different iterable types
ok   from_iterable flattens one level
ok   islice works on a generator, which cannot be sliced with []
ok   and takes start, stop and step

Running totals and pairs

accumulate and pairwise

ran at build time 
from itertools import accumulate, pairwise
import operator

readings = [3, 1, 4, 1, 5]

# accumulate: a running total, or any running application of a function.
print("running sum: ", list(accumulate(readings)))
print("running max: ", list(accumulate(readings, max)))
print("running product:", list(accumulate(readings, operator.mul)))
print("with an initial value:", list(accumulate(readings, initial=100)))

# pairwise: every adjacent pair, which replaces a zip(x, x[1:]) dance.
print()
print("pairwise:", list(pairwise(readings)))

# Which makes differences between consecutive readings a one-liner.
deltas = [b - a for a, b in pairwise(readings)]
print("deltas:  ", deltas)

test("accumulate gives a running total", list(accumulate([1, 2, 3])), [1, 3, 6])
test("with a function it runs that instead", list(accumulate([3, 1, 4], max)), [3, 3, 4])
test("initial= prepends a starting value", list(accumulate([1, 2], initial=10)), [10, 11, 13])
test("pairwise yields adjacent pairs", list(pairwise([1, 2, 3])), [(1, 2), (2, 3)])
test("so differences are a comprehension", deltas, [-2, 3, -3, 4])

What it printed, when the build ran it

running sum:  [3, 4, 8, 9, 14]
running max:  [3, 3, 4, 4, 5]
running product: [3, 3, 12, 12, 60]
with an initial value: [100, 103, 104, 108, 109, 114]

pairwise: [(3, 1), (1, 4), (4, 1), (1, 5)]
deltas:   [-2, 3, -3, 4]
ok   accumulate gives a running total
ok   with a function it runs that instead
ok   initial= prepends a starting value
ok   pairwise yields adjacent pairs
ok   so differences are a comprehension

Chunking, with batched

batched, for size-limited work

ran at build time 
from itertools import batched

ids = list(range(1, 10))

# batched splits an iterable into tuples of at most n. The last one is short.
for chunk in batched(ids, 4):
    print("  chunk:", chunk)

print()
print("as a list:", [list(c) for c in batched(ids, 4)])

# The real use: sending records to something with a request-size limit.
def send(batch):
    return f"sent {len(batch)} records"

print()
for chunk in batched(range(7), 3):
    print(" ", send(chunk))

rejoined = [x for chunk in batched(range(7), 3) for x in chunk]

test("batched produces tuples", type(next(iter(batched([1, 2], 2)))).__name__, "tuple")
test("the final batch may be short", [len(c) for c in batched(range(7), 3)], [3, 3, 1])
test("and nothing is lost or duplicated", rejoined, list(range(7)))

What it printed, when the build ran it

  chunk: (1, 2, 3, 4)
  chunk: (5, 6, 7, 8)
  chunk: (9,)

as a list: [[1, 2, 3, 4], [5, 6, 7, 8], [9]]

  sent 3 records
  sent 3 records
  sent 1 records
ok   batched produces tuples
ok   the final batch may be short
ok   and nothing is lost or duplicated

groupby, and the requirement everyone misses

groupby groups consecutive items with the same key. It does not gather matching items from across the sequence — so on unsorted input it produces a separate group each time the key changes, which looks like grouping and is not.

The same data, grouped wrongly and then correctly

ran at build time 
from itertools import groupby

sales = [
    ("north", 100),
    ("south", 250),
    ("north", 150),
    ("south", 50),
]

print("UNSORTED input:")
for region, group in groupby(sales, key=lambda row: row[0]):
    print(f"   {region}: {[amount for _, amount in group]}")

print()
print("SORTED by the same key first:")
for region, group in groupby(sorted(sales, key=lambda row: row[0]), key=lambda row: row[0]):
    print(f"   {region}: {[amount for _, amount in group]}")

unsorted_groups = [k for k, _ in groupby(sales, key=lambda r: r[0])]
sorted_groups = [k for k, _ in groupby(sorted(sales), key=lambda r: r[0])]

test("unsorted input produced four groups for two regions", unsorted_groups,
     ["north", "south", "north", "south"])
test("sorting first produces two", sorted_groups, ["north", "south"])

What it printed, when the build ran it

UNSORTED input:
   north: [100]
   south: [250]
   north: [150]
   south: [50]

SORTED by the same key first:
   north: [100, 150]
   south: [250, 50]
ok   unsorted input produced four groups for two regions
ok   sorting first produces two

Always sort by the same key you group by. The unsorted version above did not raise — it returned four groups where a reader expects two, and the totals computed from it would be plausible and wrong. If sorting is too expensive, or the key is unorderable, use a defaultdict(list) instead: one pass, no sort, and no ordering requirement.

groupby versus a defaultdict, for the same job

ran at build time 
from itertools import groupby
from collections import defaultdict

sales = [("north", 100), ("south", 250), ("north", 150), ("south", 50)]

# With groupby: sort, then group.
by_groupby = {
    region: sum(a for _, a in group)
    for region, group in groupby(sorted(sales), key=lambda r: r[0])
}

# With defaultdict: one pass, no sort needed.
totals = defaultdict(int)
for region, amount in sales:
    totals[region] += amount

print("groupby:    ", by_groupby)
print("defaultdict:", dict(totals))

test("both agree", by_groupby, dict(totals))
test("and the totals are right", dict(totals), {"north": 250, "south": 300})

What it printed, when the build ran it

groupby:     {'north': 250, 'south': 300}
defaultdict: {'north': 250, 'south': 300}
ok   both agree
ok   and the totals are right

The other groupby trap: each group is a lazy iterator that shares the underlying sequence, so it is invalidated as soon as you move to the next group. Materialise each group before advancing.

A group is only valid until the next one

ran at build time 
from itertools import groupby

data = sorted([("a", 1), ("b", 2), ("a", 3)])

# Keeping the group objects to use later: they are all empty by then.
kept = [(key, group) for key, group in groupby(data, key=lambda r: r[0])]
print("groups kept for later:", [(k, list(g)) for k, g in kept])

# Materialising inside the loop: correct.
materialised = [(key, list(group)) for key, group in groupby(data, key=lambda r: r[0])]
print("materialised in the loop:", materialised)

test("group iterators are spent once you move on",
     [(k, list(g)) for k, g in kept], [("a", []), ("b", [])])
test("materialising as you go works", materialised, [("a", [("a", 1), ("a", 3)]), ("b", [("b", 2)])])

What it printed, when the build ran it

groups kept for later: [('a', []), ('b', [])]
materialised in the loop: [('a', [('a', 1), ('a', 3)]), ('b', [('b', 2)])]
ok   group iterators are spent once you move on
ok   materialising as you go works

Combinations and products

product, permutations, combinations

ran at build time 
from itertools import product, permutations, combinations, combinations_with_replacement

colours = ["red", "green"]
sizes = ["S", "M"]

# product: every combination across iterables -- a nested loop, flattened.
print("product:", list(product(colours, sizes)))

# Which replaces this:
manual = [(c, s) for c in colours for s in sizes]
print("by hand:", manual)

items = ["a", "b", "c"]
print()
print("permutations(2): ", list(permutations(items, 2)), "<- order matters")
print("combinations(2): ", list(combinations(items, 2)), "<- order does not")
print("with_replacement:", list(combinations_with_replacement(items, 2)))

# product with repeat= is how you enumerate fixed-length sequences.
print()
print("product('01', repeat=3):", ["".join(p) for p in product("01", repeat=3)])

test("product flattens a nested loop", list(product(colours, sizes)), manual)
test("permutations count order", len(list(permutations(items, 2))), 6)
test("combinations do not", len(list(combinations(items, 2))), 3)
test("repeat= enumerates fixed-length sequences", len(list(product("01", repeat=3))), 8)

What it printed, when the build ran it

product: [('red', 'S'), ('red', 'M'), ('green', 'S'), ('green', 'M')]
by hand: [('red', 'S'), ('red', 'M'), ('green', 'S'), ('green', 'M')]

permutations(2):  [('a', 'b'), ('a', 'c'), ('b', 'a'), ('b', 'c'), ('c', 'a'), ('c', 'b')] <- order matters
combinations(2):  [('a', 'b'), ('a', 'c'), ('b', 'c')] <- order does not
with_replacement: [('a', 'a'), ('a', 'b'), ('a', 'c'), ('b', 'b'), ('b', 'c'), ('c', 'c')]

product('01', repeat=3): ['000', '001', '010', '011', '100', '101', '110', '111']
ok   product flattens a nested loop
ok   permutations count order
ok   combinations do not
ok   repeat= enumerates fixed-length sequences

The infinite ones, and how to stop them

count, cycle, repeat, takewhile, dropwhile

ran at build time 
from itertools import count, cycle, repeat, islice, takewhile, dropwhile

# count: an endless counter, with an optional step.
print("count:       ", list(islice(count(10, 5), 4)))

# cycle: repeats an iterable forever -- useful for round-robin assignment.
print("cycle:       ", list(islice(cycle(["a", "b"]), 5)))

# repeat: the same value, n times or forever.
print("repeat:      ", list(repeat("x", 3)))

# takewhile / dropwhile stop or start on a condition rather than a count.
readings = [1, 3, 8, 2, 9]
print()
print("takewhile < 5:", list(takewhile(lambda n: n < 5, readings)), "<- stops at the first failure")
print("dropwhile < 5:", list(dropwhile(lambda n: n < 5, readings)), "<- skips until the first failure")

# A real use for cycle: distributing work round-robin.
workers = cycle(["w1", "w2", "w3"])
assignment = [(task, next(workers)) for task in ["a", "b", "c", "d"]]
print()
print("round-robin:", assignment)

test("count with a step", list(islice(count(10, 5), 3)), [10, 15, 20])
test("cycle repeats", list(islice(cycle([1, 2]), 5)), [1, 2, 1, 2, 1])
test("takewhile stops at the first failure, ignoring later matches",
     list(takewhile(lambda n: n < 5, readings)), [1, 3])
test("dropwhile keeps everything after it", list(dropwhile(lambda n: n < 5, readings)), [8, 2, 9])
test("cycle distributes round-robin", [w for _, w in assignment], ["w1", "w2", "w3", "w1"])

What it printed, when the build ran it

count:        [10, 15, 20, 25]
cycle:        ['a', 'b', 'a', 'b', 'a']
repeat:       ['x', 'x', 'x']

takewhile < 5: [1, 3] <- stops at the first failure
dropwhile < 5: [8, 2, 9] <- skips until the first failure

round-robin: [('a', 'w1'), ('b', 'w2'), ('c', 'w3'), ('d', 'w1')]
ok   count with a step
ok   cycle repeats
ok   takewhile stops at the first failure, ignoring later matches
ok   dropwhile keeps everything after it
ok   cycle distributes round-robin

takewhile is not filter. It stops permanently at the first item that fails the test, so takewhile(n < 5, [1, 3, 8, 2]) gives [1, 3] and drops the 2. filter would keep it. Reach for takewhile when the sequence is ordered and you want a prefix; reach for filter or a comprehension when you want every match.

tee, and the cost it hides

tee, and why it is not free

ran at build time 
from itertools import tee

# tee gives you several independent iterators over one source.
source = iter([1, 2, 3])
a, b = tee(source, 2)

branch_a = list(a)
branch_b = list(b)
print("a:", branch_a)
print("b:", branch_b, "<- independent of a")

# But the source itself is consumed, and tee buffers internally.
leftover = list(source)
print("the original source is now:", leftover)

print()
print("tee has to remember every value one branch has seen and the other")
print("has not, so tee-ing then fully draining one branch buffers the lot.")

test("both branches see everything", (branch_a, branch_b), ([1, 2, 3], [1, 2, 3]))
test("and the original source is left with nothing", leftover, [])

What it printed, when the build ran it

a: [1, 2, 3]
b: [1, 2, 3] <- independent of a
the original source is now: []

tee has to remember every value one branch has seen and the other
has not, so tee-ing then fully draining one branch buffers the lot.
ok   both branches see everything
ok   and the original source is left with nothing

Do not use the original iterator after tee. And if you were going to drain one branch completely before touching the other, list() is simpler and uses the same memory — tee only wins when the branches advance roughly together.

The promise from lesson 6: deque

Not itertools, but this is where it belongs. Lesson 6 noted that list.pop(0) shifts every remaining element. collections.deque is the structure with cheap removal from both ends.

deque: the right structure for a queue

ran at build time 
from collections import deque
import timeit

# The same queue operation, on a list and on a deque.
def with_list(n):
    q = list(range(n))
    while q:
        q.pop(0)

def with_deque(n):
    q = deque(range(n))
    while q:
        q.popleft()

n = 20_000
list_time = timeit.timeit(lambda: with_list(n), number=3)
deque_time = timeit.timeit(lambda: with_deque(n), number=3)

print(f"draining {n:,} items from the front, three times")
print(f"  list.pop(0):    {list_time:.4f}s")
print(f"  deque.popleft(): {deque_time:.4f}s")
print(f"  the deque is about {list_time / deque_time:.0f}x faster here")

# deque also does both ends, and can be bounded.
d = deque([2, 3])
d.appendleft(1)
d.append(4)
print()
print("both ends:", list(d))

recent = deque(maxlen=3)
for value in range(6):
    recent.append(value)
print("maxlen=3 keeps only the last three:", list(recent))

test("the deque is faster for front removal", deque_time < list_time, True)
test("appendleft and append both work", list(d), [1, 2, 3, 4])
test("maxlen discards from the other end", list(recent), [3, 4, 5])

What it printed, when the build ran it

draining 20,000 items from the front, three times
  list.pop(0):    0.0804s
  deque.popleft(): 0.0035s
  the deque is about 23x faster here

both ends: [1, 2, 3, 4]
maxlen=3 keeps only the last three: [3, 4, 5]
ok   the deque is faster for front removal
ok   appendleft and append both work
ok   maxlen discards from the other end

What to take from this lesson

  • chain to treat several iterables as one; islice to slice something you cannot index.
  • accumulate for running totals, pairwise for adjacent pairs, batched for size-limited chunks.
  • groupby groups only consecutive equal keys — sort first, or use a defaultdict. And materialise each group before moving on.
  • product replaces nested loops; combinations and permutations differ on whether order matters.
  • takewhile stops at the first failure; it is not filter.
  • deque, not list, when you remove from the front.

Newsletter

New writing, when it lands.

Occasional notes on shipping software with AI agents. No schedule, no filler, unsubscribe in one click.