Python Advanced – Generators & Iterators

The Python track's post on generators covers what yield does and why laziness matters. This one goes under it: the protocol every iterator implements, generators used as coroutines rather than producers, and the standard library built on top of both.

The protocol, written out by hand

Two methods. __iter__ returns the iterator, and __next__ returns the next item or raises StopIteration:

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

    def __iter__(self):
        return self

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

print(list(Countdown(3)))     # Output: [3, 2, 1]

c = Countdown(2)
print(next(c), next(c))       # Output: 2 1

That is everything a for loop needs. No base class and no registration — the loop calls iter() then next() until StopIteration, and any object with those two methods qualifies.

Note this class returns self from __iter__, which makes it a one-shot iterator: exhaust it and looping again yields nothing. That is the same rule as a generator, and it is the reason the next section exists.

Iterable and iterator are different things

class Deck:
    """Iterable, not an iterator — a fresh iterator per loop."""
    def __init__(self, cards):
        self.cards = cards

    def __iter__(self):
        return iter(self.cards)      # a NEW iterator each time

deck = Deck(["A", "K", "Q"])

print(list(deck))     # Output: ['A', 'K', 'Q']
print(list(deck))     # Output: ['A', 'K', 'Q']

it = iter(deck)
print(next(it))       # Output: A
print(list(it))       # Output: ['K', 'Q']

A list is iterable but is not an iterator: it has no __next__, and each for loop gets a fresh iterator from it. That is why you can loop over a list twice and not over a generator.

The design rule follows: implement __iter__ returning a new iterator when the object is a collection, and return self only when the object genuinely represents one traversal — a cursor, a stream, a socket.

The same thing as a generator

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

print(list(countdown(3)))     # Output: [3, 2, 1]

Four lines against fourteen, and the state lives in a local variable rather than an attribute you have to maintain. This is why hand-written iterator classes are rare: the generator gets __iter__, __next__ and StopIteration for free.

Write the class when the object needs to be more than an iterator — when it also has methods, attributes or a repr worth having. Otherwise write the generator.

Generators receive as well as produce

yield is an expression, not a statement, and its value is whatever the caller sends in:

def running_total():
    total = 0
    while True:
        amount = yield total       # yield out, receive in
        if amount is not None:
            total += amount

acc = running_total()
print(next(acc))          # Output: 0
print(acc.send(100))      # Output: 100
print(acc.send(-25))      # Output: 75
acc.close()

That is a coroutine — a generator used as a consumer rather than a producer. The next(acc) first is not optional: a fresh generator has not reached its first yield, so there is nowhere for a sent value to land, and calling send() straight away raises TypeError.

close() raises GeneratorExit inside the generator, and throw() raises an exception of your choosing at the paused line. In practice async def has taken over most of what this was used for — see Async & Await — but it is the machinery underneath, and send still turns up in pipeline code.

yield from

def flatten(items):
    for item in items:
        if isinstance(item, list):
            yield from flatten(item)     # delegate to the recursive call
        else:
            yield item

print(list(flatten([1, [2, [3, [4]], 5], 6])))     # Output: [1, 2, 3, 4, 5, 6]

yield from delegates to another iterable. It is not merely shorthand for a loop: it also forwards send, throw and close to the sub-generator, and passes back its return value.

Recursive flattening is where it reads best — the alternative is a nested loop that re-yields each item by hand.

contextlib turns one into a context manager

from contextlib import contextmanager

@contextmanager
def timed(label):
    print(f"start {label}")
    try:
        yield label.upper()          # everything before = __enter__
    finally:
        print(f"end {label}")        # everything after = __exit__

with timed("import") as name:
    print(f"working on {name}")

# Output: start import
# Output: working on IMPORT
# Output: end import

One yield, exactly one. What comes before it is the setup, what it yields becomes the as variable, and what comes after is the teardown. The try/finally is what makes the cleanup run even when the block raises — without it, this is a context manager that leaks on exactly the path you wrote it for.

This replaces writing a class with __enter__ and __exit__ for the common case, and it is the clearest demonstration that a generator is a resumable function rather than a sequence.

itertools, past the obvious

from itertools import groupby, accumulate, pairwise, takewhile

rows = [("Checking", 100), ("Checking", 250), ("Savings", 90)]

for kind, group in groupby(rows, key=lambda r: r[0]):
    print(kind, sum(amount for _, amount in group))

# Output: Checking 350
# Output: Savings 90

print(list(accumulate([100, -25, 250])))          # Output: [100, 75, 325]
print(list(pairwise([1, 2, 3, 4])))               # Output: [(1, 2), (2, 3), (3, 4)]
print(list(takewhile(lambda n: n > 0, [3, 1, -1, 5])))   # Output: [3, 1]

accumulate is a running total — exactly the balance_after column a statement needs. pairwise (3.10) gives consecutive pairs, which is how you compute deltas between readings. takewhile stops at the first failure rather than filtering the whole input.

groupby only groups consecutive equal keys. That is the trap: it is not SQL's GROUP BY. Feed it unsorted input and you get several groups for the same key. Sort by the same key first, or use a defaultdict, which does not care about order and is usually the better tool.

Generators are not free

import timeit

setup = "data = list(range(1000))"
loop = timeit.timeit("sum(x * 2 for x in data)", setup=setup, number=1000)
comp = timeit.timeit("sum([x * 2 for x in data])", setup=setup, number=1000)

print(loop > 0 and comp > 0)      # Output: True
print(round(loop / comp, 0) >= 1) # Output: True

On a small collection the list comprehension is usually the faster of the two, because each yield costs a suspend and resume while the list version runs a tight loop in C. Generators win on memory and on time-to-first-item, not on raw throughput.

So the rule is about size and shape, not style: use a generator when the data is large, unbounded, or you may stop early. Use a list when it is small and you need it more than once.

Next: functools, operator, and reduce.