Python Itertools Recipes

from itertools import (
    chain, islice, groupby, product, combinations,
    accumulate, zip_longest, takewhile, dropwhile,
)
from typing import Iterator, TypeVar

T = TypeVar("T")


# Batch/chunk an iterable
def batched(iterable, n: int) -> Iterator[tuple]:
    it = iter(iterable)
    while batch := tuple(islice(it, n)):
        yield batch


# Flatten nested iterables
def flatten(nested: list[list[T]]) -> list[T]:
    return list(chain.from_iterable(nested))


# Sliding window
def sliding_window(seq: list[T], size: int) -> Iterator[tuple[T, ...]]:
    for i in range(len(seq) - size + 1):
        yield tuple(seq[i:i + size])


# Group consecutive items
data = [("A", 1), ("A", 2), ("B", 3), ("B", 4), ("A", 5)]
for key, group in groupby(data, key=lambda x: x[0]):
    print(f"{key}: {list(group)}")

# Running totals
numbers = [1, 2, 3, 4, 5]
running_sum = list(accumulate(numbers))           # [1, 3, 6, 10, 15]
running_max = list(accumulate(numbers, max))      # [1, 2, 3, 4, 5]

# All combinations
combos = list(combinations(["A", "B", "C", "D"], 2))
# [('A','B'), ('A','C'), ('A','D'), ('B','C'), ('B','D'), ('C','D')]

# Cartesian product
grid = list(product(range(3), range(3)))  # all (x, y) pairs

# Batch processing
items = list(range(25))
for batch in batched(items, 10):
    print(f"Processing batch of {len(batch)} items")

# Take/drop while condition
sorted_vals = [1, 3, 5, 7, 2, 4, 6]
ascending = list(takewhile(lambda x: x < 7, sorted_vals))  # [1, 3, 5]