Iterators and Generators

You've used for loops on lists, strings, dictionaries, files and range(). How does one kind of loop work with so many different things? The answer is iterators. Understanding them leads to one of Python's most useful features: generators, which produce values one at a time instead of building them all up front.

Iterables and iterators

Anything you can loop over is called an iterable: lists, strings, tuples, dictionaries, sets, files and so on.

An iterator is the object that actually walks through an iterable, handing out one item at a time. You get one with iter(), and ask it for the next item with next():

colors = ["red", "green", "blue"]

it = iter(colors)
print(next(it))
print(next(it))
print(next(it))

The iterator remembers its position: each next(it) gives the item after the last one, so you get red, then green, then blue. The list itself doesn't change. You can't call next() on the list directly; next(colors) is a TypeError, because a list isn't an iterator. iter() makes one for it.

When there are no items left, next() raises a StopIteration exception. You can catch it like any other exception:

it = iter(["only one"])
print(next(it))

try:
    next(it)
except StopIteration:
    print("StopIteration: no items left")

What a for loop really does

A for loop is just a tidy way of doing exactly that. Behind the scenes:

1. It calls iter() on what you're looping over.
2. It calls next() over and over, running the body for each item.
3. It stops quietly when StopIteration is raised.

word = "hey"

# what you write
for letter in word:
    print(letter)

# roughly what Python does for you
it = iter(word)
while True:
    try:
        letter = next(it)
    except StopIteration:
        break
    print(letter)

An iterator can only be used once. Once it has handed out every item, it's empty:

numbers = iter([1, 2, 3])

print(list(numbers))
print(list(numbers))   # used up

The second line prints []. So why can you loop over the same list as many times as you like? Because each for loop calls iter() again and gets a fresh iterator. The list holds the data; the iterator is just a bookmark.

Generators: functions that yield

A generator is the easiest way to make your own iterator. It looks like a normal function, but it uses yield instead of return:

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

for number in count_up_to(5):
    print(number)

The loop asks the generator for a value, and the function runs until it reaches yield n, which hands out 1 and pauses. The next request carries on after the yield: n += 1, back to the while check, and yield again with 2. When n reaches 6, the while loop ends, the function reaches its end, and that tells the for loop to stop (Python raises StopIteration for you).

The difference is in when the code runs:

• return sends back one value and the function is finished.
• yield sends back a value and pauses the function. The next time a value is asked for, it carries on from exactly where it stopped, with all its variables intact.

Adding print calls makes the pausing visible:

def steps():
    print("  starting")
    yield 1
    print("  resumed after 1")
    yield 2
    print("  resumed after 2")
    yield 3
    print("  finished")

gen = steps()
print("got", next(gen))
print("got", next(gen))
print("got", next(gen))

Calling steps() doesn't run any code yet. It just creates a generator object, and the body runs bit by bit as you ask for values. Notice that starting only appears with the first next(), and finished never appears at all: nobody asked for a fourth value, so the function stayed paused after yield 3. A fourth next(gen) would print finished and then raise StopIteration.

A generator is an iterator, so it can only be used once. To go through the values again, call the function again to get a new generator.

Why use generators?

They save memory

A list holds every value at once. A generator produces each value only when it's needed, so it uses almost no memory, however many values it produces:

import sys

big_list = [n * n for n in range(1_000_000)]
big_gen = (n * n for n in range(1_000_000))

print(sys.getsizeof(big_list), "bytes for the list")
print(sys.getsizeof(big_gen), "bytes for the generator")
print(sum(big_gen))

The list needs megabytes, because all million squares exist at once. The generator needs only a couple of hundred bytes, because it holds just its current position. sum(big_gen) still gets every value, one at a time. The exact numbers depend on your Python version. The trade-off: a list lets you index it (big_list[10]) and go through it again, and a generator doesn't.

They can go on forever

A generator doesn't need an end. This one produces Fibonacci numbers endlessly, and the loop decides when to stop:

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

for n in fibonacci():
    if n > 100:
        break
    print(n, end=" ")

A list could never hold an endless sequence, but a generator can, because it only ever works out the next value.

Generator expressions

A generator expression looks like a list comprehension, but with round brackets instead of square ones:

squares_list = [n * n for n in range(5)]
squares_gen = (n * n for n in range(5))

print(squares_list)
print(squares_gen)
print(list(squares_gen))

Printing the generator shows only <generator object ...>, because no values have been worked out yet. list() asks for all of them. Like any generator, it's used up afterwards, so a second list(squares_gen) gives [].

They're perfect for passing straight into functions like sum(), max() and any(). You can even drop the extra brackets:

prices = [4.99, 12.50, 3.25, 8.00]

print(round(sum(p * 1.2 for p in prices), 2))
print(max(len(w) for w in ["cat", "giraffe", "ox"]))
print(any(p > 10 for p in prices))

from itertools import islice, count

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

print(list(islice(fibonacci(), 10)))
print(list(islice(count(100, 5), 4)))

def all_items(*groups):
    for group in groups:
        yield from group

print(list(all_items([1, 2], "ab", (True,))))

Exercises

Exercise 1: Step through by hand

Use iter() and next() to print the first two items of the list, without a for loop. It should print Mon and then Tue.

days = ["Mon", "Tue", "Wed"]

Exercise 2: Even numbers generator

Write a generator evens(limit) that yields the even numbers from 0 up to and including limit. The loop below should print 0 2 4 6 8 10.

def evens(limit):
    pass

for n in evens(10):
    print(n, end=" ")

Exercise 3: Countdown generator

Write a generator countdown(n) that yields n, n - 1, … down to 1. Then print the list it produces for countdown(5): [5, 4, 3, 2, 1].

def countdown(n):
    pass

Exercise 4: Sum with a generator expression

Use a single generator expression inside sum() to add up the lengths of all the words. It should print 23.

words = ["generators", "are", "really", "neat"]

Iteration tools at a glance

A quick reference for later. The itertools examples need import itertools first.

Summary

• An iterable is anything you can loop over. An iterator hands out its items one at a time.
• iter() gets an iterator, next() gets the next item, and StopIteration means there are none left.
• A for loop calls iter() and next() for you, and an iterator can only be used once.
• A function with yield is a generator: it pauses at each yield and resumes on the next request.
• Generators save memory and can even be endless.
• (expr for x in items) is a generator expression. It's perfect inside sum(), max() and any().