itertools in Python: the Loops You Stop Writing

The loop you are about to write is probably already in the standard library. islice, chain, pairwise and groupby cover most of them, and groupby has one rule you cannot skip.

itertools is in the standard library, so there is nothing to install. Everything in it takes iterators and returns iterators, which means it composes with the generators from the last post and stays just as lazy.

islice: take a piece of something endless

You cannot slice a generator.

g = (n for n in range(10))
try:
    print(g[:5])
except TypeError as err:
    print('TypeError:', err)

It prints:

TypeError: 'generator' object is not subscriptable

islice is the slice that works by asking, rather than by index.

from itertools import islice

def naturals():
    n = 1
    while True:
        yield n
        n += 1

print(list(islice(naturals(), 5)))
print(list(islice(naturals(), 10, 15)))

It prints:

[1, 2, 3, 4, 5]
[11, 12, 13, 14, 15]

The arguments read like range: a stop on its own, or a start and a stop. It cannot count backwards, because there is no way to go backwards through an iterator.

chain: one loop over several things

from itertools import chain

a, b, c = [1, 2], (3, 4), range(5, 7)
print(list(chain(a, b, c)))

It prints:

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

A list, a tuple and a range, walked as one sequence, without building a combined list first.

When what you have is a list of lists, chain.from_iterable flattens one level:

from itertools import chain

rows = [['ada', 'grace'], ['alan'], ['edsger', 'barbara']]
print(list(chain.from_iterable(rows)))

It prints:

['ada', 'grace', 'alan', 'edsger', 'barbara']

pairwise: each item and the one before it

Differences between consecutive readings is a loop most people write with an index and an off-by-one.

from itertools import pairwise

temps = [12, 14, 13, 18, 18, 21]
for a, b in pairwise(temps):
    print(f"{a} -> {b}  {b - a:+d}")

It prints:

12 -> 14  +2
14 -> 13  -1
13 -> 18  +5
18 -> 18  +0
18 -> 21  +3

Five pairs from six readings, which is the number you wanted. pairwise needs Python 3.10 or newer.

groupby, and the rule you must not skip

groupby groups consecutive items. It does not gather everything with the same key from across the sequence — it starts a new group every time the key changes.

from itertools import groupby

people = [('ada', 'eng'), ('alan', 'math'), ('grace', 'eng'), ('emmy', 'math')]

for role, group in groupby(people, key=lambda p: p[1]):
    print(role, [name for name, _ in group])

It prints:

eng ['ada']
math ['alan']
eng ['grace']
math ['emmy']

Two roles in, four groups out. Nothing went wrong; that is the documented behaviour, and it is the reason people conclude groupby is broken.

Sort by the same key first and it does what you meant:

from itertools import groupby

people = [('ada', 'eng'), ('alan', 'math'), ('grace', 'eng'), ('emmy', 'math')]

def by_role(p):
    return p[1]

for role, group in groupby(sorted(people, key=by_role), key=by_role):
    print(role, [name for name, _ in group])

It prints:

eng ['ada', 'grace']
math ['alan', 'emmy']

The same function goes to sorted and to groupby. If those two ever disagree, you get the four-group result again.

There is a second trap. Each group is an iterator over the same underlying sequence, and it is only valid until you move to the next group:

from itertools import groupby

people = [('ada', 'eng'), ('grace', 'eng'), ('alan', 'math')]

groups = list(groupby(people, key=lambda p: p[1]))
for role, group in groups:
    print(role, list(group))

It prints:

eng []
math []

The list() walked all the way to the last group before anything was read, and every earlier group was left behind. Consume each group inside the loop, or build a real dict as you go.

If you only want counts, Counter from the collections post is shorter than any of this.

count, cycle, repeat

Three endless ones. They are only usable with something that stops — islice, a break, or a zip against a finite sequence.

from itertools import count, cycle, repeat, islice

print(list(islice(count(10, 5), 4)))
print(list(islice(cycle('ab'), 5)))
print(list(zip('abc', repeat(0))))

It prints:

[10, 15, 20, 25]
['a', 'b', 'a', 'b', 'a']
[('a', 0), ('b', 0), ('c', 0)]

cycle keeps a copy of everything it has seen, so it is the one member of this group that does grow.

combinations and product

Two nested loops you no longer have to nest.

from itertools import combinations, product

print(list(combinations('abc', 2)))
print(list(product([0, 1], repeat=2)))

It prints:

[('a', 'b'), ('a', 'c'), ('b', 'c')]
[(0, 0), (0, 1), (1, 0), (1, 1)]

combinations gives each unordered pair once. product is every nested-loop combination, and repeat=2 means two loops over the same sequence.

They all return iterators

Every example above is wrapped in list() for one reason: without it you get an object, not values.

from itertools import chain

c = chain([1, 2], [3])
print(c)
print(list(c))
print(list(c))

It prints:

<itertools.chain object at 0x7f9038d17e80>
[1, 2, 3]
[]

The address will differ on your machine. The empty second line will not — one pass, same as everything else in this series.

What to remember

• islice(it, n) is how you take from something infinite. Slicing syntax does not work on an iterator.

• chain walks several sequences as one; chain.from_iterable flattens a list of lists by one level.

• pairwise gives you each item with its predecessor, and gets the count right.

• groupby groups runs, not values. Sort by the same key first, and consume each group before moving on.

• Everything here is lazy and single-pass. Wrap it in list() when you actually want the values.