Iterables and Generators

Behind the for loop

Since L4 you have written for x in something, and it works on lists, strings, dictionaries, sets, range objects, even file objects.

None of those share a common parent. How does Python treat them uniformly?

Through a convention called the iteration protocol. Any object obeying it can be used with for — which is L10’s duck typing in action.

The protocol: two functions

nums = [10, 20, 30]

it = iter(nums)          # step one: obtain an iterator
print(it)

print(next(it))          # step two: keep asking for the next one
print(next(it))
print(next(it))

nums = [10, 20, 30]
it = iter(nums)
for _ in range(3):
    next(it)
print(next(it))          # what happens once it is exhausted

Once exhausted, next() raises StopIteration.

So for x in nums: is essentially:

nums = [10, 20, 30]

# for x in nums: print(x)  unrolls to roughly this
it = iter(nums)
while True:
    try:
        x = next(it)
    except StopIteration:
        break
    print(x)

for is syntactic sugar for that. It does the same thing for every object: get an iterator, call next repeatedly, stop at StopIteration.

Iterable versus iterator

The words look alike and mean different things:

nums = [1, 2, 3]

# a list can be walked repeatedly
print(list(nums), list(nums), list(nums))

# an iterator only once
it = iter(nums)
print(list(it))
print(list(it))          # empty the second time

data = (x * 2 for x in range(5))
print("first:", sum(data))
print("second:", sum(data))

Implementing the protocol yourself

Giving a L9-style class the ability to be iterated:

class Countdown:
    """Counts down from n to 1."""

    def __init__(self, n: int) -> None:
        self.n = n

    def __iter__(self) -> "CountdownIterator":
        return CountdownIterator(self.n)


class CountdownIterator:
    def __init__(self, n: int) -> None:
        self.current = n

    def __iter__(self) -> "CountdownIterator":
        return self                       # an iterator returns itself

    def __next__(self) -> int:
        if self.current <= 0:
            raise StopIteration           # exhausted
        value = self.current
        self.current -= 1
        return value


c = Countdown(5)
for x in c:
    print(x, end=" ")
print()
for x in c:                               # traversable again, because Countdown is an iterable
    print(x, end=" ")
print()

That is a lot of ceremony — two classes and a pile of boilerplate. A generator collapses it into three lines.

Generators: yield

Replace return with yield and the function becomes a generator function:

def countdown(n: int):
    while n > 0:
        yield n                   # hand out a value, then pause
        n -= 1

for x in countdown(5):
    print(x, end=" ")
print()

Three lines, doing exactly what those two classes did.

yield versus return

return ends a function; yield pauses it.

def demo():
    print("  starting")
    yield 1
    print("  after the first yield")
    yield 2
    print("  after the second yield")
    yield 3
    print("  finishing")

g = demo()
print("the generator exists, but not one line of the body has run")
print("next →", next(g))
print("next →", next(g))
print("next →", next(g))

Look carefully at that order:

1. Calling demo() runs no code at all; it merely creates a generator object
2. The first next() starts execution, runs to the first yield, and pauses, handing out the value
3. The next next() resumes from where it paused and runs to the next yield
4. When the function ends it raises StopIteration

Local variables survive across the pause — the most remarkable thing about generators.

def counter():
    count = 0
    while True:
        count += 1        # count is remembered between calls to next
        yield count

g = counter()
print(next(g), next(g), next(g), next(g))

What a generator saves

This is the point of the lecture.

import sys

# a list: every element built immediately and held in memory
squares_list = [x ** 2 for x in range(100_000)]

# a generator: nothing built, computed on demand
squares_gen = (x ** 2 for x in range(100_000))

print("list uses:", sys.getsizeof(squares_list), "bytes")
print("generator uses:", sys.getsizeof(squares_gen), "bytes")
print("same total:", sum(squares_list) == sum(squares_gen))

A hundred thousand elements costs hundreds of kilobytes as a list and a couple of hundred bytes as a generator.

Because a generator stores no data, only the state needed to compute the next value. Which explains something you met long ago:

import sys

r = range(10_000_000)
print("a range object uses:", sys.getsizeof(r), "bytes")
print("it stands for ten million numbers and costs almost nothing")
print("because it is not a list:", type(r))
print("computed on demand:", r[5_000_000])

L4 said “a range is not a list”. Now you know why.

Handling data larger than memory

def read_numbers(lines: list[str]):
    """Yield numbers line by line, without loading the file."""
    for line in lines:
        line = line.strip()
        if line:
            yield int(line)

# stand in for a file
fake_file = ["1", "2", "", "3", "4", "5"]

total = 0
for n in read_numbers(fake_file):
    total += n
print("total:", total)

In practice lines is a file object — and a file object is itself an iterator, which is exactly why L12 could say that for line in f uses memory independent of file size.

Infinite sequences

A generator can represent an endless sequence, because it never produces all of it at once:

def naturals():
    n = 1
    while True:
        yield n
        n += 1

def take(gen, k: int) -> list[int]:
    """Take the first k values from a generator."""
    result = []
    for x in gen:
        result.append(x)
        if len(result) == k:
            break
    return result

print(take(naturals(), 10))

def fibonacci():
    a, b = 0, 1
    while True:
        yield a
        a, b = b, a + b

fib = fibonacci()
print([next(fib) for _ in range(15)])

A list cannot do this — [n for n in naturals()] runs until memory is gone.

Generator expressions

Almost identical to a list comprehension, with parentheses instead of brackets:

nums = range(1, 11)

as_list = [x ** 2 for x in nums]          # comprehension: everything at once
as_gen = (x ** 2 for x in nums)           # generator expression: on demand

print(as_list)
print(as_gen)
print(list(as_gen))

As a function’s only argument the outer parentheses can be dropped:

nums = range(1, 1001)

print(sum(x ** 2 for x in nums))                    # parentheses omitted
print(max(len(w) for w in ["a", "abc", "ab"]))
print(any(x > 900 for x in nums))
print(all(x > 0 for x in nums))

any and all with a generator have another benefit — short-circuiting:

def checked(x: int) -> bool:
    print(f"  checking {x}")
    return x > 3

print("any:", any(checked(x) for x in [1, 2, 3, 4, 5, 6]))
print("→ it stopped at the first match; the rest were never examined")

The comprehension family

words = ["apple", "banana", "cherry", "date"]

print([w.upper() for w in words])                     # list comprehension
print({w[0] for w in words})                          # set comprehension
print({w: len(w) for w in words})                     # dict comprehension
print(tuple(len(w) for w in words))                   # tuples need tuple()

Common traps

1. A generator works once

gen = (x for x in range(5))
print(list(gen))
print(list(gen))          # empty

If you need it more than once, materialize it: data = list(gen).

2. Laziness delays exceptions

def risky():
    yield 1
    yield 2
    raise ValueError("something went wrong")
    yield 3

g = risky()
print("no error when the generator was created")
print(next(g))
print(next(g))
print(next(g))            # only now

An error in the body only surfaces when that line actually executes, which makes debugging slightly harder.

3. Late binding of the loop variable

gens = [(x * i for x in range(3)) for i in range(3)]
print([list(g) for g in gens])

This works, but the equivalent with lambdas does not — that trap belongs to L14.

itertools: iteration tools off the shelf

An entire standard library module exists for composing iterators:

import itertools

print(list(itertools.islice(itertools.count(10, 2), 5)))     # from 10, step 2, take 5
print(list(itertools.islice(itertools.cycle("ABC"), 7)))     # repeat forever
print(list(itertools.repeat("x", 3)))                        # three times

import itertools

print(list(itertools.chain([1, 2], [3, 4], [5])))            # concatenate sequences
print(list(itertools.accumulate([1, 2, 3, 4, 5])))           # running totals
print(list(itertools.combinations("ABC", 2)))                # combinations
print(list(itertools.permutations("ABC", 2)))                # permutations
print(list(itertools.product([1, 2], "AB")))                 # cartesian product

groupby groups adjacent items, and requires sorted input:

import itertools

data = [("fruit", "apple"), ("fruit", "banana"), ("veg", "cabbage"), ("veg", "carrot")]

for key, group in itertools.groupby(data, key=lambda x: x[0]):
    print(key, "→", [item[1] for item in group])

import itertools

data = ["apple", "cabbage", "banana"]
key = lambda w: "fruit" if w in ("apple", "banana") else "veg"

print("unsorted:", [(k, list(g)) for k, g in itertools.groupby(data, key=key)])
print("sorted:  ", [(k, list(g)) for k, g in itertools.groupby(sorted(data, key=key), key=key)])

collections: better containers

from collections import Counter

# Counter: counting, without the dict.get(k, 0) + 1 idiom
words = "the quick brown fox the lazy dog the end".split()
c = Counter(words)
print(c)
print(c.most_common(2))
print(c["the"], c["nonexistent"])      # missing keys give 0, not KeyError

from collections import defaultdict

# defaultdict: default values created automatically, no setdefault needed
groups = defaultdict(list)
for word in ["apple", "avocado", "banana", "blueberry"]:
    groups[word[0]].append(word)      # no need to check whether the key exists
print(dict(groups))

That is precisely the KeyError from L6’s exercise 6.

from collections import deque

# deque: efficient at both ends (list.insert(0,x) is O(n); deque is O(1))
d = deque([1, 2, 3])
d.appendleft(0)
d.append(4)
print(d)
print(d.popleft(), d.pop())
print(d)

import time
from collections import deque

n = 50_000

lst = []
t = time.time()
for i in range(n):
    lst.insert(0, i)
list_t = time.time() - t

dq = deque()
t = time.time()
for i in range(n):
    dq.appendleft(i)
deque_t = time.time() - t

print(f"list.insert(0, x) x{n}: {list_t:.4f} s")
print(f"deque.appendleft  x{n}: {deque_t:.4f} s")
print(f"deque is about {list_t / deque_t:.0f}x faster")

This is L11’s “change the data structure” in practice: O(n) becomes O(1).

Unpacking revisited

L6 covered tuple unpacking. * and ** work at call sites too:

def add(a: int, b: int, c: int) -> int:
    return a + b + c

nums = [1, 2, 3]
print(add(*nums))                # spread the list into three arguments

opts = {"a": 1, "b": 2, "c": 3}
print(add(**opts))               # spread the dict into keyword arguments

a, *rest = [1, 2, 3, 4]
print(a, rest)

first, *middle, last = [1, 2, 3, 4, 5]
print(first, middle, last)

print([*"abc", *[1, 2]])         # spread into a list literal
print({**{"a": 1}, **{"b": 2}})  # spread into a dict (the merge from L6)

Summary

• for is sugar: iter() for an iterator, repeated next(), stop on StopIteration
• Iterables can be traversed repeatedly; iterators cannot
• yield pauses a function; return ends it. Locals survive the pause
• A generator stores no data, only how to compute the next value — memory is independent of size
• range objects and file objects work the same way
• Generators can represent infinite sequences
• Generator (...) versus comprehension [...]: traverse once → generator
• any/all with a generator short-circuit
• There is no tuple comprehension; (x for x in ...) is a generator
• itertools: islice, chain, accumulate, combinations, groupby (sort first)
• collections: Counter for counting, defaultdict to skip the check, deque for O(1) at both ends

Exercises

1. Write a generator evens(n) yielding the first n even numbers, and a second squares(gen) that takes a generator and yields each value squared. Compose them.
2. Write a generator that reads a file and yields every non-blank line, stripped. Why is a generator better than returning a list here?
3. Why is the second print empty? Fix it:

   data = (x for x in range(5))
   print(sum(data))
   print(max(data))

4. Write an infinite generator primes() and take the first 20 with itertools.islice.
5. Rewrite Lab 6’s word frequency using collections.Counter and compare the line counts.
6. Use deque to build a “last N records” buffer that discards the oldest automatically once it exceeds N. (Hint: deque(maxlen=N).)
7. Explain the memory difference between these and verify it with sys.getsizeof:

   a = [x ** 2 for x in range(1_000_000)]
   b = (x ** 2 for x in range(1_000_000))

The lab problems are in Lab 13.