Iterators and generators

Lesson 29 of 31 · Part 6 · Going further · about 20 minutes

You have looped over lists, strings, dictionaries, ranges and files. A for loop works on all of them because they share one simple protocol: "give me the next item". This lesson shows that protocol, then a way to write your own sequences with a single keyword, yield.

Iterables and iterators

Anything a for loop can walk through is an iterable: a list, a string, a dictionary, a range. Give an iterable to iter() and you get an iterator: an object that hands out the items one at a time. Each call to next() gives the next one.

main.py

1lamps = ["north gate", "mill", "bridge"]2it = iter(lamps)3print(next(it))4print(next(it))5print(next(it))

Output

north gate
mill
bridge

An iterator remembers where it is, and it only goes forward. When there is nothing left, next() raises StopIteration:

main.py

1it = iter("ab")2print(next(it))3print(next(it))4print(next(it))

Output

a
b
Traceback (most recent call last):
  File "main.py", line 4, in <module>
    print(next(it))
StopIteration

	Iterable	Iterator
Examples	list, str, dict, range	what `iter()` returns; `map` and `filter` objects
What it does	can be looped over, again and again	hands out items one at a time with `next()`
Can you start again?	yes, each loop starts from the beginning	no, once used up it stays empty

What a for loop does

A for loop is iter() and next() with the details hidden. It asks for an iterator once, calls next() before every pass, and stops quietly when StopIteration comes. Written out by hand, the same loop looks like this:

main.py

1lamps = ["north gate", "mill", "bridge"]2it = iter(lamps)3while True:4    try:5        lamp = next(it)6    except StopIteration:7        break8    print("Lighting", lamp)

Output

Lighting north gate
Lighting mill
Lighting bridge

You never need to write it like this; the for loop does it for you. But it explains why an iterator works only once: a second loop over the same iterator asks for more, and there is nothing left.

Guess first

What does the last line print?

1it = iter([1, 2, 3])2for n in it:3    print(n)4print(list(it))

The for loop took every item from the iterator. list(it) asks the same iterator for more, gets StopIteration straight away, and makes an empty list. A list, unlike an iterator, could be looped over again.

Output

Generators and yield

Writing an iterator by hand is work. A generator function does it for you: it is an ordinary def with yield in it. Calling it does not run the body; it gives back a generator, which is an iterator. Each next() runs the body until the next yield, hands out that value, and pauses right there.

main.py

1def countdown(n):2    print("Starting")3    while n > 0:4        yield n5        n -= 16    print("Done")7 8gen = countdown(3)9print(next(gen))10print(next(gen))11print("Back in the main program")12print(next(gen))

Output

Starting
3
2
Back in the main program
1

Press Step and watch the current line jump back into countdown on every next(), then out again at yield. The local variable n keeps its value while the generator is paused. "Done" is never printed: the body would only reach it on a fourth next().

Usually you just loop over a generator, or hand it to list(), sum() or anything else that takes an iterable:

main.py

1def countdown(n):2    while n > 0:3        yield n4        n -= 15 6for n in countdown(3):7    print(n)8print(list(countdown(5)))9print(sum(countdown(4)))

Output

3
2
1
[5, 4, 3, 2, 1]
10

Printing the generator itself would only show something like <generator object countdown at 0x...>, with a memory address that changes every run. To see the values, loop over it or wrap it in list().

Guess first

In what order are the lines printed?

1def lamps():2    yield "mill"3    print("between")4    yield "bridge"5 6for name in lamps():7    print(name)

The loop body prints mill while the generator is paused at the first yield. The next pass resumes it: it prints between, then yields bridge.

Output

Why generators help

A list holds all of its items in memory at once. A generator holds none of them: it makes each value only when asked, then forgets it. To add up the squares of a hundred thousand numbers, you don't need a list of a hundred thousand squares:

main.py

1def squares(limit):2    for n in range(limit):3        yield n * n4 5print(sum(squares(100_000)))

Output

333328333350000

sum() asks for one square, adds it, and asks for the next. At any moment there is only one square in memory. With a list it would be a hundred thousand numbers, all stored before the adding even begins. For a big file or a long stream of data, that difference decides whether a program runs at all.

Endless generators

Because a generator makes values only on demand, it doesn't have to stop. This one hands out lamp numbers forever. That is fine, as long as whoever loops over it stops:

main.py

1def lamp_ids():2    n = 13    while True:4        yield f"L{n:03}"5        n += 16 7for lamp_id in lamp_ids():8    if lamp_id == "L004":9        break10    print(lamp_id)

Output

L001
L002
L003

Never write list(lamp_ids()): list() would keep asking and the program would never finish. To take just the first few items, the itertools module (it comes with Python) has islice, a slice for iterators. itertools.count is a ready-made endless counter:

main.py

1import itertools2 3def lamp_ids():4    n = 15    while True:6        yield f"L{n:03}"7        n += 18 9print(list(itertools.islice(lamp_ids(), 3)))10print(list(itertools.islice(itertools.count(10, 5), 4)))

Output

['L001', 'L002', 'L003']
[10, 15, 20, 25]

Generator expressions

A list comprehension in parentheses instead of square brackets is a generator expression. It works the same way, but makes its values one at a time, like a generator function, instead of building a list.

main.py

1loads = [3, 8, 2, 6, 11]2print(sum(n * 2 for n in loads))3print(max(len(word) for word in ["ore", "lamp", "warehouse"]))4print(any(n > 10 for n in loads))

Output

60
9
True

any() gives True if at least one item is true, so the last line asks whether any load is over 10. When a generator expression is the only argument, its parentheses can share the function's: sum(n * 2 for n in loads). Like every iterator, a generator expression can be used only once:

main.py

1doubled = (n * 2 for n in [1, 2, 3])2print(sum(doubled))3print(sum(doubled))

Output

12
0

Use square brackets when you need the list itself: to print it, index it or loop over it twice. Use parentheses when the values go straight into sum(), max(), any() or a loop.

Practice

Write the code, then press Check. A hint costs nothing; the answer is there when you need it.

Exercise 1

One at a time

Use iter() and next() to print the first two items of ores, each on its own line. Don't use a loop.

It should print

iron
gold

exercise1.py

1ores = ["iron", "gold", "copper"]2# Your code here

Output

Press Run or Ctrl+Enter

Exercise 2

Even numbers

Finish the generator evens so that it yields the even numbers from 0 up to and including limit. The last line should then print the list below.

It should print

[0, 2, 4, 6, 8, 10]

exercise2.py

1def evens(limit):2    pass3 4print(list(evens(10)))

Output

Press Run or Ctrl+Enter

Exercise 3

Heavy loads

Print the total of the loads that are larger than 5. Use sum() with a generator expression, not a list.

It should print

exercise3.py

1loads = [3, 8, 2, 6, 11]2print(sum(loads))

Output

Press Run or Ctrl+Enter

What you learned

An iterable can be looped over; iter() turns it into an iterator, and next() takes one item at a time.
When an iterator is used up, next() raises StopIteration; a for loop catches it and stops.
A function with yield is a generator function: each next() runs it to the next yield and pauses there.
Generators make values only when asked, so they save memory and can even go on forever; stop them with break or itertools.islice.
(x for x in ...) is a generator expression; like every iterator, it works only once.