An Introduction to Python's itertools Library

Python developers often aim to write more Pythonic code that is both readable and efficient; the itertools module in the standard library offers a collection of iterator-building functions that can replace verbose loops and improve performance.

itertools.accumulate – Returns an iterator that yields accumulated sums (or other binary functions) of the input iterable.

import itertools
x = itertools.accumulate(range(10))
print(list(x))  # [0, 1, 3, 6, 10, 15, 21, 28, 36, 45]

itertools.chain – Concatenates multiple iterables into a single iterator.

import itertools
x = itertools.chain(range(3), range(4), [3,2,1])
print(list(x))  # [0, 1, 2, 0, 1, 2, 3, 3, 2, 1]

itertools.combinations – Generates all possible r‑length combinations of elements from the input iterable without replacement.

import itertools
x = itertools.combinations(range(4), 3)
print(list(x))  # [(0,1,2), (0,1,3), (0,2,3), (1,2,3)]

itertools.combinations_with_replacement – Like combinations, but allows individual elements to be repeated.

import itertools
x = itertools.combinations_with_replacement('ABC', 2)
print(list(x))  # [('A','A'), ('A','B'), ('A','C'), ('B','B'), ('B','C'), ('C','C')]

itertools.compress – Filters elements from the data iterable, keeping those where the corresponding selector is true.

import itertools
x = itertools.compress(range(5), (True, False, True, True, False))
print(list(x))  # [0, 2, 3]

itertools.count – Creates an infinite iterator that returns evenly spaced values starting from a given number.

import itertools
x = itertools.count(start=20, step=-1)
print(list(itertools.islice(x, 0, 10)))  # [20, 19, 18, 17, 16, 15, 14, 13, 12, 11]

itertools.cycle – Repeats the elements of an iterable indefinitely.

import itertools
x = itertools.cycle('ABC')
print(list(itertools.islice(x, 0, 10)))  # ['A','B','C','A','B','C','A','B','C','A']

itertools.dropwhile – Skips elements as long as the predicate is true, then yields the rest.

import itertools
x = itertools.dropwhile(lambda e: e < 5, range(10))
print(list(x))  # [5, 6, 7, 8, 9]

itertools.filterfalse – Yields elements for which the predicate is false.

import itertools
x = itertools.filterfalse(lambda e: e < 5, (1,5,3,6,9,4))
print(list(x))  # [5, 6, 9]

itertools.groupby – Groups consecutive elements that share a common key defined by a function.

import itertools
x = itertools.groupby(range(10), lambda x: x < 5 or x > 8)
for condition, numbers in x:
    print(condition, list(numbers))
# True [0,1,2,3,4]
# False [5,6,7,8]
# True [9]

itertools.islice – Slices an iterator by start, stop, and step parameters.

import itertools
x = itertools.islice(range(10), 0, 9, 2)
print(list(x))  # [0, 2, 4, 6, 8]

itertools.permutations – Generates all possible r‑length ordered arrangements of elements.

import itertools
x = itertools.permutations(range(4), 3)
print(list(x))  # [(0,1,2), (0,1,3), ... (3,2,1)]

itertools.product – Computes the Cartesian product of input iterables.

import itertools
x = itertools.product('ABC', range(3))
print(list(x))  # [('A',0), ('A',1), ('A',2), ('B',0), ...]

itertools.repeat – Returns an iterator that yields the same value a specified number of times.

import itertools
x = itertools.repeat(0, 5)
print(list(x))  # [0, 0, 0, 0, 0]

itertools.starmap – Applies a function to the unpacked arguments from an iterable of tuples, similar to map.

import itertools
x = itertools.starmap(str.islower, 'aBCDefGhI')
print(list(x))  # [True, False, False, False, False, True, True, False, True]

itertools.takewhile – Yields elements while the predicate remains true, then stops.

import itertools
x = itertools.takewhile(lambda e: e < 5, range(10))
print(list(x))  # [0, 1, 2, 3, 4]

itertools.tee – Creates multiple independent iterators from a single original iterator.

import itertools
x1, x2 = itertools.tee(range(10), 2)
print(list(x1))
print(list(x2))
# both output [0,1,2,3,4,5,6,7,8,9]

itertools.zip_longest – Zips iterables together, filling missing values with None when the iterables have unequal lengths.

import itertools
x = itertools.zip_longest(range(3), range(5))
print(list(x))  # [(0,0), (1,1), (2,2), (None,3), (None,4)]

In conclusion, mastering itertools and its many functions can greatly simplify iterator handling in Python, leading to cleaner, more efficient code once you become comfortable using these tools.