Iterators and Generators

Ava wants to process a log file with ten million lines. If she loads every line into a list, her laptop runs out of memory, and freezes. She needs a way to handle one line at a time, forgetting the previous ones. Python's answer is the generator, a function that produces values on demand, and never holds them all at once.

• Ten million lines will not fit in memory
• Process one item at a time
• The tool: generators

What a for loop really does

A for loop works with iterators. An iterable, like a list, can give you an iterator, with the iter function. Each call to next returns the following item. When there are none left, a StopIteration is raised, and the for loop quietly ends. Understanding this unlocks the rest of the lesson.

plants = ["fern", "ivy"]
it = iter(plants)
print(next(it))
print(next(it))
try:
    next(it)
except StopIteration:
    print("no more items")

fern
ivy
no more items

Your first generator

A generator is a function that uses yield instead of return. Each time it yields, it pauses, remembering where it was, and hands one value to the caller. The next request resumes it right after the yield. Here, a generator counts down. Nothing runs until the loop asks for a value.

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

for x in countdown(3):
    print(x)
g = countdown(2)
print(next(g), next(g))

3
2
1
2 1

Generators save memory

Here is the memory difference. A list of one million squares takes about eight megabytes. The generator expression, in round brackets, takes almost nothing, because it produces each square only when asked. We can still sum it. When data is large, or endless, generators are the right choice.

import sys

big_list = [x * x for x in range(1_000_000)]
lazy = (x * x for x in range(1_000_000))
print(sys.getsizeof(big_list) > 1_000_000)
print(sys.getsizeof(lazy) < 500)
print(sum(lazy))

True
True
333332833333500000

A pipeline of generators

Generators combine into pipelines, where each stage passes items to the next, one at a time. Here, the first stage reads lines, the second filters the errors, and the third cleans them up. Memory use stays tiny, no matter how large the log is. This is how professionals process huge data.

def read(lines):
    for ln in lines:
        yield ln

def errors(lines):
    for ln in lines:
        if "ERROR" in ln:
            yield ln

def clean(lines):
    for ln in lines:
        yield ln.strip().lower()

log = ["INFO ok", "ERROR disk full ", "INFO ok", "ERROR timeout "]
for item in clean(errors(read(log))):
    print(item)

error disk full
error timeout

A generator of even numbers

Your turn. Write a generator called evens, that takes a limit, and yields the even numbers from zero up to that limit. Then loop over it, and print them. Ready? Here is one solution. A loop with a step of two, and a yield inside.

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

for e in evens(10):
    print(e)

0
2
4
6
8
10

Summary