Iterators & Generators

16 min

Explanation

Every for loop you've written works because the object being looped over is iterable — it can produce an iterator, which yields one value at a time via __next__() until it raises StopIteration.

The easiest way to write your own iterator is a generator function: any function containing yield instead of return. Calling it doesn't run the body immediately — it returns a generator object that runs up to the next yield each time you ask it for a value.

def count_up_to(n):
    i = 1
    while i <= n:
        yield i
        i += 1

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

Try it

A generator never builds the whole sequence in memory at once — it computes each value lazily, one yield at a time. That matters a lot once 'count' is huge.

Loading editor…

Explanation

Under the hood, for x in thing: is really doing:

it = iter(thing)      # calls thing.__iter__()
while True:
    try:
        x = next(it)   # calls it.__next__()
    except StopIteration:
        break
    # loop body with x

A generator function handles all of this for you automatically. Writing a class with __iter__/__next__ by hand — like Countdown below — shows you exactly what a generator is doing behind the scenes.

Exercise

Write a generator function `squares_gen(n)` that yields the squares 0², 1², ..., (n-1)², one at a time, using `yield`.

Exercise

Complete the `Countdown` class below by implementing `__next__(self)`. It should return `self.current`, decrement it, and raise `StopIteration` once `self.current` reaches 0.

Quiz

What exception must a custom __next__ method raise to signal there are no more items?

Checkpoint

You can write a generator function with yield, and understand the iter/next/StopIteration protocol it's built on.