Iterators and Generators 🌊

What a for loop really does

crew = ["Guybrush", "Elaine"]

it = iter(crew)          # ask for an iterator
print(next(it))
print(next(it))

try:
    next(it)
except StopIteration:
    print("StopIteration: that is the loop's stop signal")

Guybrush
Elaine
StopIteration: that is the loop's stop signal

Every for loop is that: call iter(), call next() until StopIteration. Anything that supports those two calls can be looped over, which is why for works identically on lists, strings, dictionaries, files and things you write yourself.

Generators: iterators without the ceremony

def countdown(n):
    """A generator: it yields values instead of returning one."""
    while n > 0:
        yield n
        n -= 1
    yield "Liftoff!"


for value in countdown(3):
    print(value)

print(type(countdown(3)))
print(list(countdown(2)))

3
2
1
Liftoff!
<class 'generator'>
[2, 1, 'Liftoff!']

One word, yield, changes everything. The function no longer runs to completion and returns a value. It runs until the first yield, hands that value out, and freezes, keeping all its local variables. The next next() resumes exactly where it stopped.

def noisy():
    print("  starting")
    yield 1
    print("  woke up again")
    yield 2
    print("  finishing")


gen = noisy()
print("nothing has run yet")
print(next(gen))
print(next(gen))

nothing has run yet
  starting
1
  woke up again
2

Why it matters: memory

import sys


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


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


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

print(f"list:      {sys.getsizeof(as_list):>10,} bytes")
print(f"generator: {sys.getsizeof(as_gen):>10,} bytes")
print(f"same total: {sum(as_list) == sum(squares_gen(1_000_000))}")

list:       8,448,728 bytes
generator:        200 bytes
same total: True

Eight megabytes versus two hundred bytes, for the same answer. The generator never holds more than one value at a time. For a million items on a modern laptop this is a curiosity; for a log file bigger than your RAM it is the entire difference between possible and impossible.

The real-world shape: processing a big file

from pathlib import Path

Path("server.log").write_text("""INFO  started
ERROR disk full
INFO  retrying
ERROR still full
INFO  done
""", encoding="utf-8")


def read_lines(path):
    """Yield each line, stripped. Never holds the whole file."""
    with open(path, encoding="utf-8") as f:
        for line in f:
            yield line.rstrip("\n")


def only_errors(lines):
    """Yield only the ERROR lines."""
    for line in lines:
        if line.startswith("ERROR"):
            yield line


def messages(lines):
    """Yield just the message part."""
    for line in lines:
        yield line.split(maxsplit=1)[1]


pipeline = messages(only_errors(read_lines("server.log")))

for message in pipeline:
    print(message)

disk full
still full

That is a pipeline. Nothing is read until the for loop pulls, and then each line flows through all three stages one at a time. The file could be a terabyte and the memory use would not change. This is exactly how Unix pipes work, and it is one of the most reusable structures in programming.

Infinite sequences, which lists cannot do

def fibonacci():
    """Every Fibonacci number. All of them. Forever."""
    a, b = 0, 1
    while True:
        yield a
        a, b = b, a + b


from itertools import islice

print(list(islice(fibonacci(), 10)))

# find the first Fibonacci number over a thousand
for n in fibonacci():
    if n > 1000:
        print("first over 1000:", n)
        break

[0, 1, 1, 2, 3, 5, 8, 13, 21, 34]
first over 1000: 1597

Generator expressions

numbers = range(1, 11)

squares_list = [n * n for n in numbers]        # builds it all
squares_gen = (n * n for n in numbers)         # builds nothing yet

print(squares_list)
print(type(squares_gen).__name__)
print(sum(squares_gen))
print(sum(squares_gen))          # exhausted! generators are one-shot

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

yield from: delegating

def chain(*iterables):
    for it in iterables:
        yield from it          # yield every value from that one


def flatten(nested):
    """Flatten any depth of nested lists."""
    for item in nested:
        if isinstance(item, list):
            yield from flatten(item)
        else:
            yield item


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

[1, 2, 'a', 'b', 3, 4]
[1, 2, 3, 4, 5, 6, 7]

Writing an iterator class

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

    def __iter__(self):
        """Return a fresh iterator each time, so this can be looped twice."""
        n = self.start
        while n > 0:
            yield n
            n -= 1


c = Countdown(3)
print(list(c))
print(list(c))          # works again, unlike a bare generator

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

Making __iter__ a generator function is the tidiest way to build a reusable iterable. Each for loop calls __iter__ again and gets a fresh generator, so the one-shot problem disappears.

itertools: the ones worth knowing

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

print(list(islice(count(10, 5), 4)))
print(list(islice(cycle("ab"), 5)))
print(list(chain([1, 2], [3])))
print(list(accumulate([1, 2, 3, 4])))
print(list(pairwise([1, 2, 3, 4])))

crew = [("deck", "Otis"), ("deck", "Meathook"), ("bridge", "Elaine")]
for station, people in groupby(crew, key=lambda pair: pair[0]):
    print(station, [name for _, name in people])

[10, 15, 20, 25]
['a', 'b', 'a', 'b', 'a']
[1, 2, 3]
[1, 3, 6, 10]
[(1, 2), (2, 3), (3, 4)]
deck ['Otis', 'Meathook']
bridge ['Elaine']

def numbers(n):
    print("  (numbers started)")
    for i in range(1, n + 1):
        yield i


def evens(source):
    for n in source:
        if n % 2 == 0:
            yield n


def squared(source):
    for n in source:
        yield n * n


pipeline = squared(evens(numbers(10)))
print("built the pipeline, nothing has run")
print(list(pipeline))

built the pipeline, nothing has run
  (numbers started)
[4, 16, 36, 64, 100]

from pathlib import Path

Path("log.txt").write_text("one\ntwo\nthree\n", encoding="utf-8")


def reversed_lines(path):
    """Yield lines last-first. Reads the whole file: see the note."""
    with open(path, encoding="utf-8") as f:
        lines = f.readlines()
    for line in reversed(lines):
        yield line.rstrip("\n")


print(list(reversed_lines("log.txt")))

['three', 'two', 'one']

def running_average(numbers):
    """Yield the mean of everything seen so far."""
    total = 0
    count = 0
    for n in numbers:
        total += n
        count += 1
        yield total / count


for average in running_average([10, 20, 30, 40]):
    print(f"{average:.2f}")

10.00
15.00
20.00
25.00