Python Generators and Itertools: Memory-Efficient Data Processing

#python #generators #tutorial #programming

Python Generators and Itertools: Memory-Efficient Data Processing

Generators process data lazily — no need to load everything into memory.

Generator Functions

def count_up(start: int, end: int):
    current = start
    while current <= end:
        yield current
        current += 1

# Doesn't load all numbers into memory!
for num in count_up(1, 1_000_000):
    if num > 5:
        break
    print(num)

Generator Expressions

# List comprehension — loads ALL into memory
squares_list = [x**2 for x in range(10_000_000)]

# Generator expression — lazy, memory efficient
squares_gen = (x**2 for x in range(10_000_000))

# sum() works with generators!
total = sum(x**2 for x in range(100))

Read Large Files Efficiently

def read_large_csv(filename: str):
    with open(filename) as f:
        header = next(f).strip().split(",")
        for line in f:
            values = line.strip().split(",")
            yield dict(zip(header, values))

# Process 10GB CSV without loading it all
for row in read_large_csv("huge_file.csv"):
    process(row)

itertools Essentials

import itertools

# chain — iterate multiple sequences
for item in itertools.chain([1, 2], [3, 4], [5, 6]):
    print(item)  # 1, 2, 3, 4, 5, 6

# batched (Python 3.12+)
for batch in itertools.batched(range(10), 3):
    print(list(batch))  # [0,1,2], [3,4,5], [6,7,8], [9]

# groupby
data = [("A", 1), ("A", 2), ("B", 3), ("B", 4)]
for key, group in itertools.groupby(data, key=lambda x: x[0]):
    print(key, list(group))

# islice — slice a generator
first_5 = list(itertools.islice(count_up(1, 100), 5))
print(first_5)  # [1, 2, 3, 4, 5]

Pipeline Pattern

def pipeline(*funcs):
    def process(data):
        for func in funcs:
            data = func(data)
        return data
    return process

def parse_numbers(lines):
    return (int(line.strip()) for line in lines)

def filter_even(numbers):
    return (n for n in numbers if n % 2 == 0)

def square(numbers):
    return (n**2 for n in numbers)

process = pipeline(parse_numbers, filter_even, square)
result = list(process(["1", "2", "3", "4", "5"]))
print(result)  # [4, 16]