Skip to main content

Generators

In this chapter, we'll explore generators - a powerful way to work with sequences lazily. We'll cover:

  • What are generators?
  • The yield keyword
  • Generator expressions
  • Lazy evaluation benefits
  • Iterators vs generators
  • Testing generators

What is a generator?​

A generator is a function that returns an iterator. Instead of returning all values at once, it yields them one at a time:

def count_up_to(n):
    i = 1
    while i <= n:
        yield i
        i += 1


# Usage
for num in count_up_to(5):
    print(num)
# Prints: 1, 2, 3, 4, 5

Why use generators?​

Memory efficiency​

A list stores all elements in memory:

# This creates a list with 1 million integers in memory
numbers = [x for x in range(1_000_000)]

A generator produces values on-demand:

# This creates only one integer at a time
numbers = (x for x in range(1_000_000))

Infinite sequences​

Generators can represent infinite sequences:

def infinite_counter():
    n = 0
    while True:
        yield n
        n += 1


# Take first 5 values
counter = infinite_counter()
for _ in range(5):
    print(next(counter))

The yield keyword​

yield is like return, but it pauses the function instead of ending it:

def my_generator():
    print("Before first yield")
    yield 1
    print("Before second yield")
    yield 2
    print("Before third yield")
    yield 3
    print("After all yields")


gen = my_generator()
print(next(gen))  # Prints "Before first yield", returns 1
print(next(gen))  # Prints "Before second yield", returns 2

Testing generators​

Write the test first​

from generators import fibonacci_gen


def test_fibonacci_generator():
    gen = fibonacci_gen()
    
    first_10 = [next(gen) for _ in range(10)]
    
    assert first_10 == [0, 1, 1, 2, 3, 5, 8, 13, 21, 34]


def test_fibonacci_is_generator():
    gen = fibonacci_gen()
    
    # Generators are iterators
    assert hasattr(gen, '__iter__')
    assert hasattr(gen, '__next__')

Implementation​

def fibonacci_gen():
    """Generate Fibonacci numbers infinitely."""
    a, b = 0, 1
    while True:
        yield a
        a, b = b, a + b

Generator expressions​

Like list comprehensions, but lazy:

# List comprehension - creates list immediately
squares_list = [x ** 2 for x in range(1000)]

# Generator expression - creates generator
squares_gen = (x ** 2 for x in range(1000))

Generator expressions use parentheses instead of brackets.

Testing generator expressions​

def test_generator_expression():
    gen = (x * 2 for x in [1, 2, 3])
    
    assert list(gen) == [2, 4, 6]
    
    # Generator is exhausted after iteration
    assert list(gen) == []

Practical examples​

Reading large files​

def read_lines(filename):
    """Read a file line by line without loading entire file."""
    with open(filename) as f:
        for line in f:
            yield line.strip()


# Test
def test_read_lines(tmp_path):
    file_path = tmp_path / "test.txt"
    file_path.write_text("line1\nline2\nline3\n")
    
    lines = list(read_lines(file_path))
    
    assert lines == ["line1", "line2", "line3"]

Processing data in chunks​

def chunked(iterable, size):
    """Yield chunks of the iterable."""
    chunk = []
    for item in iterable:
        chunk.append(item)
        if len(chunk) == size:
            yield chunk
            chunk = []
    if chunk:
        yield chunk


# Test
def test_chunked():
    data = [1, 2, 3, 4, 5, 6, 7]
    chunks = list(chunked(data, 3))
    
    assert chunks == [[1, 2, 3], [4, 5, 6], [7]]

Filtering data​

def filter_evens(numbers):
    """Yield only even numbers."""
    for n in numbers:
        if n % 2 == 0:
            yield n


# Test
def test_filter_evens():
    evens = list(filter_evens([1, 2, 3, 4, 5, 6]))
    assert evens == [2, 4, 6]

yield from​

Use yield from to delegate to another generator:

def flatten(nested_list):
    """Flatten a nested list."""
    for item in nested_list:
        if isinstance(item, list):
            yield from flatten(item)  # Recursively yield from nested
        else:
            yield item


# Test
def test_flatten():
    nested = [1, [2, 3], [4, [5, 6]], 7]
    flat = list(flatten(nested))
    
    assert flat == [1, 2, 3, 4, 5, 6, 7]

Generators vs iterators​

Iterator (class-based)​

class CountUp:
    def __init__(self, n):
        self.n = n
        self.current = 0
    
    def __iter__(self):
        return self
    
    def __next__(self):
        if self.current >= self.n:
            raise StopIteration
        self.current += 1
        return self.current

Generator (function-based)​

def count_up(n):
    for i in range(1, n + 1):
        yield i

Generators are usually simpler and more readable.

Generator state​

Generators maintain state between yields:

def running_total():
    """Yield running total of received values."""
    total = 0
    while True:
        value = yield total
        if value is not None:
            total += value


# Usage
gen = running_total()
next(gen)  # Prime the generator
print(gen.send(10))  # 10
print(gen.send(20))  # 30
print(gen.send(5))   # 35

Testing generator with send​

def test_running_total():
    gen = running_total()
    next(gen)  # Prime
    
    assert gen.send(10) == 10
    assert gen.send(20) == 30
    assert gen.send(5) == 35

Generator pipelines​

Chain generators for data processing:

def read_numbers(filename):
    with open(filename) as f:
        for line in f:
            yield int(line.strip())


def double(numbers):
    for n in numbers:
        yield n * 2


def filter_large(numbers, threshold):
    for n in numbers:
        if n > threshold:
            yield n


# Pipeline
numbers = read_numbers("data.txt")
doubled = double(numbers)
large = filter_large(doubled, 100)

for n in large:
    print(n)

Testing pipelines​

def test_generator_pipeline():
    source = iter([1, 2, 3, 4, 5])
    doubled = double(source)
    large = filter_large(doubled, 5)
    
    result = list(large)
    
    assert result == [6, 8, 10]

Common generator utilities​

Python's itertools module provides generator utilities:

from itertools import islice, chain, cycle, count

# Take first n items
first_5 = list(islice(infinite_counter(), 5))

# Chain iterables
combined = chain([1, 2], [3, 4])

# Cycle through items
colors = cycle(["red", "green", "blue"])

# Count infinitely
numbers = count(start=10, step=5)  # 10, 15, 20, ...

Testing with itertools​

from itertools import islice


def test_infinite_generator_with_islice():
    def count_forever():
        n = 0
        while True:
            yield n
            n += 1
    
    first_5 = list(islice(count_forever(), 5))
    
    assert first_5 == [0, 1, 2, 3, 4]

A complete example: CSV processor​

# test_csv_processor.py
import pytest
from csv_processor import process_csv


def test_process_csv(tmp_path):
    csv_file = tmp_path / "data.csv"
    csv_file.write_text("name,age\nAlice,30\nBob,25\nCharlie,35\n")
    
    rows = list(process_csv(csv_file))
    
    assert rows == [
        {"name": "Alice", "age": "30"},
        {"name": "Bob", "age": "25"},
        {"name": "Charlie", "age": "35"},
    ]


def test_process_csv_is_lazy(tmp_path):
    csv_file = tmp_path / "data.csv"
    csv_file.write_text("name,age\nAlice,30\nBob,25\n")
    
    gen = process_csv(csv_file)
    
    # Generator hasn't read file yet
    first_row = next(gen)
    assert first_row == {"name": "Alice", "age": "30"}

Implementation:

# csv_processor.py
def process_csv(filename):
    """Process CSV file row by row."""
    with open(filename) as f:
        headers = None
        for line in f:
            values = line.strip().split(",")
            if headers is None:
                headers = values
            else:
                yield dict(zip(headers, values))

Wrapping up​

We've covered:

  • Generators - Functions with yield that return iterators
  • yield keyword - Pause and produce a value
  • Generator expressions - Lazy comprehensions with ()
  • yield from - Delegate to another generator
  • Generator pipelines - Chain generators for processing
  • itertools - Standard library generator utilities

Key takeaways​

  1. Use generators for large or infinite sequences
  2. Generators are memory-efficient - one value at a time
  3. Generators are iterators - use with for, next(), list()
  4. Use yield from for delegation
  5. Test generators by consuming with list() or next()

When to use generators​

  • Processing large files line by line
  • Working with infinite sequences
  • Chaining data transformations
  • Any case where you don't need all values at once

Generators make your code more memory-efficient and often clearer by expressing computation as a sequence of transformations.