Bytebook Compiler

Real Programs

Iterators & Generators

How for loops really work, the iterator protocol, writing generators with yield, lazy evaluation and infinite sequences, and generator pipelines.

3 min read6 runnable examples4-question quiz

How a for loop works

for x in collection calls iter(collection) to get an iterator, then calls next() on it repeatedly until it raises StopIteration:

Python
nums = [10, 20, 30]
it = iter(nums)
print(next(it), next(it), next(it))
try:
    next(it)
except StopIteration:
    print("exhausted")
Output
10 20 30
exhausted
Try it Yourself

An iterable is anything you can loop over (lists, strings, dicts, files, ranges). An iterator is the object that actually remembers the position.

Writing an iterator class

Python
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(5)))
for n in Countdown(3):
    print(n, end=" ")
print("liftoff!")
Output
[5, 4, 3, 2, 1, 0]
3 2 1 0 liftoff!
Try it Yourself

Generators: the easy way

A generator function uses yield instead of return. Each yield hands out one value and pauses the function, keeping all its local variables until the next value is requested:

Python
def countdown(start):
    while start >= 0:
        yield start
        start -= 1

def fibonacci():
    a, b = 0, 1
    while True:              # infinite — fine, values are produced only on demand
        yield a
        a, b = b, a + b

print(list(countdown(5)))
fib = fibonacci()
print([next(fib) for _ in range(10)])

gen = countdown(2)
print(next(gen), next(gen), next(gen))
Output
[5, 4, 3, 2, 1, 0]
[0, 1, 1, 2, 3, 5, 8, 13, 21, 34]
2 1 0
Try it Yourself

Why generators matter: lazy evaluation

A generator produces values one at a time, so it uses almost no memory no matter how many values it yields:

Python
import sys

squares_list = [n * n for n in range(1_000_000)]
squares_gen = (n * n for n in range(1_000_000))
print("list:", sys.getsizeof(squares_list), "bytes")
print("generator:", sys.getsizeof(squares_gen), "bytes")
print(sum(squares_gen))
Output · yours may differ
list: 8448728 bytes
generator: 208 bytes
333332833333500000
Try it Yourself

Generator pipelines

Chain generators like assembly-line stages. Each item flows through all stages before the next item starts:

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

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

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

data = """Asha,88

Ravi,95
Meera,79
"""
results = parse(non_empty(read_lines(data)))
for name, score in results:
    print(f"{name}: {score}")
Output
Asha: 88
Ravi: 95
Meera: 79
Try it Yourself

itertools

Python
import itertools

print(list(itertools.islice(itertools.count(1, 3), 5)))        # 1, 4, 7, 10, 13
print(list(itertools.islice(itertools.cycle("AB"), 5)))
print(list(itertools.chain([1, 2], [3], [4, 5])))
print([(k, len(list(g))) for k, g in itertools.groupby("aaabccdddd")])
print(list(itertools.product("xy", [1, 2])))
Output
[1, 4, 7, 10, 13]
['A', 'B', 'A', 'B', 'A']
[1, 2, 3, 4, 5]
[('a', 3), ('b', 1), ('c', 2), ('d', 4)]
[('x', 1), ('x', 2), ('y', 1), ('y', 2)]
Try it Yourself
Iterator class Generator function Generator expression
Syntax __iter__ + __next__ def … yield (x for x in xs)
Length most code short one line
State kept in attributes kept automatically kept automatically

Test yourself

4 questions · pick an answer to check it instantly

  1. 1

    What does yield do?

  2. 2

    What exception signals that an iterator has no more values?

  3. 3

    What is printed?

    Python
    def g():
    yield 1
    yield 2
it = g()
print(next(it) + next(it))
  4. 4

    Why is (x * x for x in range(10**8)) safe to create while the list version may not be?

Esc
↑↓ move↵ open