Iterators & Generators

Every for loop you have written runs on the same machinery: an object hands out one value at a time until it says stop. Today you build things that do the handing out.

The iterator protocol

What a for loop really does Python

numbers = [10, 20, 30]

iterator = iter(numbers)      # the for loop calls this
print(next(iterator))
print(next(iterator))
print(next(iterator))

try:
    print(next(iterator))     # nothing left to hand out
except StopIteration:
    print("StopIteration - this is where a for loop stops")

for x in thing: means: call iter(thing), then next() repeatedly until StopIteration. Anything that supports those two calls can be looped over.

Generators

Writing that protocol by hand is tedious. A function containing yield becomes a generator and does it for you.

yield in action Python

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

gen = countdown(3)
print(gen)                # nothing has run yet

print(next(gen))          # runs up to the first yield, then pauses
print(next(gen))
print(next(gen))

for value in countdown(3):
    print("  got", value)

The key idea: yield pauses the function and hands a value out. The next next() resumes exactly where it stopped, with all local variables intact.

Why it matters: memory

A list of a million versus a generator Python

import sys

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

def squares_gen(n):
    for i in range(n):
        yield i ** 2

as_list = squares_list(1_000_000)
as_gen = squares_gen(1_000_000)

print("list:     ", sys.getsizeof(as_list), "bytes")
print("generator:", sys.getsizeof(as_gen), "bytes")

print(sum(squares_gen(1_000_000)))

The list holds a million integers at once; the generator holds one. Both produce the same sum. This is why reading a file line by line works on files larger than your memory — file objects are iterators.

Infinite sequences

A generator need never finish, because nothing is computed until asked for.

An endless generator, safely consumed Python

def fibonacci():
    a, b = 0, 1
    while True:
        yield a
        a, b = b, a + b        # advance the pair - both from the OLD values

# Take only what you need:
for index, value in enumerate(fibonacci()):
    if index >= 10:
        break
    print(value, end=" ")
print()

from itertools import islice
print(list(islice(fibonacci(), 10)))   # islice = "slice, but for iterators"

Two lines there are worth slowing down on.

`a, b = b, a + b`

This is the unpacking from Day 3, doing real work. Python evaluates the whole right-hand side first, using the values the names have now, and only then assigns. So b and a + b are both computed from the old pair, and the two names move forward together — no temporary variable, and no danger of the second calculation seeing the first one’s result.

Written as two separate statements it quietly gives the wrong answer:

One statement, not two Python

a, b = 0, 1
a, b = b, a + b
print("together:  ", a, b)     # 1 1 - correct

a, b = 0, 1
a = b                          # a is now 1
b = a + b                      # so this adds 1 + 1, not 0 + 1
print("one by one:", a, b)     # 1 2 - wrong

`islice`

A generator has no length and no index — it only knows how to produce the next value — so the usual slice syntax does not apply to it:

Slicing something that has no positions Python

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

from itertools import islice

print(list(islice(fibonacci(), 5)))        # the first five
print(list(islice(fibonacci(), 5, 10)))    # skip five, then take five

print(fibonacci()[:5])                     # TypeError - no positions to slice

islice(iterable, stop) and islice(iterable, start, stop) read like range, and they pull values one at a time and stop — which is what makes them safe on an endless generator. list(...) around the result is only there to print it; without it you would get another lazy iterator.

Pipelines

Generators chain together, each stage pulling from the one before. Nothing is buffered in between.

A three-stage pipeline Python

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

def drop_blanks(lines):
    for line in lines:
        if line.strip():
            yield line

def parse(lines):
    for line in lines:
        name, _, score = line.partition(",")
        yield name.strip(), int(score)

raw = """ada, 95

alan, 88
grace, 91
"""

for name, score in parse(drop_blanks(read_lines(raw))):
    print(f"{name:6} {score}")

Each function does one thing, and each could be tested on its own by handing it a list of lines.

The interesting part is what happens if the data stops being a string. Only read_lines knows where the lines come from — drop_blanks and parse just take an iterable of lines and do not care what is producing them. So swapping the source means rewriting one function and leaving the other two untouched:

Same pipeline, different source Python

def drop_blanks(lines):        # unchanged
    for line in lines:
        if line.strip():
            yield line

def parse(lines):              # unchanged
    for line in lines:
        name, _, score = line.partition(",")
        yield name.strip(), int(score)

def read_file(path):           # the only new function
    with open(path) as file:
        for line in file:
            yield line.rstrip("\n")

with open("scores.csv", "w") as file:
    file.write("ada, 95\n\nalan, 88\ngrace, 91\n")

for name, score in parse(drop_blanks(read_file("scores.csv"))):
    print(f"{name:6} {score}")

And now the file can be any size at all. Iterating a file hands over one line at a time, and each stage passes that single line along before asking for the next, so the whole chain only ever holds one line — whether the file is four lines or forty million.

Which is why the source is the piece you swap

The string version above is not size-proof: text.splitlines() builds a list of every line, and the string itself was already in memory. Laziness has to start at the source — the rest of the pipeline is lazy either way.

yield from

yield hands out one value. yield from hands out every value of an iterable, one at a time, and then carries on. It is exactly a loop that yields each item, written in one line:

yield from thing          # is the same as:

for item in thing:
    yield item

The thing can be anything you can loop over — a string, a range, a list, or another generator. Handing the work to a second generator like that is called delegating, which is where the syntax gets its name.

Delegating to another generator Python

def letters():
    yield from "abc"           # yields "a", then "b", then "c"

def letters_the_long_way():    # identical, spelled out
    for letter in "abc":
        yield letter

print(list(letters()))
print(list(letters_the_long_way()))

def both():
    yield from letters()       # delegate to another generator
    yield from range(3)        # then to a range

print(list(both()))

yield and yield from are not interchangeable

yield "abc" produces one value — the string itself. yield from "abc" produces three — one per character. Whenever a generator gives you a single lumpy item where you expected several, this is usually why.

One value, or its contents Python

def one_value():
    yield "abc"

def three_values():
    yield from "abc"

print(list(one_value()))
print(list(three_values()))

Where it earns its keep

Delegating is most useful when a generator calls itself. This one walks a list that may contain other lists, to any depth, and yields the plain values:

Flattening a nested list Python

def flatten(items):
    for item in items:
        if isinstance(item, list):
            yield from flatten(item)   # the recursive call is a generator too
        else:
            yield item

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

Without yield from, that inner line has to become a loop over the recursive call — for inner in flatten(item): yield inner — and every level of nesting adds another loop to read past. The one-line form says what is happening: whatever that call produces, pass it straight out.

itertools

The standard library ships a toolbox of ready-made generators.

A few worth knowing Python

from itertools import count, cycle, islice, chain, groupby, pairwise

print(list(islice(count(10, 5), 4)))          # 10, 15, 20, 25
print(list(islice(cycle("ab"), 5)))
print(list(chain([1, 2], [3, 4])))
print(list(pairwise([1, 2, 3, 4])))           # consecutive pairs

words = ["apple", "avocado", "banana", "blueberry", "cherry"]
for letter, group in groupby(words, key=lambda w: w[0]):
    print(letter, list(group))

A generator is consumed once

Once you have iterated a generator, it is empty — there is no rewind. If you need the values twice, either collect them into a list or call the generator function again to get a fresh one.

Exhaustion Python

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

print(list(gen))
print(list(gen))      # empty - already consumed

Before you write your own

Predict. Four lines. When does start get printed?

What will this print? Python

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

gen = countdown(3)
print("made it")
print(next(gen))
print(list(gen))

Check your prediction

made it
start
3
[2, 1]

Calling countdown(3) runs none of its body — it only makes the generator. The body starts on the first next, runs to the first yield, and pauses. list(gen) then takes whatever is left: 3 has already been used.

Fix. This should print the sum of the even numbers below 7, then the largest: 12, then 6. The second line crashes.

Used up Python

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

numbers = evens(7)
print(sum(numbers))
print(max(numbers))

Check your fix

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

numbers = list(evens(7))
print(sum(numbers))
print(max(numbers))

12
6

A generator can be walked through only once. sum used it up, so max got nothing — max() iterable argument is empty. When you need the values twice, keep them in a list (or call evens(7) again).

Exercise

A log filter

Write a generator errors_only(lines) that yields only the lines containing ERROR, and another with_numbers(lines) that yields (line_number, line) pairs counting from 1 in the original log. Chain them so the numbers refer to the original file, not the filtered result. Your output should be:

2: ERROR disk full
4: ERROR timeout

log = """INFO starting up
ERROR disk full
INFO retrying
ERROR timeout
INFO done"""

# Your code here

Hint 1: the idea

Two small generators joined in a pipeline, each with one job: one attaches a number to every line, the other lets through only the error lines. Order matters. Number the lines before filtering, or the numbers will count errors instead of lines.

Hint 2: a first step

Write with_numbers first: loop over enumerate(lines, start=1) and yield number, line. Check it on its own — for pair in with_numbers(log.splitlines()): print(pair) should show all five lines, numbered 1 to 5. Then write a filter that takes those pairs.

Show one solution

log = """INFO starting up
ERROR disk full
INFO retrying
ERROR timeout
INFO done"""

def with_numbers(lines):
    for number, line in enumerate(lines, start=1):
        yield number, line

def errors_only(numbered):
    for number, line in numbered:
        if "ERROR" in line:
            yield number, line

for number, line in errors_only(with_numbers(log.splitlines())):
    print(f"{number}: {line}")

Numbering first, then filtering, is what keeps the line numbers honest — filter first and you would be numbering the errors, not the log.

What you learned

for runs on iter() and next(), ending at StopIteration.
A function with yield is a generator: it pauses and resumes, keeping its state.
Generators use constant memory regardless of how many values they produce.
They can be infinite, and chained into pipelines.
yield from passes on every value of an iterable — a loop that yields each item, in one line.
A generator can only be consumed once.