Python Generators and yield: Processing Large Data Efficiently

The short answer

A generator is a function with yield that returns values one at a time, on demand. Until you ask for the next item, no code runs. That is why a generator does not load the whole collection into memory: processing a multi-gigabyte file takes roughly as much memory as processing a single line.

Iterable, iterator, generator

You can loop over a list as many times as you want. An iterator or generator works only once: once exhausted, it is empty.

How yield works

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

gen = countdown(3)
next(gen)  # 3
next(gen)  # 2

Calling countdown(3) computes nothing — it returns a generator object. Each next() runs the body up to the next yield, hands out the value and freezes the function state. When the body ends, StopIteration is raised and a for loop stops on its own.

Generator expressions

Square brackets build a list, parentheses build a generator:

squares_list = [x * x for x in range(10_000_000)]  # all in memory
squares_gen = (x * x for x in range(10_000_000))   # one by one

total = sum(x * x for x in range(10_000_000))     # parentheses optional

Use a generator expression when you need the result once — for sum, max, any or writing to a file.

A lazy pipeline for a big file

The real strength of generators is chaining. Each step takes a stream in and gives a stream out:

import csv

def read_rows(path):
    with open(path, newline="", encoding="utf-8") as f:
        yield from csv.DictReader(f)

def only_paid(rows):
    for row in rows:
        if row["status"] == "paid":
            yield row

def amounts(rows):
    for row in rows:
        yield float(row["amount"])

total = sum(amounts(only_paid(read_rows("orders.csv"))))

At no step is the whole file in memory: a row is read, filtered, converted and added right away. The same approach works for logs:

def errors(path):
    with open(path, encoding="utf-8") as f:
        for line in f:
            if " ERROR " in line:
                yield line.rstrip()

For money in real code, prefer Decimal over float — float is used here only for brevity.

yield from

yield from iterable yields every item of another iterator. It is shorter than a loop and handy for recursion and for splitting a generator into parts:

def walk(tree):
    yield tree.value
    for child in tree.children:
        yield from walk(child)

yield from also correctly forwards send(), exceptions and the return value of the inner generator.

Common pitfalls

• One-time iteration. A second pass over the same generator yields nothing:

gen = (x for x in range(3))
list(gen)  # [0, 1, 2]
list(gen)  # []

If you need the data twice, create the generator again, or store it in a list if it fits in memory.

• Closed file. If a function opens a file with with and returns a generator without yielding inside the block, the file closes before reading. Keep yield inside the with block.
• No len() or indexing. A generator has no length and gen[5] fails. Use itertools.islice for slices.
• Lazy errors. An exception is raised when an item is read, not when the generator is created, so the traceback may point somewhere unexpected.
• Debugging with print(gen) shows only <generator object ...>, not the data.

Generator or list

• Data is large or infinite — generator.
• You need one pass and an aggregate — generator.
• You need repeated passes, sorting, len or indexes — list.
• Data is small — the difference is negligible, pick whatever reads better.