Iterators and generators

Two words that sound alike: iterable and iterator

An iterable is anything you can use in a for loop: a list, a string, a tuple, a dict, a set, a range, an open file. An iterator is a different, smaller object. Its only job is to give you the next item each time you ask, and to remember how far it has reached.

iter() and next()

The built-in function iter() asks an iterable for a fresh iterator. The built-in function next() asks an iterator for its next item. When nothing is left, next() raises a special exception called StopIteration. That exception is not a bug; it is the signal for "the end".

marks = [72, 88, 95]
it = iter(marks)      # get an iterator from the list
print(next(it))       # 72
print(next(it))       # 88
print(next(it))       # 95
try:
    next(it)          # nothing left
except StopIteration:
    print("No more items")

A for loop is a short way of writing the steps above. Python calls iter() on your iterable once, then calls next() again and again, and stops quietly when it sees StopIteration. Knowing this explains a lot of Python behaviour that looks like magic.

cities = ["Pune", "Jaipur", "Kochi"]
it = iter(cities)
while True:
    try:
        city = next(it)
    except StopIteration:
        break
    print(city)

A protocol is an agreed set of method names. Any object becomes an iterator if it has two dunder methods: __iter__, which returns the iterator itself, and __next__, which returns the next item or raises StopIteration. You rarely write these by hand, because generators (below) do it for you.

class Countdown:
    def __init__(self, start):
        self.current = start

    def __iter__(self):
        return self

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

print(list(Countdown(3)))   # [3, 2, 1]

Generator functions and yield

A generator function is a normal def that contains the word yield. Calling it does not run the body. Instead it returns a generator object, which is an iterator. Each next() runs the body until the next yield, hands out that value, and pauses. Local variables are kept safe until the next call.

def countdown(n):
    while n > 0:
        yield n
        n -= 1

gen = countdown(3)
print(gen)          # <generator object countdown at ...>
print(next(gen))    # 3
print(list(gen))    # [2, 1]  (the rest of the values)

Lazy evaluation means a value is worked out only when someone asks for it. A list stores all its items in memory at once. A generator stores only its current position, so it uses the same small memory for ten items or ten million items.

import sys

squares_list = [x * x for x in range(1_000_000)]   # list comprehension
squares_gen = (x * x for x in range(1_000_000))    # generator expression

print(sys.getsizeof(squares_list))   # millions of bytes
print(sys.getsizeof(squares_gen))    # around 200 bytes
print(sum(x * x for x in range(1_000_000)))

yield from some_iterable yields every item of another iterable, one by one. It is a short form of "for item in some_iterable: yield item", and it is handy when one generator is built from smaller ones.

def all_orders(shop_a, shop_b):
    yield from shop_a
    yield from shop_b

print(list(all_orders(["dosa", "idli"], ["vada pav"])))