Generators and yield

yield turns a function into a generator

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


gen = countdown(3)
print(type(gen))         # <class 'generator'>

print(next(gen))         # 3
print(next(gen))         # 2
print(next(gen))         # 1
# print(next(gen))       # StopIteration

Calling countdown(3) runs no code at all. It builds a generator object. The body starts executing only on the first next(), runs until it reaches a yield, hands the value back and freezes.

Pausing and resuming

def traced():
    print("  A: start")
    yield 1
    print("  B: after the first yield")
    yield 2
    print("  C: after the second yield")


gen = traced()
print("created, nothing has run yet")
print("got", next(gen))
print("got", next(gen))
try:
    next(gen)
except StopIteration:
    print("finished")

created, nothing has run yet
  A: start
got 1
  B: after the first yield
got 2
  C: after the second yield
finished

The function's local variables, and its position in the code, survive between calls. That is what makes a generator different from a function that returns a list.

Why it matters: memory

import sys


def squares_list(n):
    return [i * i for i in range(n)]


def squares_generator(n):
    for i in range(n):
        yield i * i


print(sys.getsizeof(squares_list(1_000_000)))        # about 8 MB
print(sys.getsizeof(squares_generator(1_000_000)))   # about 200 bytes

The list holds a million integers. The generator holds a position and a recipe. Both produce the same values when iterated.

def read_large_file(path):
    """Yield one line at a time; the file is never loaded whole."""
    with open(path, encoding="utf-8") as handle:
        for line in handle:
            yield line.rstrip("\n")


def errors_only(lines):
    for line in lines:
        if "ERROR" in line:
            yield line


def first_field(lines):
    for line in lines:
        yield line.split(",")[0]


# A pipeline: nothing is stored, each line flows through all three stages
for value in first_field(errors_only(read_large_file("app.log"))):
    print(value)

This processes a file of any size in constant memory. Building three intermediate lists instead would need three copies of the data.

Generator expressions

squares_list = [n * n for n in range(10)]        # a list, built now
squares_gen = (n * n for n in range(10))         # a generator, lazy

print(squares_list)
print(squares_gen)
print(list(squares_gen))

# Inside a call, the brackets are optional
print(sum(n * n for n in range(1000)))
print(max((len(w) for w in ["a", "bcd"]), default=0))
print(", ".join(str(n) for n in range(5)))
print(any(n > 100 for n in range(1000)))         # stops at 101

yield from

def chain(*iterables):
    for iterable in iterables:
        for item in iterable:
            yield item


def chain_shorter(*iterables):
    for iterable in iterables:
        yield from iterable          # delegate to the sub-iterable


print(list(chain_shorter([1, 2], "ab", range(3))))

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


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

Infinite generators

def naturals():
    n = 1
    while True:
        yield n
        n += 1


def take(iterable, n):
    for i, value in enumerate(iterable):
        if i >= n:
            return
        yield value


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


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


print(list(take(fibonacci(), 10)))

Never call list() on an infinite generator. Always bound it with itertools.islice, a break, or a helper like take above.

Generators keep state

def running_average():
    """A generator that maintains state across values."""
    total = 0
    count = 0
    while True:
        value = yield (total / count if count else 0)
        if value is not None:
            total += value
            count += 1


averager = running_average()
next(averager)                    # prime it: run to the first yield

print(averager.send(10))          # 10.0
print(averager.send(20))          # 15.0
print(averager.send(30))          # 20.0

yield is an expression as well as a statement. value = yield x sends x out and waits for a value to be sent back in with .send(). This turns a generator into a small coroutine, and it must be primed with one next() first.

Closing and cleanup

def managed():
    print("acquiring")
    try:
        yield 1
        yield 2
        yield 3
    finally:
        print("releasing")          # runs even if the caller stops early


gen = managed()
print(next(gen))
gen.close()                         # triggers the finally block

def read_lines(path):
    with open(path, encoding="utf-8") as handle:
        for line in handle:
            yield line
    # the file closes when the generator is exhausted or closed


for line in read_lines("notes.txt"):
    if "STOP" in line:
        break                       # the with block still closes the file

A worked pipeline

def read_records(rows):
    for row in rows:
        yield row.strip()


def skip_blank_and_comments(rows):
    for row in rows:
        if row and not row.startswith("#"):
            yield row


def parse(rows):
    for row in rows:
        parts = row.split(",")
        if len(parts) != 3:
            continue
        name, dept, salary = parts
        try:
            yield {"name": name, "dept": dept, "salary": int(salary)}
        except ValueError:
            continue


def above(records, threshold):
    for record in records:
        if record["salary"] > threshold:
            yield record


raw = [
    "# staff list",
    "Meera,eng,90000",
    "",
    "Arun,design,75000",
    "broken row",
    "Sara,eng,82000",
    "Ravi,eng,not-a-number",
]

pipeline = above(parse(skip_blank_and_comments(read_records(raw))), 80_000)

for record in pipeline:
    print(f"{record['name']:<8}{record['dept']:<10}{record['salary']:>8,}")

Each stage does one thing, each is independently testable, and no stage ever holds more than one record. The same code works on four rows or forty million.

Generators versus returning a list

def evens_list(n):
    result = []
    for i in range(n):
        if i % 2 == 0:
            result.append(i)
    return result


def evens_gen(n):
    for i in range(n):
        if i % 2 == 0:
            yield i


print(evens_list(10))               # [0, 2, 4, 6, 8]
print(evens_gen(10))                # <generator object ...>
print(list(evens_gen(10)))          # [0, 2, 4, 6, 8]

Use a generator when the result is consumed once, could be large, or might be infinite. Use a list when the caller needs its length, needs to index it, or will walk it more than once.

Common mistakes

• Printing a generator and seeing <generator object> instead of values.
• Iterating a generator twice and getting nothing the second time.
• Calling len() on a generator.
• Mixing return value and yield and expecting the value from iteration; it becomes the StopIteration value instead.
• Calling list() on an infinite generator.
• Forgetting to prime a generator before using .send().
• Assuming the body runs when the generator is created.

Best practices

• Use a generator whenever the caller only needs to loop once.
• Chain small generators into a pipeline instead of building intermediate lists.
• Use a generator expression when the result goes straight into sum, any, join or a loop.
• Use yield from for delegation and recursion.
• Wrap resources in with inside the generator so cleanup happens even on early exit.

Practice

1. Write a generator producing the first n prime numbers.
2. Compare the memory used by a list and a generator of one million items.
3. Build a three stage pipeline that reads, filters and transforms a list of records.
4. Write a generator that reads a file and yields only lines matching a pattern, and prove the file closes on an early break.
5. Explain why gen = (x for x in range(3)); print(sum(gen), sum(gen)) prints 3 0.

Conclusion

A generator is a function that pauses. It produces values on demand, keeps its local state between them, and uses constant memory regardless of how much data flows through. For anything large, streamed or infinite, it is not an optimisation - it is the only workable approach.