itertools and functools

itertools

Every itertools function returns a lazy iterator. Nothing is computed until you ask for it, so these work on sequences too large to fit in memory - and on infinite ones.

Infinite iterators

import itertools

for i in itertools.count(10, 2):        # 10, 12, 14, ... forever
    if i > 16:
        break
    print(i, end=" ")
print()

cycled = itertools.cycle(["a", "b", "c"])
print([next(cycled) for _ in range(7)])     # a b c a b c a

print(list(itertools.repeat("x", 3)))       # ['x', 'x', 'x']

Never write list(itertools.count()). It never stops. Always pair an infinite iterator with islice, zip or a break.

Combining

import itertools

print(list(itertools.chain([1, 2], [3, 4], [5])))          # [1, 2, 3, 4, 5]
print(list(itertools.chain.from_iterable([[1, 2], [3]])))  # [1, 2, 3]

print(list(itertools.zip_longest([1, 2, 3], "ab", fillvalue="-")))
# [(1, 'a'), (2, 'b'), (3, '-')]  - zip would have stopped at two

print(list(itertools.product([1, 2], "ab")))
# [(1, 'a'), (1, 'b'), (2, 'a'), (2, 'b')]  - nested loops, flattened

print(list(itertools.product([0, 1], repeat=3)))    # every 3 bit combination

Selecting

import itertools

numbers = [1, 3, 5, 2, 7, 4, 9]

print(list(itertools.islice(numbers, 2, 5)))          # [5, 2, 7] - slicing an iterator
print(list(itertools.takewhile(lambda n: n % 2, numbers)))   # [1, 3, 5] - stops at 2
print(list(itertools.dropwhile(lambda n: n % 2, numbers)))   # [2, 7, 4, 9]
print(list(itertools.filterfalse(lambda n: n % 2, numbers))) # [2, 4] - the evens
print(list(itertools.compress("abcdef", [1, 0, 1, 0, 1, 0])))  # ['a', 'c', 'e']

takewhile stops at the first failure; filter checks everything. That difference matters on a long or infinite sequence.

Accumulating

import itertools
import operator

sales = [100, 250, 75, 300]

print(list(itertools.accumulate(sales)))                      # running total
print(list(itertools.accumulate(sales, operator.mul)))        # running product
print(list(itertools.accumulate(sales, max)))                 # running maximum

Grouping

import itertools

data = [
    ("eng", "Meera"),
    ("eng", "Arun"),
    ("design", "Sara"),
    ("design", "Ravi"),
]

for department, members in itertools.groupby(data, key=lambda row: row[0]):
    print(department, [name for _, name in members])

groupby groups consecutive equal keys, exactly like the Unix uniq command. It does not gather scattered matches. Sort by the same key first, or the results will be wrong in a way that is easy to miss.

import itertools

data = [("eng", "Meera"), ("design", "Sara"), ("eng", "Arun")]

# Wrong: eng appears twice as separate groups
for key, group in itertools.groupby(data, key=lambda r: r[0]):
    print(key, list(group))

print("---")

# Right
data.sort(key=lambda r: r[0])
for key, group in itertools.groupby(data, key=lambda r: r[0]):
    print(key, [name for _, name in group])

Combinatorics

import itertools

items = ["a", "b", "c"]

print(list(itertools.permutations(items, 2)))
# order matters: (a,b) and (b,a) are both present

print(list(itertools.combinations(items, 2)))
# order does not matter: only (a,b), (a,c), (b,c)

print(list(itertools.combinations_with_replacement(items, 2)))
# includes (a,a)

print(len(list(itertools.permutations(range(8)))))     # 40320

These grow explosively. Ten items have 3.6 million permutations; twelve have 479 million. Count before you build a list.

Practical recipes

import itertools


def chunked(iterable, size):
    """Yield fixed size chunks from any iterable, lazily."""
    iterator = iter(iterable)
    while chunk := list(itertools.islice(iterator, size)):
        yield chunk


print(list(chunked(range(10), 3)))     # [[0,1,2], [3,4,5], [6,7,8], [9]]


def pairwise(iterable):
    """Yield overlapping pairs. Also available as itertools.pairwise in 3.10+."""
    a, b = itertools.tee(iterable)
    next(b, None)
    return zip(a, b)


print(list(pairwise([1, 2, 3, 4])))    # [(1, 2), (2, 3), (3, 4)]

temperatures = [20, 22, 21, 25]
print([b - a for a, b in pairwise(temperatures)])     # daily changes

functools

cache and lru_cache

from functools import cache, lru_cache


@cache                       # Python 3.9+; unlimited size
def fib(n):
    return n if n < 2 else fib(n - 1) + fib(n - 2)


print(fib(100))              # instant
print(fib.cache_info())
fib.cache_clear()


@lru_cache(maxsize=128)      # keeps only the 128 most recent
def slow_lookup(key):
    print("computing", key)
    return key.upper()


print(slow_lookup("a"))      # computing a -> A
print(slow_lookup("a"))      # A, no computation

Two rules: the arguments must be hashable, so no lists or dictionaries; and the function must be pure, giving the same result for the same arguments with no side effects. Caching a function that reads a file or the clock will return stale answers forever.

partial

from functools import partial


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


square = partial(power, exponent=2)
two_to_the = partial(power, 2)

print(square(7), two_to_the(10))         # 49 1024


def log(level, message):
    print(f"[{level}] {message}")


warn = partial(log, "WARNING")
error = partial(log, "ERROR")
warn("disk nearly full")

reduce

from functools import reduce
import operator

numbers = [1, 2, 3, 4]

print(reduce(operator.mul, numbers))            # 24
print(reduce(operator.add, numbers, 100))       # 110, with a starting value
print(reduce(lambda a, b: a if a > b else b, numbers))    # 4

# Prefer the built ins where they exist
print(sum(numbers), max(numbers), min(numbers))

Use reduce only when no built in covers the operation - a product, a custom merge, a fold over dictionaries. A loop is usually clearer than a clever reduce.

wraps

from functools import wraps


def logged(func):
    @wraps(func)                     # copy the name, docstring and signature
    def wrapper(*args, **kwargs):
        print(f"calling {func.__name__}")
        return func(*args, **kwargs)
    return wrapper


@logged
def greet(name):
    """Return a greeting."""
    return f"Hello, {name}"


print(greet.__name__)      # greet    - without @wraps this would be 'wrapper'
print(greet.__doc__)       # Return a greeting.

Every decorator you write should use @wraps. Without it the wrapped function loses its identity, which breaks documentation tools, debuggers and test discovery. The decorators note covers this fully.

total_ordering

from functools import total_ordering


@total_ordering
class Version:
    def __init__(self, major, minor):
        self.major = major
        self.minor = minor

    def __eq__(self, other):
        return (self.major, self.minor) == (other.major, other.minor)

    def __lt__(self, other):
        return (self.major, self.minor) < (other.major, other.minor)

    def __repr__(self):
        return f"v{self.major}.{self.minor}"


versions = [Version(1, 2), Version(1, 10), Version(0, 9)]
print(sorted(versions))                      # [v0.9, v1.2, v1.10]
print(Version(1, 2) >= Version(1, 1))        # True - generated for you

Define __eq__ and one ordering method, and total_ordering fills in the other three.

singledispatch

from functools import singledispatch


@singledispatch
def describe(value):
    return f"some object: {value!r}"


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


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


@describe.register
def _(value: str):
    return f"text of length {len(value)}"


print(describe(42))          # an integer: 42
print(describe([1, 2]))      # a list of 2 items
print(describe("hi"))        # text of length 2
print(describe(3.5))         # some object: 3.5

One function name, several implementations chosen by the type of the first argument. It replaces a chain of isinstance checks and lets new types be registered later without editing the original.

Common mistakes

• Calling list() on an infinite iterator.
• Using groupby without sorting first.
• Consuming an iterator twice; the second pass is empty.
• Caching a function that reads changing state.
• Passing a list to a cached function and meeting TypeError: unhashable type.
• Writing a decorator without @wraps.
• Building a permutation list of more than about ten items.

Best practices

• Use itertools to avoid building intermediate lists.
• Sort before groupby, on the same key.
• Use cache for pure, expensive, repeatedly called functions.
• Use partial instead of a lambda that only fixes an argument.
• Always use @wraps in a decorator.
• Prefer sum, max and min over reduce.

Practice

1. Split a list of 1000 records into batches of 50 using islice, lazily.
2. Group a list of transactions by month, remembering to sort first.
3. Generate every three letter combination from the alphabet and count them without building a list.
4. Add @cache to a recursive function and measure the difference.
5. Write a singledispatch function formatting integers, lists and dictionaries differently.

Conclusion

itertools gives you lazy building blocks for sequences: chain, islice, product, groupby, accumulate. functools transforms functions themselves: cache them, fix their arguments, fold them, and - crucially - preserve their identity when wrapping with @wraps.