Iterators and Generators

Iteration in Python runs on one small protocol, and generators make implementing it almost free. Together they let you process data that is infinite, or far larger than memory, with ordinary-looking code.

Iterables vs iterators

They are different things, and the distinction explains most iteration surprises.

• An iterable can produce an iterator. It implements __iter__. Lists, tuples, dicts, sets, strings, and files are all iterable.
• An iterator produces values one at a time. It implements __next__ and __iter__ (returning itself). It is consumed as you go and cannot be rewound.

xs = [1, 2, 3]        # an iterable
it = iter(xs)         # an iterator over it

next(it)   # => 1
next(it)   # => 2
next(it)   # => 3
next(it)   # StopIteration

StopIteration is how an iterator says “done”. It is not an error — the for statement catches it internally.

A for loop is exactly this:

for x in xs:
    body(x)

# desugars to
it = iter(xs)
while True:
    try:
        x = next(it)
    except StopIteration:
        break
    body(x)

The critical consequence: an iterator is single-use.

it = iter([1, 2, 3])
list(it)     # => [1, 2, 3]
list(it)     # => []          — already exhausted

Whereas an iterable produces a fresh iterator each time:

xs = [1, 2, 3]
list(xs)     # => [1, 2, 3]
list(xs)     # => [1, 2, 3]

next() accepts a default instead of raising:

next(it, None)      # => None when exhausted

Writing an iterator by hand

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

    def __iter__(self):
        return CountdownIterator(self.start)     # fresh state each time


class CountdownIterator:
    def __init__(self, current):
        self.current = current

    def __iter__(self):
        return self

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


list(Countdown(3))     # => [3, 2, 1]
list(Countdown(3))     # => [3, 2, 1]  — reusable, because __iter__ makes a new iterator

That is a lot of code for very little. Generators collapse it to four lines.

Generators

A function containing yield is a generator function. Calling it runs no code — it returns a generator object. Each next() runs the body until the next yield, hands back that value, and freezes the function’s entire state (locals, instruction pointer) until the next call.

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

list(countdown(3))     # => [3, 2, 1]

g = countdown(3)
g              # => <generator object countdown at 0x...>  — nothing has run yet
next(g)        # => 3
next(g)        # => 2

Returning from a generator (or falling off the end) raises StopIteration. A return value inside a generator attaches that value to the exception, which yield from can retrieve.

Generators are iterators, so they are single-use:

g = countdown(3)
list(g)     # => [3, 2, 1]
list(g)     # => []

Caution

If you need to iterate a generator’s output twice, materialise it (data = list(gen)) or call the generator function again. There is no rewind.

Lazy evaluation

Nothing is computed until it is requested. That means generators can be infinite.

def naturals():
    n = 0
    while True:
        yield n
        n += 1

from itertools import islice
list(islice(naturals(), 5))     # => [0, 1, 2, 3, 4]

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

Danger

list(naturals()) never returns and eventually exhausts memory. Always bound an infinite generator with islice, takewhile, zip against something finite, or a break.

Generator expressions

Comprehension syntax with parentheses gives a generator instead of a list.

squares = (n * n for n in range(1_000_000))     # allocates nothing
sum(squares)

Parentheses can be omitted when it is the sole argument to a call:

sum(n * n for n in range(100))
any(line.startswith("ERROR") for line in log)
max((len(w), w) for w in words)

any and all short-circuit, so this stops at the first match rather than scanning everything:

if any(user.is_admin for user in users):     # stops at the first admin
    ...

Compare the cost of eager and lazy pipelines over a large file:

# Eager: whole file in memory, twice
lines = open("huge.log", encoding="utf-8").readlines()
errors = [l for l in lines if "ERROR" in l]
count = len(errors)

# Lazy: one line in memory at a time
with open("huge.log", encoding="utf-8") as f:
    count = sum(1 for line in f if "ERROR" in line)

The file object is itself a lazy iterator of lines — never call .readlines() on something large.

Chaining generators into a pipeline

Each stage is lazy, so data flows through one item at a time.

def read_lines(path):
    with open(path, encoding="utf-8") as f:
        yield from f

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

def parse(lines):
    for line in lines:
        key, _, value = line.partition("=")
        yield key.strip(), value.strip()


config = dict(parse(strip_blank(read_lines("app.conf"))))

Memory use is constant regardless of file size, and each stage is independently testable.

Caution

A generator holding an open resource only closes it when the generator is exhausted or garbage-collected. In read_lines above, abandoning the iterator early leaves the with block suspended until the generator is closed — CPython’s reference counting usually handles it promptly, but contextlib.closing or explicit gen.close() makes it deterministic.

yield from

yield from iterable delegates to a sub-iterator: it yields everything the sub-iterator produces, and forwards send/throw correctly.

def flatten(nested):
    for item in nested:
        if isinstance(item, list):
            yield from flatten(item)      # recursion made trivial
        else:
            yield item

list(flatten([1, [2, [3, [4]]], 5]))     # => [1, 2, 3, 4, 5]

Without it you would write for x in flatten(item): yield x — equivalent for simple cases, but not for the two-way protocol.

yield from also captures the sub-generator’s return value:

def inner():
    yield 1
    return "done"

def outer():
    result = yield from inner()
    print(result)      # => 'done'

Sending values in

Generators are two-way. gen.send(value) resumes the generator and makes the paused yield expression evaluate to value.

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

acc = accumulator()
next(acc)          # prime it — run to the first yield => 0
acc.send(10)       # => 10
acc.send(5)        # => 15

gen.throw(ExcType) raises inside the generator at the paused yield; gen.close() raises GeneratorExit there, which is how cleanup in a finally gets to run.

Note

This coroutine-style use of generators is largely historical — it was the foundation of asyncio before async/await existed. You will meet it in older code and in @contextlib.contextmanager, which relies on throw to deliver the body’s exception to your generator.

When to use what

Situation	Use
Data fits in memory, needed repeatedly	A list
Large or streaming data, one pass	A generator
Infinite or unbounded sequence	A generator
Simple transform/filter of an iterable	A generator expression
Complex per-item logic, multiple yields	A generator function
An object that must be iterable many times	A class with __iter__ returning a fresh generator

That last pattern is the clean way to make a reusable iterable:

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

    def __iter__(self):
        n = self.start
        while n > 0:
            yield n
            n -= 1

Each for loop calls __iter__ and gets a brand-new generator, so it is reusable — with none of the manual state machine.

itertools highlights

itertools is a set of fast, memory-efficient iterator building blocks implemented in C.

from itertools import (
    chain, islice, groupby, product, permutations, combinations,
    count, cycle, repeat, accumulate, takewhile, dropwhile,
    zip_longest, tee, starmap, compress, pairwise,
)

chain — concatenate iterables lazily.

list(chain([1, 2], "ab", (3,)))            # => [1, 2, 'a', 'b', 3]
list(chain.from_iterable([[1, 2], [3]]))   # => [1, 2, 3]   — flatten one level

islice — slice any iterable, including infinite ones. No negative indices.

list(islice(count(), 5))          # => [0, 1, 2, 3, 4]
list(islice(range(20), 5, 15, 2)) # start, stop, step

groupby — group consecutive items sharing a key.

data = [("a", 1), ("a", 2), ("b", 3)]
for key, group in groupby(data, key=lambda t: t[0]):
    print(key, list(group))
# a [('a', 1), ('a', 2)]
# b [('b', 3)]

Danger

groupby only groups adjacent items — it is not SQL GROUP BY. Sort by the same key first, or you will get repeated groups. For unsorted aggregation use collections.defaultdict(list). Also note each group is an iterator that is invalidated as soon as you advance to the next group: convert it with list() before moving on.

product, permutations, combinations — combinatorics without nested loops.

list(product("ab", [1, 2]))          # => [('a',1), ('a',2), ('b',1), ('b',2)]
list(product([0, 1], repeat=3))      # all 3-bit combinations
list(permutations("abc", 2))         # ordered pairs, 6 of them
list(combinations("abc", 2))         # => [('a','b'), ('a','c'), ('b','c')]

Infinite generators.

count(10, 2)          # 10, 12, 14, ...
cycle("ab")           # a, b, a, b, ...
repeat(None, 3)       # None, None, None

accumulate — running totals, or any binary function.

list(accumulate([1, 2, 3, 4]))                  # => [1, 3, 6, 10]
import operator
list(accumulate([1, 2, 3, 4], operator.mul))    # => [1, 2, 6, 24]
list(accumulate([3, 1, 4], max))                # => [3, 3, 4]

takewhile / dropwhile — stop or skip based on a predicate.

list(takewhile(lambda n: n < 5, [1, 3, 6, 1]))   # => [1, 3]   — stops at 6
list(dropwhile(lambda n: n < 5, [1, 3, 6, 1]))   # => [6, 1]

zip_longest, pairwise, tee.

list(zip_longest([1, 2, 3], "ab", fillvalue="?"))  # => [(1,'a'), (2,'b'), (3,'?')]
list(pairwise([1, 2, 3, 4]))                        # => [(1,2), (2,3), (3,4)]  — 3.10+
a, b = tee(iterable, 2)                             # two independent iterators

Caution

tee buffers everything one branch has consumed but the other has not. If the branches drift far apart, memory grows to hold the gap — and you must stop using the original iterator once you have teed it.

Also useful: functools.reduce for folds, and heapq.nlargest/nsmallest for top-k without a full sort.

Memory in practice

import sys

sys.getsizeof([n for n in range(1_000_000)])     # ~8 MB of pointers
sys.getsizeof((n for n in range(1_000_000)))     # ~200 bytes

The generator holds only its frame and current state; the list holds a million references plus the integer objects. For a pipeline that touches each item once, the list is pure waste.

Where a list is still right: you need len(), random access, more than one pass, or to sort. Generators give up all of those.

Key points

• Iterable → has __iter__; iterator → has __next__ and is single-use.
• for is iter() plus next() in a loop, catching StopIteration.
• A generator function returns a generator; each yield suspends and resumes the whole frame.
• Generator expressions are lazy comprehensions — use them as function arguments.
• yield from delegates to a sub-iterator and forwards its return value.
• Make a reusable iterable by returning a fresh generator from __iter__.
• itertools covers chaining, slicing, grouping, and combinatorics in C.
• groupby needs sorted input and its groups expire when you advance.