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))

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))   # already used up

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 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.

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))

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))

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().