How to Build a Generator Function in Python – Step-by-Step Guide for Data Science 2026

How to Build a Generator Function in Python – Step-by-Step Guide for Data Science 2026

Building your own generator functions with yield is one of the most valuable skills for handling large-scale data in Python. Unlike regular functions that return once, generator functions can pause and resume, producing values one at a time with minimal memory usage.

TL;DR — Core Rules

• Use def and yield instead of return
• The function automatically becomes a generator when it contains yield
• Call it like a normal function — it returns a generator object

1. Simple Generator Function

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

# Usage
for number in count_up_to(5):
    print(number)

2. Real Data Science Generator Functions

import pandas as pd

# Example 1: Row-by-row processor
def process_sales_rows(df):
    for row in df.itertuples():
        profit = row.amount * 0.25
        category = "Premium" if profit > 500 else "Standard"
        yield {
            "customer_id": row.customer_id,
            "amount": row.amount,
            "profit": round(profit, 2),
            "category": category
        }

# Example 2: Chunked file reader with enrichment
def read_and_enrich_large_csv(file_path, chunk_size=50000):
    for chunk in pd.read_csv(file_path, chunksize=chunk_size):
        chunk["profit"] = chunk["amount"] * 0.25
        chunk["log_amount"] = chunk["amount"].apply(lambda x: round(x**0.5, 2) if x > 0 else 0)
        yield chunk

# Create a small demo file so this example runs as-is
# (small stand-in for huge_sales.csv; in your project this is your real large file)
import numpy as np
rng = np.random.default_rng(42)
n = 60
pd.DataFrame({
    "order_id": range(1001, 1001 + n),
    "order_date": pd.date_range("2026-01-01", periods=n, freq="3D"),
    "customer_id": rng.integers(101, 121, n),
    "customer_name": rng.choice([" ana", "ben ", "chen", "dana"], n),
    "product_id": rng.integers(1, 9, n),
    "region": rng.choice(["North", "South", "East", "West"], n),
    "category": rng.choice(["Electronics", "Computers", "Books", "Home"], n),
    "payment_method": rng.choice(["card", "cash", "wallet"], n),
    "segment": rng.choice(["consumer", "business"], n),
    "quantity": rng.integers(1, 6, n),
    "amount": rng.uniform(50, 2500, n).round(2),
    "discount": rng.choice([0.0, 0.05, 0.1, 0.2], n),
}).assign(
    price=lambda d: (d["amount"] / d["quantity"]).round(2),
    profit=lambda d: (d["amount"] * rng.uniform(0.05, 0.3, n)).round(2),
).assign(
    ship_date=lambda d: d["order_date"] + pd.Timedelta(days=2),
    status=lambda d: np.where(d["amount"] > 2000, "priority", "standard"),
).to_csv("huge_sales.csv", index=False)

# Usage
for enriched_chunk in read_and_enrich_large_csv("huge_sales.csv"):
    print(f"Processed chunk with {len(enriched_chunk)} rows")
    # Save or further process this chunk

3. Advanced Generator with Multiple Yields & State

def batch_processor(data, batch_size=100):
    batch = []
    for item in data:
        batch.append(item)
        if len(batch) >= batch_size:
            yield batch
            batch = []
    if batch:
        yield batch

# Usage
large_dataset = range(250)   # demo data
for batch in batch_processor(large_dataset):
    print(f"Processing batch of {len(batch)} items")

4. Best Practices for Building Generators in 2026

• Keep generator functions focused on one clear responsibility
• Use descriptive names and document what is being yielded
• Prefer itertuples() over iterrows() inside generators for speed
• Use yield from when delegating to another generator
• Test generators with next() and small inputs first

Building a Generator Step by Step

Building a generator function follows a simple recipe: write a normal function that would collect results in a list, then replace the append with yield and remove the final return. The function now produces items on demand.

Add parameters for anything that might change, such as batch size or thresholds, and keep setup and clean-up inside the function so resources are handled in one place.

Test the generator by taking a few items with next() or islice before plugging it into a larger pipeline.

Worked Example: From List Function to Generator

A batching function is converted from list-building to yielding, and both are compared.

def batches_list(items, size):
    out, batch = [], []
    for item in items:
        batch.append(item)
        if len(batch) == size:
            out.append(batch)
            batch = []
    if batch:
        out.append(batch)
    return out

def batches(items, size):
    batch = []
    for item in items:
        batch.append(item)
        if len(batch) == size:
            yield batch
            batch = []
    if batch:
        yield batch

data = range(10)
print(batches_list(data, 4))
print(list(batches(data, 4)))
print(next(batches(data, 4)))

Both versions produce the same three batches, the last one shorter. The generator version can also be used on an endless stream, because it never needs all the batches at once.

Common Mistakes

• Forgetting to yield the final partial batch.
• Reusing the same batch list object after yielding it.
• Mixing return value and yield in confusing ways.
• Not testing generators on small inputs first.

Quick Summary

• Replace append-and-return with yield.
• Parameterise sizes and thresholds.
• Keep setup and clean-up inside the function.
• Test with next() or islice first.

Conclusion

Building custom generator functions is a key skill for modern data science. In 2026, they allow you to process massive datasets, build reusable pipelines, and keep memory usage low. Start simple, practice the yield pattern, and gradually move your data processing code from full lists to powerful, lazy generators.

Next steps:

• Take one of your existing data processing scripts and rewrite the core loop as a custom generator function

Related Tutorials

• Building Custom Generator Functions in Python – Advanced Memory-Efficient Patterns 2026
• Using Generator Functions in Python – Practical Patterns for Data Science 2026