itertools and functools

Two standard-library modules hold most of the tools that people write by hand and get slightly wrong. itertools works on iterators, so everything in it is lazy and fits the generator pipelines of the last lessons. functools works on functions: it builds new functions from old ones. This lesson is a tour of the tools you will actually use, and it leaves out the rest.

itertools: joining and slicing

chain walks several iterables one after another as if they were one. chain.from_iterable does the same for an iterable of iterables, which is how you flatten one level of nesting:

from itertools import chain, islice

weekdays = ["mon", "tue"]
weekend = ("sat", "sun")
print(list(chain(weekdays, weekend, "x")))

lists = [[1, 2], [3], [4, 5]]
print(list(chain.from_iterable(lists)))

['mon', 'tue', 'sat', 'sun', 'x']
[1, 2, 3, 4, 5]

islice slices an iterator the way [start:stop:step] slices a list. It is the answer to “take the first five items of a generator”, which ordinary slicing cannot do:

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


print(list(islice(naturals(), 5)))
print(list(islice(naturals(), 3, 6)))
print(list(islice("abcdefgh", 0, None, 3)))

[1, 2, 3, 4, 5]
[4, 5, 6]
['a', 'd', 'g']

The arguments follow slice rules: one number is the stop, two are start and stop, and None for the stop means “to the end”. There are no negative values, since an iterator cannot know how long it is. Slicing consumes the iterator: items skipped by islice are gone.

itertools: groupby

groupby collects consecutive items that share a key. The word consecutive matters, and it is where nearly everyone trips first:

from itertools import groupby

words = ["apple", "avocado", "banana", "blueberry", "cherry", "apricot"]

for letter, group in groupby(words, key=lambda w: w[0]):
    print(letter, list(group))

a ['apple', 'avocado']
b ['banana', 'blueberry']
c ['cherry']
a ['apricot']

apricot starts a second a group because it is not next to the other a words. If you want one group per key, sort by the same key first:

for letter, group in groupby(sorted(words), key=lambda w: w[0]):
    print(letter, list(group))

a ['apple', 'apricot', 'avocado']
b ['banana', 'blueberry']
c ['cherry']

Each group is itself a lazy iterator that shares the underlying data. Use it, or turn it into a list, before the loop moves to the next key; afterwards it is empty.

Consecutive pairs

3.10+Requires version 3.10 or later pairwise yields each item together with the one after it. It replaces the old zip(items, items[1:]) and works on any iterable, including a generator:

from itertools import pairwise

temps = [18, 21, 20, 25]
print(list(pairwise(temps)))
print([b - a for a, b in pairwise(temps)])

[(18, 21), (21, 20), (20, 25)]
[3, -1, 5]

itertools.pairwise was added in Python 3.10, according to the itertools documentation. On 3.8 or 3.9, write zip(items, items[1:]) for lists.

Fixed-size batches

3.12+Requires version 3.12 or later batched cuts an iterable into tuples of a given length. The last tuple is shorter if the items do not divide evenly. It is the tool for sending 100 ids at a time to an API:

from itertools import batched

ids = range(1, 8)
for batch in batched(ids, 3):
    print(batch)

(1, 2, 3)
(4, 5, 6)
(7,)

itertools.batched was added in Python 3.12, per the itertools documentation. On older versions, this four-line function does the same job with islice:

from itertools import islice


def batched(iterable, n):
    """Same job as itertools.batched, for Python before 3.12."""
    iterator = iter(iterable)
    while batch := tuple(islice(iterator, n)):
        yield batch

The loop keeps taking n items from one shared iterator until nothing is left, at which point the empty tuple is false and the loop stops.

functools: partial

partial takes a function and fixes some of its arguments, returning a new function that needs only the rest. Use it where you would otherwise write a small lambda that only forwards arguments:

from functools import partial


def power(base, exponent):
    return base**exponent


square = partial(power, exponent=2)
cube = partial(power, exponent=3)
print(square(5), cube(2))

25 8

It also helps when an API wants a function of one argument, such as map or a callback, and you have a function of two:

to_binary = partial(int, base=2)
print(to_binary("1011"))
print(list(map(to_binary, ["1", "10", "11"])))
print(square.func.__name__, square.keywords)

11
[1, 2, 3]
power {'exponent': 2}

The .func, .args and .keywords attributes let you see what a partial object wraps.

functools: reduce

reduce folds a sequence into one value by applying a two-argument function repeatedly: first to the first two items, then to that result and the third item, and so on.

from functools import reduce
import operator

numbers = [1, 2, 3, 4]
print(reduce(operator.add, numbers))
print(reduce(operator.mul, numbers))
print(reduce(lambda total, n: total * 10 + n, numbers))
print(reduce(operator.add, [], 0))

10
24
1234
0

The third line builds the number 1234 from its digits: ((1*10+2)*10+3)*10+4. The last line shows the third argument, the starting value, which is what you get for an empty sequence. Without it, reduce raises TypeError on an empty input.

Many folds already have a clearer name: sum, max, min, all, any, "".join, math.prod. Use reduce only when none of them fits, and consider a plain for loop if the function is long.

functools: caching results

A function whose result depends only on its arguments can remember answers. lru_cache stores them in a dictionary keyed by the arguments, so the arguments must be hashable. A bare @lru_cache without parentheses works from Python 3.8, and maxsize limits how many results are kept, discarding the least recently used first:

from functools import lru_cache


@lru_cache(maxsize=2)
def double(n):
    print("computing", n)
    return n * 2


double(1)
double(2)
double(1)
double(3)
double(2)
print(double.cache_info())

computing 1
computing 2
computing 3
computing 2
CacheInfo(hits=1, misses=4, maxsize=2, currsize=2)

The second double(1) was a hit and printed nothing. After double(3), the cache was full, and 2 was the least recently used entry, so it was discarded and had to be computed again at the end.

cache

3.9+Requires version 3.9 or later If you do not want a size limit, @functools.cache is the shorter spelling of an unbounded lru_cache. It was added in Python 3.9 (see the functools documentation); on 3.8 use @lru_cache(maxsize=None).

from functools import cache

calls = 0


@cache
def slow_square(n):
    global calls
    calls += 1
    return n * n


print(slow_square(4), slow_square(4), slow_square(5))
print("calls:", calls)
print(slow_square.cache_info())

16 16 25
calls: 2
CacheInfo(hits=1, misses=2, maxsize=None, currsize=2)

Two real calls, one cache hit. Call slow_square.cache_clear() to empty the cache. Cache only functions that are pure: with no side effects, and with the same result every time for the same arguments. An unbounded cache that sees millions of different arguments will fill your memory.

functools: singledispatch

Sometimes one operation needs different code for different argument types, and a chain of isinstance checks is clumsy. singledispatch turns a function into a family: the original is the fallback, and register adds versions chosen by the type of the first argument:

from functools import singledispatch


@singledispatch
def describe(value):
    return f"something else: {value!r}"


@describe.register
def _(value: int):
    return f"the integer {value}"


@describe.register(list)
def _(value):
    return f"a list of {len(value)} items"


@describe.register(str)
def _(value):
    return f"the text {value!r}"


print(describe(42))
print(describe([1, 2, 3]))
print(describe("hi"))
print(describe(2.5))
print(describe(True))

the integer 42
a list of 3 items
the text 'hi'
something else: 2.5
the integer True

The type can be given as an annotation on the first parameter (as for int) or as an argument to register (as for list and str). The last line shows that subclasses are matched too: bool is a subclass of int. Registering new types does not require editing the original function, which makes this useful for extension points. The underscore name _ is a convention for a function you never call by that name.

Common pitfalls

• groupby without sorting. It only groups neighbours, so unsorted data gives repeated keys.
• Reading a groupby group late. Convert each group to a list inside the loop, or you will see empty groups.
• Caching impure functions or unhashable arguments. A cached function returns stale results if it depends on something else, such as the time, and raises TypeError: unhashable type: 'list' if given a list.
• Using @cache on methods. The cache keeps a reference to self, so instances are never freed. Prefer module-level functions, or functools.cached_property for per-instance values.
• Reaching for reduce first. If sum, max or join fits, use it.

3 calls
['user1', 'user2', 'user3', 'user4', 'user5', 'user6', 'user7', 'user8', 'user9', 'user10']

user_ids = list(range(1, 11))
calls = 0


def fetch_users(batch):
    """Pretend API call that accepts at most 4 ids at once."""
    global calls
    calls += 1
    assert len(batch) <= 4
    return [f"user{i}" for i in batch]

from itertools import batched

user_ids = list(range(1, 11))
calls = 0


def fetch_users(batch):
    """Pretend API call that accepts at most 4 ids at once."""
    global calls
    calls += 1
    assert len(batch) <= 4
    return [f"user{i}" for i in batch]


users = []
for batch in batched(user_ids, 4):
    users.extend(fetch_users(batch))

print(calls, "calls")
print(users)

3 calls
['user1', 'user2', 'user3', 'user4', 'user5', 'user6', 'user7', 'user8', 'user9', 'user10']

55
23416728348467685
81

def fib(n):
    if n < 2:
        return n
    return fib(n - 1) + fib(n - 2)

from functools import cache


@cache
def fib(n):
    if n < 2:
        return n
    return fib(n - 1) + fib(n - 2)


print(fib(10))
print(fib(80))
print(fib.cache_info().misses)

55
23416728348467685
81

from itertools import groupby

for key, group in groupby("aabaa"):
    print(key, len(list(group)))

a 2
b 1
a 2

Summary

• itertools gives lazy tools: chain joins, islice slices, groupby groups neighbours. pairwise is 3.10+Requires version 3.10 or later and batched is 3.12+Requires version 3.12 or later.
• functools.partial pre-fills arguments, reduce folds a sequence, and lru_cache (or 3.9+Requires version 3.9 or later cache) remembers results of pure functions.
• singledispatch selects a function by the type of its first argument.
• Next: closures and the lambda-in-a-loop bug.
• Going further: the advanced lesson The collections module continues with more specialised containers such as deque and ChainMap.