Iterators and generators

Every for loop you have written relies on one small protocol: ask an object for an iterator, then ask the iterator for the next item until it says it has none left. Once you know that protocol, you can write your own sequences that produce values lazily, one at a time, as the loop asks for them.

Lazy sequences let you process a file larger than your memory, describe an endless stream of IDs, or chain several processing steps without building a list between them. Generators make them short to write, and the itertools module supplies ready made building blocks.

You will be able to:

• Tell an iterable from an iterator and drive one by hand with iter() and next()
• Explain what a for loop does behind the scenes
• Write a small iterator class and, more often, a generator function
• Use generator expressions, yield from and infinite generators, and reason about laziness and memory
• Choose the right itertools function for counting, slicing, chaining, grouping, windows and combinations
• Build a lazy pipeline that processes data one item at a time

Iterables and iterators

An iterable is any object you can loop over: lists, tuples, strings, dicts, sets, ranges and open files are all iterables. Calling iter() on an iterable gives you an iterator, an object that hands out the items one at a time. Calling next() on the iterator returns the next item, and when there are none left it raises StopIteration. next(iterator, default) returns default instead of raising.

Because an iterator's __iter__ returns the iterator itself, you can loop over an iterator too, but only once.

Pitfall. map, filter, zip and generators are iterators, so they are consumed by the first thing that reads them. A second pass finds nothing:

squares = map(lambda n: n * n, [1, 2, 3]) → list(squares) gives [1, 4, 9], and a second list(squares) gives [].

If you need the values twice, store them in a list first.

A range is different: it is a lazy iterable, not an iterator, so iter(range(3)) returns a new iterator each time and you can loop over the same range as often as you like.

What a for loop does

A for loop calls iter() once, then calls next() before each pass of the body. When next() raises StopIteration, the loop catches it, runs the else block if there is one, and ends. You never see the exception.

This while loop does the same work as for item in data: print(item):

Writing an iterator class

To make your own iterator, write a class with two special methods: __iter__, which returns the iterator itself, and __next__, which returns the next value or raises StopIteration. Classes are covered properly in Classes and objects; here you need only the shape.

That is a lot of ceremony for "count down", and the object has to remember by hand, in self.current, where it is between calls. A generator does the same job in a few lines and remembers its place for you.

Generators

A generator function is a function that contains yield. Calling it doesn't run the body; it returns a generator, which is an iterator. Each next() runs the body until the next yield, hands out that value and pauses, keeping all local variables. When the body finishes, the generator raises StopIteration for you.

Notice that starting appears only at the first next(), not when countdown(3) is called, and that the for loop carries on from where the manual next() calls stopped.

Generator expressions

A generator expression looks like a list comprehension in parentheses: (n * n for n in numbers). It builds a generator instead of a list. When it is the only argument of a call, you can drop the extra parentheses: sum(n * n for n in numbers).

Laziness and memory

A list holds all its items at once. A generator holds only its current state, so its size doesn't depend on how many items it will produce.

On 64 bit CPython 3.15 this printed about 8.4 million bytes for the list and about 200 for the generator; your output will differ. sys.getsizeof counts only the list's own array of references, not the int objects in it, so the real gap is even larger.

Laziness also means work happens only when someone asks for a value. any() stops at the first true item, so the generator below never examines 5 or 7:

The price of laziness: a generator has no len(), no indexing, and only one pass. Errors in the generating code appear when the values are consumed, not when the generator is created.

yield from

yield from iterable yields every item of another iterable in turn. It replaces a for loop whose body is only yield item, and it is handy for recursion:

Infinite generators

A generator doesn't have to end. An endless stream is fine as long as the consumer takes only what it needs, for example with itertools.islice, zip or break:

Pitfall. list(), sum(), sorted() or a plain for loop over an infinite generator never finishes. Always bound it first.

The itertools toolbox

The itertools module provides fast, lazy building blocks. All of them return iterators.

Counting, cycling, slicing and chaining

zip stops at its shortest input, which is what makes it safe to pair a finite list with an infinite cycle.

Running totals and grouping

accumulate yields running results. groupby yields a key and an iterator over each run of consecutive items with the same key, so you normally sort by the same key first.

Pitfall. Each group shares the underlying iterator with groupby. When groupby moves to the next key, the previous group is gone, so turn a group into a list straight away if you need to keep it.

Windows and chunks

pairwise (since 3.10) yields overlapping pairs, ideal for differences between neighbours. batched (since 3.12) cuts an iterable into tuples of n items; the last one may be shorter, and since 3.13 strict=True raises ValueError instead.

Combinatorics

These grow fast: permutations of 10 items yields 3,628,800 tuples. If you need only the count, use math.perm and math.comb.

Generator pipelines

Small generators chain into a pipeline: each stage takes an iterable, does one job and yields results to the next. Nothing runs until the last stage pulls a value, and then each line flows through all the stages before the next line is taken. Memory use stays flat however large the file is. The first stage opens the file with with open(...), which closes it again when the block ends; Errors and exceptions explains with.

Each stage is easy to test on its own with a small list, and you can swap one out, for example reading from a network stream instead of a file, without touching the others.

send and close, briefly

Generators have two more methods you will meet occasionally:

• gen.send(value) resumes the generator and makes the paused yield expression evaluate to value. You must start the generator with next(gen) (or send(None)) first; sending anything else to a fresh generator raises TypeError: can't send non-None value to a just-started generator.
• gen.close() raises GeneratorExit at the paused yield, so any finally block runs and the generator ends.

These methods were once the basis for coroutines. Modern asynchronous code uses async def and await instead (see Concurrency), so in everyday code you will mostly write plain generators.

Check yourself

1. What is the difference between an iterable and an iterator?
2. What does a for loop do when its iterator raises StopIteration?
3. nums = (n for n in range(4)). Why does sum(nums) return 6 the first time and 0 the second?
4. What happens when you call a generator function?
5. When would you choose a list comprehension over a generator expression?
6. Why do you usually sort data before passing it to groupby?
7. What is the difference between pairwise and batched?
8. What does close() do to a paused generator?

Suggested answers

1. An iterable can produce an iterator through iter(), and can usually be looped over many times. An iterator produces items through next(), remembers its position, and is used up after one pass.
2. It catches the exception, runs the loop's else block if there is one, and ends the loop.
3. A generator expression is an iterator. The first sum consumes all the values; the second finds the iterator already exhausted.
4. Nothing in the body runs yet. You get a generator object; the body runs up to the first yield on the first next().
5. When you need the values more than once, need len() or indexing, or the data is small and a list is simpler to inspect.
6. groupby groups only consecutive items with equal keys. Sorting by the same key puts all equal items next to each other, giving one group per key.
7. pairwise yields overlapping pairs, each item paired with the next one. batched yields nonoverlapping chunks of n items.
8. It raises GeneratorExit at the paused yield, so finally blocks run, and the generator finishes; further next() calls raise StopIteration.

Exercises

These exercises are optional practice. Pick the ones that help you.

1. Run length encoding. Using itertools.groupby, write encode(text) that writes each run of a repeated character as the character followed by the run length. Example: encode("aaabccdddd") gives "a3b1c2d4", and encode("") gives "". Runs are counted separately: encode("aabaa") gives "a2b1a2".
2. Fibonacci. Write an infinite generator fibonacci() that yields 0, 1, 1, 2, 3, 5, 8, ... Then write first_with_digits(n) that returns the first Fibonacci number with at least n digits. Example: the first ten values are [0, 1, 1, 2, 3, 5, 8, 13, 21, 34], and first_with_digits(3) gives 144.
3. Chunked reader. Write a generator read_chunks(path, size) that opens a UTF-8 text file and yields its content in strings of at most size characters, without reading the whole file at once. For a file containing abcdefghij and size=4, it yields "abcd", "efgh", "ij". An empty file yields nothing. A size below 1 raises ValueError.
4. Sliding window. Write windows(iterable, k) that yields tuples of k consecutive items from any iterable, including a generator. windows([1, 2, 3, 4, 5], 3) yields (1, 2, 3), (2, 3, 4), (3, 4, 5). If there are fewer than k items it yields nothing. A k below 1 must raise ValueError as soon as windows is called, not at the first next().

Solutions

1. Run length encoding.

2. Fibonacci.

3. Chunked reader.

file.read(size) returns an empty string at the end of the file, which ends the while loop. The two argument form iter(callable, sentinel) does the same: iter(lambda: file.read(size), "") calls the function until it returns "". Here the size check runs at the first next(), which is fine because the test consumes the generator; exercise 4 shows how to check eagerly.

4. Sliding window.

A generator function's body, including any argument check, runs only at the first next(). Splitting the work into a plain function that validates and a generator that yields makes the error appear at the call. deque(maxlen=k) drops the oldest item automatically when a new one arrives (see Collections and comprehensions).