Iterators and generators

Part of the free LogicWiz Python course, Module 12: Iteration and the standard toolkit. The examples and practice run in the app.

The canteen's sales log has grown to a crore lines. Loading it all into a list would fill Aarav's phone memory. A generator reads one line, hands it over, and waits.

You'll learn to

• Explain how a `for` loop really works with `iter()`, `next()` and `StopIteration`
• Write an iterator class with `__iter__` and `__next__`
• Write generators with `yield` and generator expressions
• Use laziness to save memory, and remember a generator works only once

Learn

Iterable vs iterator

An iterable is anything you can loop over: a list, string, dict, file, range.

An iterator is the object that actually walks through it, one item at a time, remembering where it is.

• iter(thing) gives you an iterator
• next(it) gives the next item
• when nothing is left, next raises StopIteration

A for loop does exactly this for you, and quietly stops at StopIteration.

A for loop, by hand

tokens = ["A1", "A2"]
it = iter(tokens)
print(next(it))
print(next(it))
try:
    next(it)
except StopIteration:
    print("no more tokens")

The iterator protocol: __iter__ and __next__

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

    def __iter__(self):
        return self              # the object is its own iterator

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

for x in Countdown(3):
    print(x)

Generators: yield does the work

Writing a class is a lot of code. A generator function is a normal def that uses yield instead of return.

• Calling it does not run the body. It returns a generator object.
• Each next() runs the body until the next yield, hands that value out, and pauses there, keeping all its local variables.
• When the body ends, the generator raises StopIteration for you.

Watch the pause

def countdown(n):
    print("starting")
    while n > 0:
        yield n
        n -= 1
    print("done")

g = countdown(3)
print("made the generator")
print(next(g))
print(next(g))
print(list(g))

"starting" appears only at the first next(). Nothing runs until someone asks for a value.

Generator expressions save memory

import sys

as_list = [n * n for n in range(100000)]
as_gen = (n * n for n in range(100000))

print(sys.getsizeof(as_list), "bytes for the list")
print(sys.getsizeof(as_gen), "bytes for the generator")
print(sum(as_gen))
print(sum(x for x in range(1, 11)))

Round brackets make a generator. Inside a call like sum(...), you can drop the extra brackets.

An endless token machine

def token_numbers(start):
    n = start
    while True:          # never ends, and that's fine
        yield n
        n += 1

tokens = token_numbers(101)
print(next(tokens), next(tokens), next(tokens))

A list can never be infinite. A generator can, because it only makes the next value when asked.

A generator can be used only once. After you loop through it (or call sum, list, max on it), it is empty. A second loop gives nothing, with no error. If you need the values twice, make a list: data = list(gen), or call the generator function again.

In one line

for uses iter() and next() until StopIteration; a generator (yield or (x for x in …)) makes values lazily, saves memory and can be used only once.

Interview corner

What is the difference between an iterable and an iterator?

An iterable has __iter__ and can produce an iterator (lists, strings, dicts). An iterator has __next__ (and __iter__ returning itself) and remembers its position; next() gives the next item and raises StopIteration at the end. A list is iterable but not an iterator: next([1, 2]) is a TypeError.

What does yield do? How is a generator different from a function that returns a list?

yield makes the function a generator: calling it returns a generator object without running the body. Each next() runs to the next yield, hands out the value and pauses with its state saved. Compared with returning a list, a generator is lazy, uses constant memory, can be infinite, and can be consumed only once.

When would you choose a generator over a list?

When the data is big or endless (log files, streams), when you only need one pass (feeding sum, max, any, or another loop), or when you may stop early. Use a list when you need indexing, len(), or several passes over the same data.

Why does looping over the same generator a second time give nothing?

A generator is an iterator: it moves forward only and is exhausted after the first full pass. Later loops get StopIteration immediately, so they run zero times silently. Store the values with list(gen) or call the generator function again for a fresh one.