Python Generators and Context Managers: Memory-Efficient Code Patterns

!DOCTYPE html> Python Generators and Context Managers: Memory-Efficient and Clean Code

Python Generators and Context Managers: Memory-Efficient and Clean Code

Two of Python's most powerful but underused features. Generators let you process billions of rows without running out of RAM. Context managers guarantee cleanup even when exceptions occur.

Generator Basics: yield vs return

A generator function uses yield instead of return. Calling it returns a lazy iterator -- values are computed one at a time, on demand.

# Regular function -- builds entire list in memory
def squares_list(n):
    return [x**2 for x in range(n)]  # ALL in RAM at once

# Generator -- one value at a time, constant memory
def squares_gen(n):
    for x in range(n):
        yield x**2  # pauses here, resumes on next()

import sys
lst = squares_list(1_000_000)
gen = squares_gen(1_000_000)

print(sys.getsizeof(lst))  # ~8,000,056 bytes (8 MB)
print(sys.getsizeof(gen))  # 112 bytes -- always!

# Both work identically in a for loop
for sq in squares_gen(10):
    print(sq, end=" ")
# 0 1 4 9 16 25 36 49 64 81

Generator Expressions vs List Comprehensions

# List comprehension -- eager, all in memory
squares_list = [x**2 for x in range(1_000_000)]      # 8MB

# Generator expression -- lazy, constant memory
squares_gen  = (x**2 for x in range(1_000_000))      # 112 bytes

# Use generator expressions when:
# - You only iterate once
# - The dataset is large
# - You're passing to sum(), max(), any(), all(), etc.

total = sum(x**2 for x in range(1_000_000))          # no list created!
first_even = next(x for x in range(100) if x % 2 == 0)

# Chain generator expressions as pipeline (zero intermediate lists)
result = sum(
    x**2
    for x in range(1_000_000)
    if x % 3 == 0
)

Memory Comparison

Approach1M integers10M integersIterable once?
List [...]~8 MB~80 MBMultiple times
Generator (...)112 bytes112 bytesOnce only
Generator function112 bytes112 bytesOnce (re-call to reuse)

send(), throw(), and close()

# send() -- pass a value INTO a running generator (coroutine pattern)
def running_average():
    total = 0
    count = 0
    average = 0
    while True:
        value = yield average       # yield sends out avg, receives next value
        if value is None:
            break
        total += value
        count += 1
        average = total / count

gen = running_average()
next(gen)             # prime the generator (advance to first yield)
print(gen.send(10))   # 10.0
print(gen.send(20))   # 15.0
print(gen.send(30))   # 20.0

# throw() -- inject an exception into the generator
def safe_gen():
    try:
        while True:
            value = yield
    except ValueError as e:
        print(f"Caught: {e}")
        yield "recovered"

g = safe_gen()
next(g)
print(g.throw(ValueError, "bad input"))  # Caught: bad input -> "recovered"

yield from: Delegation

# yield from delegates to a sub-generator
def flatten(nested):
    for item in nested:
        if isinstance(item, list):
            yield from flatten(item)  # recurse into nested lists
        else:
            yield item

data = [1, [2, 3, [4, 5]], 6, [7, [8, 9]]]
print(list(flatten(data)))
# [1, 2, 3, 4, 5, 6, 7, 8, 9]

# yield from also passes return values from sub-generators
def sub():
    yield 1
    yield 2
    return "sub_done"   # this return value is accessible via StopIteration

def main():
    result = yield from sub()   # result = "sub_done"
    print(f"Sub returned: {result}")
    yield 3

Generator Pipelines

import csv
from pathlib import Path

# Process a 10GB CSV with constant memory using generator pipeline

def read_csv_rows(filepath: str):
    """Stage 1: Read rows lazily"""
    with open(filepath, newline='') as f:
        reader = csv.DictReader(f)
        yield from reader

def filter_active(rows):
    """Stage 2: Filter"""
    for row in rows:
        if row['status'] == 'active':
            yield row

def parse_amounts(rows):
    """Stage 3: Transform"""
    for row in rows:
        row['amount'] = float(row['amount'])
        yield row

def above_threshold(rows, threshold: float):
    """Stage 4: Filter again"""
    for row in rows:
        if row['amount'] > threshold:
            yield row

# Build pipeline -- nothing executes yet
pipeline = read_csv_rows("transactions.csv")
pipeline = filter_active(pipeline)
pipeline = parse_amounts(pipeline)
pipeline = above_threshold(pipeline, 1000.0)

# Execute -- processes one row at a time, O(1) memory
total = sum(row['amount'] for row in pipeline)
print(f"Total high-value active transactions: {total:,.2f}")

Context Managers: __enter__ and __exit__

class DatabaseTransaction:
    def __init__(self, connection):
        self.conn = connection
        self.cursor = None

    def __enter__(self):
        self.cursor = self.conn.cursor()
        self.conn.begin()          # start transaction
        return self.cursor

    def __exit__(self, exc_type, exc_val, exc_tb):
        if exc_type is None:
            self.conn.commit()     # success -- commit
            print("Transaction committed")
        else:
            self.conn.rollback()   # exception -- rollback
            print(f"Transaction rolled back due to: {exc_val}")
        self.cursor.close()
        return False               # False = re-raise exception; True = suppress

# Usage
with DatabaseTransaction(conn) as cursor:
    cursor.execute("INSERT INTO orders VALUES (1, 'widget', 99.99)")
    cursor.execute("UPDATE inventory SET stock = stock - 1 WHERE id = 42")
    # If anything raises here, rollback is automatic

@contextmanager Decorator

from contextlib import contextmanager, asynccontextmanager
import time, logging

@contextmanager
def timer(label: str):
    start = time.perf_counter()
    try:
        yield                           # code inside 'with' block runs here
    finally:
        elapsed = time.perf_counter() - start
        print(f"{label}: {elapsed:.4f}s")

with timer("database query"):
    results = db.execute("SELECT * FROM large_table")

@contextmanager
def managed_tempdir():
    import tempfile, shutil
    tmpdir = tempfile.mkdtemp()
    try:
        yield tmpdir
    finally:
        shutil.rmtree(tmpdir, ignore_errors=True)

with managed_tempdir() as tmpdir:
    process_files(tmpdir)
# tmpdir is always deleted, even on exception

@asynccontextmanager
async def db_session(engine):
    async with engine.begin() as conn:
        try:
            yield conn
            await conn.commit()
        except Exception:
            await conn.rollback()
            raise

Production Use Cases

Chunked File Processing

@contextmanager
def chunked_reader(filepath: str, chunk_size: int = 8192):
    """Read large file in chunks without loading into memory"""
    with open(filepath, 'rb') as f:
        while True:
            chunk = f.read(chunk_size)
            if not chunk:
                break
            yield chunk

def count_bytes(filepath: str) -> int:
    total = 0
    with chunked_reader(filepath) as chunks:
        for chunk in chunks:   # wait -- contextmanager doesn't work like this
            total += len(chunk)
    return total

# Correct pattern -- generator function for iteration:
def read_chunks(filepath: str, chunk_size: int = 8192):
    with open(filepath, 'rb') as f:
        while chunk := f.read(chunk_size):
            yield chunk

total_bytes = sum(len(chunk) for chunk in read_chunks("bigfile.bin"))

Connection Pool Context Manager

from contextlib import contextmanager
from queue import Queue

class ConnectionPool:
    def __init__(self, create_conn, size=10):
        self._pool = Queue(maxsize=size)
        for _ in range(size):
            self._pool.put(create_conn())

    @contextmanager
    def acquire(self):
        conn = self._pool.get()
        try:
            yield conn
        finally:
            self._pool.put(conn)   # always returns to pool

pool = ConnectionPool(lambda: create_db_connection(), size=5)

with pool.acquire() as conn:
    results = conn.execute("SELECT 1")
[YES] Rule of thumb: If you're building a list just to iterate over it once -- use a generator instead. If you have setup/teardown code -- use a context manager. These two patterns eliminate entire categories of memory leaks and resource bugs.

Key Takeaways: Generators are lazy iterators -- use them for large data. Generator pipelines compose cleanly with zero intermediate memory. Context managers guarantee teardown code runs. Use @contextmanager for simple cases, class-based for complex state.

PythonGeneratorsContext ManagersMemory EfficiencyAdvanced Python

Post a Comment

0 Comments