Command Palette

Search for a command to run...

Hectal
PHASE 9Intermediate ~32 min· topic 4 of 9

Topic 9.4

itertools

In one line

itertools is the standard library's toolbox for iterators: infinite counters, chaining, slicing, grouping, combinations, running totals, and splitting into batches. Every tool is lazy and works on any iterable.

Think of it like this

A set of kitchen gadgets for a conveyor belt. One gadget slices the belt into pieces, one joins two belts end to end, one sorts items into trays by type, one makes every pairing of two menus. You do not build these from knives each time. itertools is the box of ready-made gadgets for streams of values.

Words you'll meet

New words in this topic, in plain English. Come back here whenever one feels fuzzy.

itertools
The standard-library module of fast, lazy tools for working with iterators.
Infinite iterator
An iterator that never ends by itself, such as count() or cycle().
Run
A stretch of neighbouring items with the same key. groupby makes one group per run.
Key function
A function that says what to group or sort by, for example len or lambda w: w[0]. A lambda is a small function with no name (Topic 5.6).
Cartesian product
Every possible pairing of items from two or more collections.
Permutation
An arrangement where order matters: ab and ba are different.
Combination
A selection where order does not matter: ab and ba are the same.
Batch
A small group of items handled together, such as 100 documents sent in one API request.

Step by step

01Counting, repeating, cycling

These three never stop, so you take a piece with islice. islice(count(10, 5), 4) means "start at 10, step 5, take 4".

main.pywhole filepython
from itertools import count, cycle, islice, repeat

print(list(islice(count(10, 5), 4)))   # [10, 15, 20, 25]
print(list(islice(cycle("AB"), 5)))    # ['A', 'B', 'A', 'B', 'A']
print(list(repeat("hi", 3)))           # ['hi', 'hi', 'hi']

02chain and islice

chain treats several collections as one stream, without copying them into a new list. islice(range(100), 2, 10, 3) means start at index 2, stop before 10, step 3. It cannot take negative numbers, because an iterator does not know where its end is.

main.pywhole filepython
from itertools import chain, islice

print(list(chain([1, 2], (3,), "ab")))      # [1, 2, 3, 'a', 'b']
print(list(islice(range(100), 2, 10, 3)))   # [2, 5, 8]

03accumulate: running totals

accumulate([1, 2, 3, 4]) gives 1, 3, 6, 10: each output is the total so far. Pass a function for something else: accumulate(data, max) is the running maximum. It is the lazy cousin of functools.reduce (Topic 9.5), which only returns the final value.

main.pywhole filepython
from itertools import accumulate

print(list(accumulate([1, 2, 3, 4])))        # [1, 3, 6, 10]
print(list(accumulate([3, 1, 4, 1, 5], max)))  # [3, 3, 4, 4, 5]

04groupby needs sorted input

groupby only merges neighbours. For ["a", "b", "a", "b", "a"] every item is its own run, giving five groups of size 1. Sorting first makes equal keys neighbours.

The break-it section below shows both results. If you need a count of each key and do not need the groups, collections.Counter (Topic 4.10) is simpler.

groupby needs sorted inputdiagram
Rendering diagram…

05product, permutations, combinations

product is a nested loop. permutations("abc", 2) lists ordered pairs without repeating an item. combinations("abcd", 2) lists unordered pairs: ("a", "b") appears but ("b", "a") does not. Counting them: write n! for n x (n-1) x ... x 1. Permutations of n items taken r at a time number n!/(n-r)!, and combinations are that divided by r!. The numbers get huge fast, so these suit small inputs, such as testing every pair of options.

06pairwise and batched

pairwise (3.10) turns [1, 4, 9, 16] into (1, 4), (4, 9), (9, 16). The differences between neighbours are [b - a for a, b in pairwise(data)].

batched (3.12) cuts a stream into tuples of at most n items. The last batch is shorter if the items do not divide evenly. n must be at least 1 or it raises ValueError. Before 3.12 people wrote this by hand with islice in a loop.

07Combine them

Because each tool takes an iterable and returns an iterator, they stack: batched(chain(file1, file2), 100) batches the lines of two files lazily, and islice(cycle(colours), 10) gives the first ten colours of a repeating palette. This is the same pipeline idea as Topic 9.3, with ready-made stages.

Try it yourself

  1. 1

    Predict the batches

    In the second example change batched(documents, 2) to batched(documents, 3). Write down the batches you expect (the last one is shorter) and then run it.

  2. 2

    Count the combinations

    Replace combinations("abcd", 2) with len(list(combinations(range(10), 3))). Check the answer against 10 x 9 x 8 / (3 x 2 x 1) = 120.

  3. 3

    Differences between neighbours

    Print [b - a for a, b in pairwise([1, 4, 9, 16])]. The result [3, 5, 7] shows how fast the numbers grow at each step.

Code & diagrams

Counting, chaining and running totals New tab

Change the step in count(10, 5) and the function passed to accumulate (try min, or operator.mul after importing operator).

Sign in to run this example in your browser.

Expected output

count:      [10, 15, 20, 25]
cycle:      ['A', 'B', 'A', 'B', 'A']
repeat:     ['hi', 'hi', 'hi']
chain:      [1, 2, 3, 'a', 'b']
accumulate: [1, 3, 6, 10]
running max: [3, 3, 4, 4, 5]
islice slice: [2, 5, 8]
Grouping, combinations and batches Python 3.12+ New tab

batched needs Python 3.12 or later (the browser runs 3.13). The batches mimic splitting documents into groups for an embedding API.

Sign in to run this example in your browser.

Expected output

a ['apple', 'avocado']
b ['banana', 'blueberry']
c ['cherry']
pairwise:      [(1, 4), (4, 9), (9, 16)]
product:       [('a', 1), ('a', 2), ('b', 1), ('b', 2)]
permutations:  [('a', 'b'), ('a', 'c'), ('b', 'a'), ('b', 'c'), ('c', 'a'), ('c', 'b')]
combinations:  [('a', 'b'), ('a', 'c'), ('a', 'd'), ('b', 'c'), ('b', 'd'), ('c', 'd')]
send to embedding API: ('d1', 'd2')
send to embedding API: ('d3', 'd4')
send to embedding API: ('d5',)
Which itertools tool?diagram
Rendering diagram…
A realistic groupby: totals per customer New tab

Orders arrive in any order, as they would from an API or a log. Without sorting, groupby splits each customer into several groups. Sorting by the same key first fixes it. operator.itemgetter("customer") is a ready-made function that fetches that key from each dictionary.

Sign in to run this example in your browser.

Expected output

unsorted groups: ['Meera', 'Ravi', 'Meera', 'Ravi', 'Asha']
Asha: 1 orders, total 300
Meera: 2 orders, total 370
Ravi: 2 orders, total 130

Break it on purpose

Errors are the best teachers. Make each change, read the error, guess what went wrong, then reveal the answer.

Break #1

Use groupby on unsorted data

Group a list whose equal items are not next to each other.

main.pywhole filepython
from itertools import groupby

tags = ["a", "b", "a", "b", "a"]
for tag, group in groupby(tags):
    print(tag, len(list(group)))

print("--- sorted first ---")
for tag, group in groupby(sorted(tags)):
    print(tag, len(list(group)))
terminal
$ python main.py
── what you'll see ──
a 1
b 1
a 1
b 1
a 1
--- sorted first ---
a 3
b 2

Break #2

Batch size zero

Call batched with a size of 0.

main.pywhole filepython
from itertools import batched

print(list(batched([1, 2, 3], 0)))
terminal
$ python main.py
── what you'll see ──
Traceback (most recent call last):
File "main.py", line 3, in <module>
print(list(batched([1, 2, 3], 0)))
~~~~~~~^^^^^^^^^^^^^^^^
ValueError: n must be at least one

Myth vs fact

Myth

groupby groups equal items like SQL GROUP BY.

Fact

It only groups neighbours. Without sorting by the same key you get several groups for one key.

Myth

itertools functions return lists.

Fact

They return lazy iterators that are used up after one pass. Wrap in list(...) to see or reuse the values.

Myth

combinations and permutations are cheap.

Fact

The number of results grows very fast. permutations(range(12)) has about 479 million items. Use them only on small inputs, or consume lazily and stop early.

Pro corner

Extra depth for experienced readers. New to this? Skip it for now and come back later.

  • ▸

    The itertools functions are written in C, so they run at the speed of built-ins and use no Python-level frames. A pipeline of chain, islice and batched is usually faster and clearer than an equivalent hand-written loop.

  • ▸

    itertools.tee(it, n) makes n independent iterators from one, but it buffers every item that one copy has read and another has not. If one copy runs far ahead, memory grows. Often it is simpler to list() the data once.

  • ▸

    batched before Python 3.12 does not exist; the old recipe is while batch := tuple(islice(it, n)): yield batch (the := walrus operator is Topic 2.9). Python 3.13 and later also let you pass strict=True to batched, which raises ValueError: batched(): incomplete batch if the last batch is short (on 3.12 the extra argument is a TypeError), so only use it if all your target versions support it. batched always yields tuples, never lists; convert with list(batch) if your API wants a list.

  • ▸

    islice consumes the source: after first = list(islice(it, 3)), the iterator it has moved on three places, and the next islice(it, 3) returns the following three items. That is useful for reading a header and then the body, and a bug when you expected a peek. Other tools worth knowing: chain.from_iterable(list_of_lists) flattens one level lazily; zip_longest pads the shorter input instead of stopping (compare zip(strict=True), Phase 2, which raises); accumulate(data, initial=0) (3.8) also yields the starting value; starmap(f, rows) is map where each row is unpacked into arguments.

  • ▸

    groupby returns groups that share the underlying iterator. Advancing to the next key invalidates the previous group, so list(groupby(data)) followed by reading the groups gives empty groups. Consume each group inside the loop.

Remember this

  1. 1

    Every itertools function returns a lazy iterator, so the tools compose and use little memory. They accept any iterable (lists, files, generators). To see the values you wrap the result in list(...) or loop over it. Remember that they are one-shot iterators, like the ones in Topic 9.1.

  2. 2

    Infinite iterators. count(start, step) counts forever (count(10, 5) gives 10, 15, 20, ...). cycle(items) repeats a sequence forever. repeat(value, times) repeats a value (forever if you leave out times). Because they never stop, pair them with islice or zip.

  3. 3

    Joining and cutting. chain(a, b, c) loops over several iterables as one long stream. islice(iterable, stop) or islice(iterable, start, stop, step) is slicing for iterators: it works on generators, where [:] does not. accumulate(items) gives running totals, and with a function (accumulate(items, max)) a running maximum or anything else.

  4. 4

    Grouping. groupby(items, key=...) gives (key, group) pairs for each run of neighbouring items with the same key. It does not sort. If equal keys are scattered, you get several groups with the same key. Sort by the same key first to get one group per key. Each group is itself a lazy iterator that is invalidated when you move to the next key, so consume it (for example list(group)) before moving on.

  5. 5

    Combinatorics. product(a, b) gives every pairing (a nested loop in one call). permutations(items, r) gives every ordered selection of r items. combinations(items, r) gives every unordered selection. They produce tuples in a predictable order. Watch the size: 20 items taken 10 at a time has 184,756 combinations, and permutations grow even faster.

  6. 6

    Pairs and batches. pairwise(items) (Python 3.10) yields overlapping neighbours (a, b), (b, c), ...: perfect for differences or checking a sequence is sorted. batched(items, n) (Python 3.12) yields tuples of up to n items, with a shorter last one. In AI code batched is how you send documents to an embedding API in groups of, say, 100 instead of one request per document.

Explain it without notes

01

What does itertools.groupby do, and what must you do to the data first?

02

Why are the itertools tools lazy, and what are the consequences?

03

What is the difference between permutations and combinations? Give the result of each for "abc" taken 2 at a time.

04

Why would you use batched in an AI application?

05

You call list(islice(it, 3)) twice on the same iterator it. What do you get, and why?

Practice

01

Use islice and count to print the first five multiples of 7 as a list.

02

Use accumulate to turn the daily sales [5, 3, 8, 2] into running totals and print them.

03

Group the words ["ant", "ape", "bee", "bat", "cow"] by first letter and print each letter with the number of words. (Hint: are the words sorted by first letter?)

04

Use pairwise to print the differences between neighbours of [2, 5, 9, 14].

Trade-offs

  • ↔

    Speed and brevity versus readability. A dense chain of itertools calls can be hard for a newcomer to read. Name the intermediate steps, and add a short comment when the stages are not obvious.

  • ↔

    Laziness versus debuggability. You cannot inspect an iterator without consuming it. While developing, wrap a stage in list(...) to look at it, then remove that.

Done when you can

  • Done when you can bound an infinite iterator with islice.

  • Done when you can explain why groupby needs sorted input, and fix the bug.

  • Done when you can tell permutations, combinations and product apart.

  • Done when you can use accumulate and pairwise for running totals and differences.

  • Done when you can cut a stream into groups with batched and say which Python version it needs.