ruk·si

🐍 Python
itertools

Updated at 2018-11-19 19:50

Standard library itertools module contains a lot of useful iterable tools.

Filtering:

import itertools

def is_vowel(c: str) -> bool:
    return c.lower() in 'aeiou'

# yields predicate to each item, returning the item if True
assert list(filter(is_vowel, 'Aardvark')) == ['A', 'a', 'a']

# same as field but returns if False
assert list(itertools.filterfalse(is_vowel, 'Aardvark')) == ['r', 'd', 'v', 'r', 'k']

# returns items AFTER predicate is True
assert list(itertools.dropwhile(is_vowel, 'Aardvark')) == ['r', 'd', 'v', 'a', 'r', 'k']

# returns items as long as predicate is True
assert list(itertools.takewhile(is_vowel, 'Aardvark')) == ['A', 'a']

# compare two generators side-by-side and return left element if right is True
assert list(itertools.compress('Aardvark', (1, 0, 0, 0, 0, 0, 1))) == ['A', 'r']

# same as list[:] but works with any iterable and is lazy
assert list(itertools.islice('Aardvark', 3)) == ['A', 'a', 'r']

Mapping:

import itertools
import operator

sample = [5, 4, 2, 8]

# yields accumulated sums (by default)
assert list(itertools.accumulate(sample)) == [5, 9, 11, 19]

# yields function return forward
assert list(itertools.accumulate(sample, min)) == [5, 4, 2, 2]
assert list(itertools.accumulate(sample, max)) == [5, 5, 5, 8]

# yields (index/key, value) pairs starting from index of your choosing
assert list(enumerate('cat', 1)) == [(1, 'c'), (2, 'a'), (3, 't')]

# map applies function to iterables and yields the result
# itertools.starmap is the same but does func(*iterator)
assert list(map(str.upper, 'cat')) == ['C', 'A', 'T']
assert list(map(operator.mul, range(4), range(4))) == [0, 1, 4, 9]

Merging:

import itertools

# yields all items from the first iterable, then from the next, etc.
assert list(itertools.chain('ABC', range(2))) == ['A', 'B', 'C', 0, 1]

# yields all item produced by the given iterator
assert list(itertools.chain.from_iterable(enumerate('ABC'))) == [0, 'A', 1, 'B', 2, 'C']

# yields items as sets until one stops iterating
assert list(zip('AB', range(3))) == [('A', 0), ('B', 1)]

# yields items as sets until all stop iterating
assert list(itertools.zip_longest('AB', range(3))) == [('A', 0), ('B', 1), (None, 2)]

Producers:

import itertools

# yields next increment infinitely, can be customized
ct = itertools.count()
assert next(ct) == 0 and next(ct) == 1 and next(ct) == 2

# yields items from an iterable, cycling around in the end
cy = itertools.cycle('AB')
assert next(cy) == 'A' and next(cy) == 'B' and next(cy) == 'A'

# yields the given value, can be limited how many times
r = itertools.repeat('AB')
assert next(r) == 'AB' and next(r) == 'AB'

Combinatorics:

import itertools

assert list(itertools.combinations('ABC', 2)) == [
    ('A', 'B'), ('A', 'C'), ('B', 'C')
]
assert list(itertools.combinations_with_replacement('ABC', 2)) == [
    ('A', 'A'), ('A', 'B'), ('A', 'C'),
    ('B', 'B'), ('B', 'C'), ('C', 'C')
]
assert list(itertools.permutations('ABC', 2)) == [
    ('A', 'B'), ('A', 'C'), ('B', 'A'),
    ('B', 'C'), ('C', 'A'), ('C', 'B')
]
assert list(itertools.product('ABC', repeat=2)) == [
    ('A', 'A'), ('A', 'B'), ('A', 'C'),
    ('B', 'A'), ('B', 'B'), ('B', 'C'),
    ('C', 'A'), ('C', 'B'), ('C', 'C')
]

Grouping:

import itertools

def is_g(char) -> str:
    return 'g' if char.lower() == 'g' else 'not-g'

groups = {char: list(items) for char, items in itertools.groupby('LLlAGg', is_g)}
assert groups['not-g'] == ['L', 'L', 'l', 'A']
assert groups['g'] == ['G', 'g']

Tee duplicates an iterable.

import itertools

g1, g2 = itertools.tee(('AB'))
assert list(g1) == ['A', 'B']
assert list(g2) == ['A', 'B']

assert list(zip(*itertools.tee('AB'))) == [('A', 'A'), ('B', 'B')]

Reducers:

# returns True if all items are truthy
assert all([True, True])
assert not all([True, False])
assert not all([False, False])

# returns True if any of the items is truthy
assert any([True, True])
assert any([True, False])
assert not any([False, False])

Source

  • Fluent Python, Luciano Ramalho