Python Generators: Lazy Iteration with yield

Suppose you need to scan a 10 GB log file for errors. Read it into a list and your program runs out of memory before it finds the first one. Process it one line at a time and it runs comfortably on a laptop.

Generators are Python's tool for exactly that: producing values one at a time, only when they're needed. In this article you'll learn how yield works, why generators use almost no memory, how to chain them into pipelines, and the gotchas that trip people up.

The problem with building everything up front#

A list comprehension builds the whole list before you can use a single item:

import sys

squares_list = [x * x for x in range(1_000_000)]
squares_gen = (x * x for x in range(1_000_000))

print(sys.getsizeof(squares_list))  # 8448728  (about 8 MB)
print(sys.getsizeof(squares_gen))   # 104      (the same for any size)

Swapping the square brackets for parentheses gives you a generator expression. It doesn't compute anything yet. It only remembers how to produce the next value, so its size stays tiny however many items it will eventually yield. (The exact byte counts depend on your Python version.)

Your first generator function#

Any function that contains yield is a generator function. Calling it doesn't run its body. It returns a generator object that runs the body step by step as you ask for values:

def countdown(n):
    print("Starting")
    while n > 0:
        yield n
        n -= 1
    print("Done")

gen = countdown(3)
print(gen)  # <generator object countdown at 0x7f...>  (nothing printed yet)

Ask for values with next():

>>> next(gen)
Starting
3
>>> next(gen)
2
>>> next(gen)
1
>>> next(gen)
Done
Traceback (most recent call last):
  ...
StopIteration

In practice you rarely call next() yourself. A for loop calls it for you and stops cleanly at StopIteration:

for n in countdown(3):
    print(n)

How yield pauses a function#

A normal function runs from top to bottom and its local variables vanish when it returns. A generator is different:

1. next() runs the body until it reaches a yield.
2. The yielded value is handed to the caller and the function pauses, keeping its local variables and its exact position.
3. The next next() resumes right after that yield.
4. When the function returns, the generator raises StopIteration and is finished.

That pause is the whole trick. Between two next() calls the generator does no work and holds only its own local variables.

Generator expressions#

For simple cases you don't need a function at all. A generator expression has the same shape as a list comprehension, with parentheses:

total = sum(x * x for x in range(1_000_000))
has_admin = any(user.is_admin for user in users)
longest = max(len(word) for word in words)

When the generator expression is the only argument to a function, you can drop its own parentheses, as above. Functions like sum(), any(), all(), min() and max() consume it item by item without ever building a list.

Building lazy pipelines#

Generators really shine when you chain them. Each stage pulls items from the previous one, one at a time, so the whole pipeline uses a constant amount of memory however big the input is.

Here we find the first five server errors in a web server's access log:

from itertools import islice

def read_lines(path):
    with open(path, encoding="utf-8") as file:
        for line in file:
            yield line.rstrip("\n")

def parse(lines):
    for line in lines:
        parts = line.split()
        if len(parts) >= 9:
            yield {"ip": parts[0], "path": parts[6], "status": int(parts[8])}

def server_errors(records):
    return (r for r in records if r["status"] >= 500)

lines = read_lines("access.log")
records = parse(lines)
first_five = islice(server_errors(records), 5)

for record in first_five:
    print(record["ip"], record["path"])

Nothing happens until the for loop asks for the first record. Then each stage does just enough work to hand over one item. And once five errors have been found, the file stops being read, even if it has another 9 GB to go.

Infinite generators#

Since values are only computed on demand, a generator can describe a sequence that never ends:

from itertools import islice

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

print(list(islice(fibonacci(), 10)))
# [0, 1, 1, 2, 3, 5, 8, 13, 21, 34]

while True would hang a normal function forever. Here it's fine, because the generator pauses after every value and islice() simply stops asking after ten. The itertools module has ready-made infinite generators too: count(), cycle() and repeat().

Delegating with yield from#

yield from hands over to another iterable and yields everything it produces. It's perfect for recursive structures:

def flatten(items):
    for item in items:
        if isinstance(item, list):
            yield from flatten(item)
        else:
            yield item

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

Without it you'd write for x in flatten(item): yield x. yield from is shorter and also passes through send() and close() calls correctly.

Gotchas#

A generator can only be used once. After it's exhausted, it stays empty:

squares = (x * x for x in range(3))
print(list(squares))  # [0, 1, 4]
print(list(squares))  # []  (already used up)

If you need the values more than once, either store them in a list or call the generator function again to get a fresh generator.

No len() and no indexing. A generator doesn't know how many items it will produce, so len(gen) and gen[0] raise TypeError. Use itertools.islice() to take a slice, or convert to a list when the data is small.

Errors show up late. Because the body runs lazily, a bug inside a generator only appears when the item is actually requested, which can be far away from the line that created it.

List or generator?#

A good rule: if you'll loop over the values once, start with a generator. Switch to a list when you need to reuse, count or index them.

Recap#

• A function with yield returns a generator; its body runs only when values are requested.
• Each next() runs to the next yield, then the function pauses with its local variables intact.
• Generator expressions (x for x in ...) are the lazy cousin of list comprehensions.
• Chain generators into pipelines that process huge inputs in constant memory.
• Use islice() to take a few items from long or infinite generators, and yield from to delegate.
• Generators are single-use, have no len(), and can hold resources open while paused.

Next time you're about to build a big list only to loop over it once, try replacing the brackets with parentheses. Often that's all it takes.