Python Generators and itertools for Memory-Efficient Code

A list holds every item in memory. A generator produces items one at a time, on demand. For large files, streams or infinite sequences, that's the difference between working and running out of memory.

yield

def read_lines(path):
    with open(path) as f:
        for line in f:
            yield line.rstrip("\n")

for line in read_lines("huge.log"):   # one line in memory at a time
    if "ERROR" in line:
        print(line)

Generator expressions

squares = [n * n for n in range(10_000_000)]   # builds a huge list
squares = (n * n for n in range(10_000_000))   # lazy, almost no memory

total = sum(n * n for n in range(10_000_000))  # parentheses optional inside a call

Pipelines

import json

def parse(lines):
    for line in lines:
        yield json.loads(line)

def only_errors(events):
    return (e for e in events if e["level"] == "error")

errors = only_errors(parse(read_lines("events.ndjson")))
for e in errors:
    alert(e)

itertools highlights

from itertools import islice, chain, groupby, batched, count, pairwise

first_10 = list(islice(read_lines("huge.log"), 10))

everything = chain(read_lines("a.log"), read_lines("b.log"))

for batch in batched(user_ids, 500):        # Python 3.12+
    send_bulk_email(batch)

for day, events in groupby(sorted_events, key=lambda e: e["date"]):
    print(day, sum(1 for _ in events))      # input must be sorted by the key

for prev, curr in pairwise(prices):
    print(curr - prev)

ids = count(start=1000)                     # infinite counter

Things to remember

• Generators can be consumed only once.
• len() and indexing don't work; convert to a list if you need them (and the data fits).
• yield from other_generator() delegates to another generator.
• Async code has async generators: async def with yield, consumed by async for.