AI/Biz

Python Primer·Chapter 8

Iterators, Generators and Context Managers

How Python processes things one at a time without loading everything at once, and the with statement that guarantees cleanup. This is the chapter that lets you read a file bigger than your memory.

Two ideas here, and both are about resources.

The first is laziness: producing values one at a time, on demand, rather than building the whole collection up front. That is what lets you process a fifty gigabyte file on a laptop.

The second is cleanup: making sure files get closed and connections get released, even when something goes wrong halfway through.

Iteration, Underneath

When you write a for loop, Python does something specific:

for x in [1, 2, 3]:
    print(x)

It asks the list for an iterator, then repeatedly asks that iterator for the next value until there are none left. You can do it by hand:

it = iter([1, 2, 3])
next(it)      # 1
next(it)      # 2
next(it)      # 3
next(it)      # StopIteration

You will rarely write that. But knowing the loop works this way explains two things that otherwise seem arbitrary.

An iterator is used up. Once exhausted, it is finished:

squares = (n ** 2 for n in range(3))     # a generator expression

list(squares)     # [0, 1, 4]
list(squares)     # []   ← already consumed

That catches people constantly. If you need the values twice, make a list.

Anything can be iterable. A list, a string, a dictionary, a file, a database cursor, a network response. The loop does not care where the values come from, which is why the same syntax reads a file and a list.

Generators

A generator is a function that produces values one at a time. Instead of return, it uses yield.

def count_up_to(n):
    i = 1
    while i <= n:
        yield i          # hand out a value and PAUSE here
        i += 1

for x in count_up_to(3):
    print(x)             # 1, 2, 3

When the function hits yield, it hands back a value and freezes. Its local variables stay exactly as they were. When the next value is requested, it resumes on the following line.

flowchart LR
    subgraph EAGER["A normal function"]
        E1["build the whole list"] --> E2["return it"] --> E3["all of it in memory"]
    end
    subgraph LAZY["A generator"]
        L1["yield one value"] --> L2["pause, keep state"]
        L2 -->|"next value asked for"| L1
        L2 --> L3["one value in memory<br/>at a time"]
    end
    style EAGER fill:#f0efe9,stroke:#5c6670
    style LAZY fill:#eef1fc,stroke:#2141c8

The difference in memory is the entire point:

import sys

eager = [n ** 2 for n in range(1_000_000)]      # square brackets: a list
lazy  = (n ** 2 for n in range(1_000_000))      # parentheses: a generator

sys.getsizeof(eager)     # about 8,000,000 bytes
sys.getsizeof(lazy)      # about 200 bytes

Both produce the same values. One holds a million numbers; the other holds a recipe.

Reading a large file

This is the pattern that matters, and it is why the chapter exists.

# Loads the whole file into memory. Fails on a large one.
with open('huge.csv') as f:
    lines = f.readlines()
for line in lines:
    process(line)

# Reads one line at a time. Works on any size.
with open('huge.csv') as f:
    for line in f:
        process(line)

A file object is already an iterator over its lines. The second version never holds more than one line, and it is also shorter.

Chaining generators lets you build a pipeline where nothing is materialised:

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

def parse(lines):
    for line in lines:
        yield line.split(',')

def valid_only(rows):
    for row in rows:
        if len(row) == 5:
            yield row

for row in valid_only(parse(read_lines('huge.csv'))):
    ...      # one row flows through the whole chain at a time

Each stage pulls from the one before it on demand. Memory use stays flat regardless of file size.

When not to use one

Generators are not free of downsides:

  • You cannot get the length without consuming it.
  • You cannot index into it. No gen[5].
  • It is single use.

If the data comfortably fits in memory and you need it more than once, use a list. Laziness is for when size or cost makes it necessary.

itertools

The standard library has a module of iterator tools worth knowing exists.

from itertools import islice, chain, groupby, product, count

# Take the first n items from anything, including an infinite generator
first_five = list(islice(some_generator, 5))

# Treat several sequences as one
for x in chain([1, 2], [3, 4]):
    ...

# Every combination of two lists
for city, year in product(['Delhi', 'Sydney'], [2023, 2024]):
    ...

# Count without end — useful with islice or a break
for i in count(start=1):
    if done: break

islice is the one you will reach for most. It is how you peek at the first few rows of an enormous file without reading it all:

with open('huge.csv') as f:
    for line in islice(f, 5):
        print(line)

Context Managers

The other half of this chapter. A with block guarantees that cleanup happens.

with open('data.csv') as f:
    data = f.read()
# the file is closed here, guaranteed

The guarantee holds even if an exception is raised inside the block. Compare with the manual version:

f = open('data.csv')
data = f.read()          # if this raises, the next line never runs
f.close()                # and the file stays open

Leaked file handles are a real problem in long-running processes. The operating system limits how many you can have open, and a script that leaks one per iteration will eventually fail with a message that does not obviously point at the cause.

You will meet with in several places:

with open('out.csv', 'w') as f:            # files
    ...

with sqlite3.connect('data.db') as conn:   # database connections
    ...

with pd.ExcelWriter('report.xlsx') as w:   # writing multiple sheets
    df1.to_excel(w, sheet_name='Summary')
    df2.to_excel(w, sheet_name='Detail')

Several at once, in one statement:

with open('in.csv') as src, open('out.csv', 'w') as dst:
    for line in src:
        dst.write(line.upper())

Writing your own

The easiest way is a decorated generator. Everything before the yield is setup, everything after is cleanup, and the cleanup runs no matter what.

from contextlib import contextmanager
import time

@contextmanager
def timed(label):
    start = time.perf_counter()
    try:
        yield
    finally:
        print(f"{label}: {time.perf_counter() - start:.3f}s")

with timed("loading data"):
    df = pd.read_csv('big.csv')
# loading data: 2.417s

That try/finally is what makes the guarantee work. finally always runs, whether the block completed or raised.

This pattern is genuinely useful in analysis code. A timing context manager around each stage of a pipeline tells you where the time goes with almost no effort.

Putting It Together

A realistic pattern combining everything: process a file too large to load, in batches, with timing.

from itertools import islice
from contextlib import contextmanager
import time

@contextmanager
def timed(label):
    start = time.perf_counter()
    try:
        yield
    finally:
        print(f"{label}: {time.perf_counter() - start:.2f}s")

def batched(iterable, size):
    """Yield lists of `size` items from any iterable."""
    it = iter(iterable)
    while batch := list(islice(it, size)):
        yield batch

with timed("processing"), open('huge.csv') as f:
    next(f)                                  # skip the header
    for batch in batched(f, 10_000):
        rows = [line.split(',') for line in batch]
        # ... handle 10,000 rows at a time

The := in that while loop is the walrus operator. It assigns and tests in one step, so the loop runs while batched is still producing non-empty lists.

Memory use stays at ten thousand rows regardless of whether the file has a thousand rows or a billion.

What Carries Forward

A generator produces values on demand. Parentheses instead of square brackets turns a list comprehension into one.

Iterating a file gives you lines one at a time. This is how you handle files larger than memory, and it is also the shorter way to write it.

Iterators are consumed once. If you need the values twice, build a list.

Always open files with with. The cleanup guarantee holds even when something raises.

Next: errors, modules, and the parts of the standard library that come up constantly in data work.