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.
Contents
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
| Approach | 1M integers | 10M integers | Iterable once? |
|---|---|---|---|
List [...] | ~8 MB | ~80 MB | Multiple times |
Generator (...) | 112 bytes | 112 bytes | Once only |
| Generator function | 112 bytes | 112 bytes | Once (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
0 Comments