Generators, yield, and lazy pipelines

Writing your own iterator the easy way

You could implement __iter__ and __next__ on a class, but Python has a shortcut: a generator function. Use yield instead of return and the function becomes a value factory that can pause and resume:

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

Calling countdown(3) runs no code yet. It returns a generator object, which is an iterator. Each next call runs the body until the next yield, hands out that value, and freezes right there, local variables intact. When the function body ends, StopIteration is raised for you. This is called lazy evaluation: values are produced only when asked for.

Generators earn their keep when the data cannot fit in memory: reading a multi-gigabyte log file line by line, or streaming millions of database rows. Production data pipelines are assembled from exactly the pieces in this lesson.

Watching a generator pause and resume

The print calls inside countdown reveal the timing. The body does not begin running when the generator is created, and it advances only as far as the next yield each time a value is requested.

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

gen = countdown(3)
print("created, nothing ran yet")
print(next(gen))
print(next(gen))
for rest in gen:
    print(rest)

Output

created, nothing ran yet
  starting
3
2
1
  done

Two details are worth pausing on. starting prints after created, nothing ran yet, proving the body was frozen until the first next. And the for loop picks up exactly where the manual next calls left off, at 1, because it is consuming the same generator rather than a fresh one.

Generator expressions and pipelines

Write a comprehension with parentheses instead of square brackets and you get a generator expression, the lazy cousin of the list comprehension from lesson 1-1:

squares = (n * n for n in range(1_000_000))

No million-item list is built. Values stream out one at a time, so memory stays tiny. Feed one straight into a consuming function and you may drop the extra parentheses:

total = sum(n * n for n in range(1000))

Chain stages to build pipelines: filter with one generator, transform with another, then let sum, max, or a for loop pull values through the whole chain. Use a list when you need indexing or multiple passes, use a generator when you stream through the data once.

evens

evens(limit) is a generator function that yields the even numbers from 0 up to but not including limit. Because it is an iterable, a generator expression can consume it directly inside a single sum call.

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

print(list(evens(10)))
print(sum(e * e for e in evens(10)))

Output

[0, 2, 4, 6, 8]
120

range(0, limit, 2) steps by two, which skips the odd numbers without any test at all. Testing n % 2 == 0 inside a plain range(limit) would also work and is worth recognizing when you meet it. The sum checks out by hand: 0² + 2² + 4² + 6² + 8² = 0 + 4 + 16 + 36 + 64 = 120.

powers_of_two

powers_of_two(limit) yields 1, 2, 4, 8, and so on for as long as the value stays below limit. It keeps its own value variable and doubles it after each yield, which is a shape a range cannot express.

def powers_of_two(limit):
    value = 1
    while value < limit:
        yield value
        value *= 2

print(list(powers_of_two(50)))
print(sum(powers_of_two(50)))

Output

[1, 2, 4, 8, 16, 32]
63

The order inside the loop matters: yield value comes first, then value *= 2, so the very first value handed out is 1. Note also that each powers_of_two(50) call constructs a brand-new generator with its own value. That is why the second line still sees all six numbers even though the first line consumed a generator to exhaustion.

Square brackets build, parentheses promise

[n * n for n in big] and (n * n for n in big) look almost identical and behave very differently. The first builds the entire list in memory immediately. The second creates a generator that computes each value only when something asks for it.

Both can be fed to a for loop or to sum, so they are interchangeable for a single pass. The generator wins whenever big is large or infinite, and the list wins whenever you need to walk the data more than once.