Python itertools Recipes

from itertools import islice
from collections import deque
 
def sliding_window(iterable, n):
    """sliding_window('ABCDEF', 3) → ABC BCD CDE DEF"""
    it = iter(iterable)
    window = deque(islice(it, n), maxlen=n)
    if len(window) == n:
        yield tuple(window)
    for x in it:
        window.append(x)
        yield tuple(window)
 
# Usage
data = [1, 2, 3, 4, 5, 6]
for window in sliding_window(data, 3):
    print(window)
# (1, 2, 3)
# (2, 3, 4)
# (3, 4, 5)
# (4, 5, 6)

from itertools import pairwise
 
# Built-in since 3.10
list(pairwise('ABCD'))  # [('A', 'B'), ('B', 'C'), ('C', 'D')]
 
# For older versions
def pairwise(iterable):
    a, b = tee(iterable)
    next(b, None)
    return zip(a, b)

from itertools import zip_longest
 
def grouper(iterable, n, fillvalue=None):
    """grouper('ABCDEFG', 3, 'x') → ABC DEF Gxx"""
    args = [iter(iterable)] * n
    return zip_longest(*args, fillvalue=fillvalue)
 
list(grouper('ABCDEFG', 3, 'x'))
# [('A', 'B', 'C'), ('D', 'E', 'F'), ('G', 'x', 'x')]
 
# Without fill value (strict chunks)
def batched(iterable, n):
    """batched('ABCDEFG', 3) → ABC DEF G"""
    from itertools import islice
    it = iter(iterable)
    while batch := tuple(islice(it, n)):
        yield batch
 
list(batched('ABCDEFG', 3))
# [('A', 'B', 'C'), ('D', 'E', 'F'), ('G',)]

from itertools import chain
 
def flatten(list_of_lists):
    """Flatten one level of nesting."""
    return chain.from_iterable(list_of_lists)
 
list(flatten([[1, 2], [3, 4], [5]]))  # [1, 2, 3, 4, 5]
 
# Deep flatten
def deep_flatten(nested):
    for item in nested:
        if isinstance(item, (list, tuple)):
            yield from deep_flatten(item)
        else:
            yield item
 
list(deep_flatten([1, [2, [3, 4]], 5]))  # [1, 2, 3, 4, 5]

from itertools import filterfalse
 
def unique_everseen(iterable, key=None):
    """List unique elements, preserving order."""
    seen = set()
    if key is None:
        for element in filterfalse(seen.__contains__, iterable):
            seen.add(element)
            yield element
    else:
        for element in iterable:
            k = key(element)
            if k not in seen:
                seen.add(k)
                yield element
 
list(unique_everseen('AAAABBBCCDAABBB'))  # ['A', 'B', 'C', 'D']
list(unique_everseen('ABBcCAD', key=str.lower))  # ['A', 'B', 'c', 'D']

from itertools import islice
 
def take(n, iterable):
    """Return first n items."""
    return list(islice(iterable, n))
 
def drop(n, iterable):
    """Drop first n items."""
    return islice(iterable, n, None)
 
def tail(n, iterable):
    """Return last n items."""
    from collections import deque
    return iter(deque(iterable, maxlen=n))
 
take(3, range(10))       # [0, 1, 2]
list(drop(3, range(10))) # [3, 4, 5, 6, 7, 8, 9]
list(tail(3, range(10))) # [7, 8, 9]

from itertools import islice
 
def nth(iterable, n, default=None):
    """Return nth item or default."""
    return next(islice(iterable, n, None), default)
 
nth(range(10), 3)     # 3
nth(range(10), 100)   # None
nth(range(10), 100, 'missing')  # 'missing'

from itertools import islice
from collections import deque
 
def consume(iterator, n=None):
    """Advance iterator n steps (or exhaust if n is None)."""
    if n is None:
        deque(iterator, maxlen=0)
    else:
        next(islice(iterator, n, n), None)
 
it = iter(range(10))
consume(it, 3)
print(next(it))  # 3

from itertools import filterfalse, tee
 
def partition(predicate, iterable):
    """Split into false and true elements."""
    t1, t2 = tee(iterable)
    return filterfalse(predicate, t1), filter(predicate, t2)
 
is_odd = lambda x: x % 2
evens, odds = partition(is_odd, range(10))
list(evens)  # [0, 2, 4, 6, 8]
list(odds)   # [1, 3, 5, 7, 9]

from itertools import cycle, islice
 
def roundrobin(*iterables):
    """roundrobin('ABC', 'D', 'EF') → A D E B F C"""
    iterators = [iter(it) for it in iterables]
    while iterators:
        for i, it in enumerate(iterators[:]):
            try:
                yield next(it)
            except StopIteration:
                iterators.remove(it)
 
list(roundrobin('ABC', 'D', 'EF'))  # ['A', 'D', 'E', 'B', 'F', 'C']

from itertools import chain, combinations
 
def powerset(iterable):
    """powerset([1,2,3]) → () (1,) (2,) (3,) (1,2) (1,3) (2,3) (1,2,3)"""
    s = list(iterable)
    return chain.from_iterable(combinations(s, r) for r in range(len(s) + 1))
 
list(powerset([1, 2, 3]))
# [(), (1,), (2,), (3,), (1, 2), (1, 3), (2, 3), (1, 2, 3)]

def first_true(iterable, default=False, predicate=None):
    """Return first true value or default."""
    return next(filter(predicate, iterable), default)
 
first_true([0, '', [], 'hello', 'world'])  # 'hello'
first_true([0, '', []], default='nothing')  # 'nothing'
first_true([1, 3, 5, 4, 2], predicate=lambda x: x % 2 == 0)  # 4

def quantify(iterable, predicate=bool):
    """Count how many times predicate is True."""
    return sum(map(predicate, iterable))
 
quantify([True, False, True, True])  # 3
quantify(range(10), lambda x: x > 5)  # 4

from itertools import groupby
 
def all_equal(iterable):
    """Check if all elements are equal."""
    g = groupby(iterable)
    return next(g, True) and not next(g, False)
 
all_equal('AAAA')  # True
all_equal('AABA')  # False
all_equal([])      # True

from itertools import chain
 
def interleave(*iterables):
    """interleave('ABC', '12') → A 1 B 2 C"""
    return chain.from_iterable(zip(*iterables))
 
list(interleave('ABC', '123'))  # ['A', '1', 'B', '2', 'C', '3']

from itertools import groupby, islice
from operator import itemgetter
 
# Process log entries
logs = [
    {'level': 'ERROR', 'msg': 'Failed'},
    {'level': 'INFO', 'msg': 'Started'},
    {'level': 'ERROR', 'msg': 'Timeout'},
    {'level': 'INFO', 'msg': 'Completed'},
]
 
# Group by level
keyfunc = itemgetter('level')
for level, entries in groupby(sorted(logs, key=keyfunc), keyfunc):
    print(f"{level}: {list(entries)}")
 
# Process in batches
def process_in_batches(items, batch_size=100):
    it = iter(items)
    while batch := list(islice(it, batch_size)):
        process_batch(batch)