Python Generators: yield, Generator Expressions and Lazy Evaluation

A generator is the easiest way to create an iterator. It looks like a normal function, but uses yield instead of return. Each yield hands back one value and pauses the function; the next request resumes it exactly where it left off. Values are produced one at a time, only when needed, which can save huge amounts of memory.

A first generator#

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

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

gen = countdown(2)
print(type(gen).__name__, next(gen), next(gen))

3
2
1
generator 2 1

Calling countdown(3) does not run the body. It returns a generator object. The code only runs as values are requested by a loop or next().

Pause and resume#

The prints show exactly when each part runs:

def steps():
    print("start")
    yield 1
    print("middle")
    yield 2
    print("end")

g = steps()
print(next(g))
print(next(g))
try:
    next(g)
except StopIteration:
    print("StopIteration")

start
1
middle
2
end
StopIteration

All local variables are kept between pauses. When the function finishes, the generator raises StopIteration, just like any iterator.

Generator expressions#

Like a list comprehension but with parentheses, a generator expression produces values lazily instead of building a list:

import sys

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

print(sys.getsizeof(squares_list), sys.getsizeof(squares_gen))
print(sum(n * n for n in range(1_000_000)))

8448728 200
333332833333500000

The list holds a million numbers in memory (about 8 MB for the list structure alone); the generator stays tiny regardless of size. (Exact byte counts vary between Python versions.) When a function takes a single generator argument, you can drop the extra parentheses, as in sum(n * n for n in ...).

Processing large files lazily#

Generators let you process a file of any size line by line. Given this log file:

INFO start
ERROR disk full
INFO retry
ERROR timeout

def read_lines(path):
    with open(path) as f:
        for line in f:
            yield line.strip()

errors = [line for line in read_lines("big.log") if line.startswith("ERROR")]
print(errors)

['ERROR disk full', 'ERROR timeout']

Only one line is in memory at a time, so this works the same for a 1 KB file or a 10 GB log.

Infinite sequences#

Because values are produced on demand, a generator can run forever. Take just what you need with itertools.islice:

import itertools

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

print(list(itertools.islice(fibonacci(), 10)))

[0, 1, 1, 2, 3, 5, 8, 13, 21, 34]

Pipelines#

Chained generators form a pipeline where each item flows through every step before the next item starts:

def pipeline():
    nums = range(1, 11)
    evens = (n for n in nums if n % 2 == 0)
    squared = (n * n for n in evens)
    return sum(squared)

print(pipeline())

220

No intermediate lists are created: 4 + 16 + 36 + 64 + 100 = 220.

yield from#

yield from hands over to another iterable, producing all its values:

def numbers():
    yield from range(3)
    yield from "ab"

print(list(numbers()))

[0, 1, 2, 'a', 'b']

Common mistakes#

Practice#

def chunks(items, size):
    for i in range(0, len(items), size):
        yield items[i:i + size]

for c in chunks(list(range(1, 11)), 4):
    print(c)

[1, 2, 3, 4]
[5, 6, 7, 8]
[9, 10]

Key takeaway#