itertools & functools

itertools provides fast building blocks for working with iterators — combining, grouping, slicing and generating sequences lazily. functools provides tools for working with functions — caching, partial application and more. Both replace a lot of hand-written loops.

Infinite iterators

PY infinite.py
from itertools import count, cycle, repeat, islice

print(list(islice(count(10, 5), 4)))
print(list(islice(cycle(["red", "green"]), 5)))
print(list(repeat("x", 3)))

for ticket, name in zip(count(100), ["Ana", "Ben", "Cleo"]):
    print(ticket, name)
Output
[10, 15, 20, 25]
['red', 'green', 'red', 'green', 'red']
['x', 'x', 'x']
100 Ana
101 Ben
102 Cleo

count, cycle and repeat can go on forever, so they're combined with something that stops — islice (slicing for iterators) or zip with a finite sequence.

Combining iterables

PY combining.py
from itertools import chain, zip_longest, pairwise, batched

print(list(chain([1, 2], (3, 4), "ab")))
print(list(chain.from_iterable([[1, 2], [3], [4, 5]])))
print(list(zip_longest("abc", [1, 2], fillvalue="-")))
print(list(pairwise([10, 13, 19, 30])))
print(list(batched(range(7), 3)))
Output
[1, 2, 3, 4, 'a', 'b']
[1, 2, 3, 4, 5]
[('a', 1), ('b', 2), ('c', '-')]
[(10, 13), (13, 19), (19, 30)]
[(0, 1, 2), (3, 4, 5), (6,)]
  • chain joins iterables end to end; chain.from_iterable flattens one level.
  • zip_longest keeps going until the longest input is done.
  • pairwise gives consecutive pairs — handy for differences: [b - a for a, b in pairwise(values)].
  • batched splits into chunks — perfect for sending data in batches to an API.

Filtering and accumulating

PY filtering.py
from itertools import takewhile, dropwhile, filterfalse, compress, accumulate
import operator

temps = [12, 15, 19, 23, 21, 17, 25]
print(list(takewhile(lambda t: t < 20, temps)))
print(list(dropwhile(lambda t: t < 20, temps)))
print(list(filterfalse(lambda t: t % 2, temps)))
print(list(compress("ABCDE", [1, 0, 1, 0, 1])))
print(list(accumulate([5, 3, 8, 2])))
print(list(accumulate([5, 3, 8, 2], max)))
print(list(accumulate([1, 2, 3, 4], operator.mul)))
Output
[12, 15, 19]
[23, 21, 17, 25]
[12]
['A', 'C', 'E']
[5, 8, 16, 18]
[5, 5, 8, 8]
[1, 2, 6, 24]

accumulate produces running totals — or running maximums, products, or any function you give it.

Grouping with groupby

PY groupby.py
from itertools import groupby

sales = [
    {"month": "Jan", "amount": 100},
    {"month": "Jan", "amount": 50},
    {"month": "Feb", "amount": 80},
    {"month": "Mar", "amount": 30},
    {"month": "Mar", "amount": 90},
]

for month, rows in groupby(sales, key=lambda r: r["month"]):
    print(month, sum(r["amount"] for r in rows))

print([(k, len(list(g))) for k, g in groupby("aaabccdd")])
Output
Jan 150
Feb 80
Mar 120
[('a', 3), ('b', 1), ('c', 2), ('d', 2)]

groupby groups consecutive items with the same key, so data usually needs sorting by that key first. For unsorted data, a defaultdict(list) (lesson 27) is often simpler.

Combinatorics

PY combinatorics.py
from itertools import product, permutations, combinations

print(list(product("AB", [1, 2])))
print(["".join(p) for p in permutations("abc", 2)])
print(list(combinations(["tea", "cake", "scone"], 2)))
print(len(list(product(range(10), repeat=4))))
Output
[('A', 1), ('A', 2), ('B', 1), ('B', 2)]
['ab', 'ac', 'ba', 'bc', 'ca', 'cb']
[('tea', 'cake'), ('tea', 'scone'), ('cake', 'scone')]
10000

product replaces nested loops; permutations considers order, combinations doesn't.

functools: caching

PY caching.py
from functools import cache, lru_cache

@cache
def ways_to_climb(steps):
    if steps <= 1:
        return 1
    return ways_to_climb(steps - 1) + ways_to_climb(steps - 2)

@lru_cache(maxsize=128)
def slow_lookup(key):
    return key.upper()

print(ways_to_climb(80))
slow_lookup("a"); slow_lookup("a"); slow_lookup("b")
print(slow_lookup.cache_info())
Output
37889062373143906
CacheInfo(hits=1, misses=2, maxsize=128, currsize=2)

@cache remembers every result forever; @lru_cache(maxsize=…) keeps only the most recently used results, bounding memory. Only cache functions whose result depends solely on their arguments.

functools: partial, reduce, cached_property

PY more_functools.py
from functools import partial, reduce, cached_property
import operator

to_binary = partial(int, base=2)
print(to_binary("1011"))

print(reduce(operator.mul, range(1, 6)))

class Report:
    def __init__(self, rows):
        self.rows = rows

    @cached_property
    def total(self):
        print("(computing total once)")
        return sum(self.rows)

r = Report([3, 4, 5])
print(r.total, r.total)
Output
11
120
(computing total once)
12 12

cached_property computes an attribute the first time it's accessed and stores the result on the object.

Try it yourself

Try it yourself
Console output
PreviousRegular Expressions