Skip to content

Commit

Permalink
Minor readability improvements for the itertools recipes (gh-127928)
Browse files Browse the repository at this point in the history
  • Loading branch information
rhettinger authored Dec 13, 2024
1 parent 5dd775b commit 292067f
Showing 1 changed file with 35 additions and 39 deletions.
74 changes: 35 additions & 39 deletions Doc/library/itertools.rst
Original file line number Diff line number Diff line change
Expand Up @@ -30,11 +30,6 @@ For instance, SML provides a tabulation tool: ``tabulate(f)`` which produces a
sequence ``f(0), f(1), ...``. The same effect can be achieved in Python
by combining :func:`map` and :func:`count` to form ``map(f, count())``.

These tools and their built-in counterparts also work well with the high-speed
functions in the :mod:`operator` module. For example, the multiplication
operator can be mapped across two vectors to form an efficient dot-product:
``sum(starmap(operator.mul, zip(vec1, vec2, strict=True)))``.


**Infinite iterators:**

Expand Down Expand Up @@ -843,12 +838,11 @@ and :term:`generators <generator>` which incur interpreter overhead.

.. testcode::

import collections
import contextlib
import functools
import math
import operator
import random
from collections import deque
from contextlib import suppress
from functools import reduce
from math import sumprod, isqrt
from operator import itemgetter, getitem, mul, neg

def take(n, iterable):
"Return first n items of the iterable as a list."
Expand All @@ -863,11 +857,11 @@ and :term:`generators <generator>` which incur interpreter overhead.
"Return function(0), function(1), ..."
return map(function, count(start))

def repeatfunc(func, times=None, *args):
"Repeat calls to func with specified arguments."
def repeatfunc(function, times=None, *args):
"Repeat calls to a function with specified arguments."
if times is None:
return starmap(func, repeat(args))
return starmap(func, repeat(args, times))
return starmap(function, repeat(args))
return starmap(function, repeat(args, times))
def flatten(list_of_lists):
"Flatten one level of nesting."
Expand All @@ -885,13 +879,13 @@ and :term:`generators <generator>` which incur interpreter overhead.
def tail(n, iterable):
"Return an iterator over the last n items."
# tail(3, 'ABCDEFG') → E F G
return iter(collections.deque(iterable, maxlen=n))
return iter(deque(iterable, maxlen=n))

def consume(iterator, n=None):
"Advance the iterator n-steps ahead. If n is None, consume entirely."
# Use functions that consume iterators at C speed.
if n is None:
collections.deque(iterator, maxlen=0)
deque(iterator, maxlen=0)
else:
next(islice(iterator, n, n), None)

Expand Down Expand Up @@ -919,8 +913,8 @@ and :term:`generators <generator>` which incur interpreter overhead.
# unique_justseen('AAAABBBCCDAABBB') → A B C D A B
# unique_justseen('ABBcCAD', str.casefold) → A B c A D
if key is None:
return map(operator.itemgetter(0), groupby(iterable))
return map(next, map(operator.itemgetter(1), groupby(iterable, key)))
return map(itemgetter(0), groupby(iterable))
return map(next, map(itemgetter(1), groupby(iterable, key)))

def unique_everseen(iterable, key=None):
"Yield unique elements, preserving order. Remember all elements ever seen."
Expand All @@ -941,13 +935,14 @@ and :term:`generators <generator>` which incur interpreter overhead.
def unique(iterable, key=None, reverse=False):
"Yield unique elements in sorted order. Supports unhashable inputs."
# unique([[1, 2], [3, 4], [1, 2]]) → [1, 2] [3, 4]
return unique_justseen(sorted(iterable, key=key, reverse=reverse), key=key)
sequenced = sorted(iterable, key=key, reverse=reverse)
return unique_justseen(sequenced, key=key)

def sliding_window(iterable, n):
"Collect data into overlapping fixed-length chunks or blocks."
# sliding_window('ABCDEFG', 4) → ABCD BCDE CDEF DEFG
iterator = iter(iterable)
window = collections.deque(islice(iterator, n - 1), maxlen=n)
window = deque(islice(iterator, n - 1), maxlen=n)
for x in iterator:
window.append(x)
yield tuple(window)
Expand Down Expand Up @@ -981,7 +976,7 @@ and :term:`generators <generator>` which incur interpreter overhead.
"Return all contiguous non-empty subslices of a sequence."
# subslices('ABCD') → A AB ABC ABCD B BC BCD C CD D
slices = starmap(slice, combinations(range(len(seq) + 1), 2))
return map(operator.getitem, repeat(seq), slices)
return map(getitem, repeat(seq), slices)

def iter_index(iterable, value, start=0, stop=None):
"Return indices where a value occurs in a sequence or iterable."
Expand All @@ -995,39 +990,40 @@ and :term:`generators <generator>` which incur interpreter overhead.
else:
stop = len(iterable) if stop is None else stop
i = start
with contextlib.suppress(ValueError):
with suppress(ValueError):
while True:
yield (i := seq_index(value, i, stop))
i += 1

def iter_except(func, exception, first=None):
def iter_except(function, exception, first=None):
"Convert a call-until-exception interface to an iterator interface."
# iter_except(d.popitem, KeyError) → non-blocking dictionary iterator
with contextlib.suppress(exception):
with suppress(exception):
if first is not None:
yield first()
while True:
yield func()
yield function()


The following recipes have a more mathematical flavor:

.. testcode::

def powerset(iterable):
"Subsequences of the iterable from shortest to longest."
# 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))

def sum_of_squares(iterable):
"Add up the squares of the input values."
# sum_of_squares([10, 20, 30]) → 1400
return math.sumprod(*tee(iterable))
return sumprod(*tee(iterable))
def reshape(matrix, cols):
def reshape(matrix, columns):
"Reshape a 2-D matrix to have a given number of columns."
# reshape([(0, 1), (2, 3), (4, 5)], 3) → (0, 1, 2), (3, 4, 5)
return batched(chain.from_iterable(matrix), cols, strict=True)
return batched(chain.from_iterable(matrix), columns, strict=True)

def transpose(matrix):
"Swap the rows and columns of a 2-D matrix."
Expand All @@ -1038,7 +1034,7 @@ The following recipes have a more mathematical flavor:
"Multiply two matrices."
# matmul([(7, 5), (3, 5)], [(2, 5), (7, 9)]) → (49, 80), (41, 60)
n = len(m2[0])
return batched(starmap(math.sumprod, product(m1, transpose(m2))), n)
return batched(starmap(sumprod, product(m1, transpose(m2))), n)

def convolve(signal, kernel):
"""Discrete linear convolution of two iterables.
Expand All @@ -1059,16 +1055,16 @@ The following recipes have a more mathematical flavor:
n = len(kernel)
padded_signal = chain(repeat(0, n-1), signal, repeat(0, n-1))
windowed_signal = sliding_window(padded_signal, n)
return map(math.sumprod, repeat(kernel), windowed_signal)
return map(sumprod, repeat(kernel), windowed_signal)

def polynomial_from_roots(roots):
"""Compute a polynomial's coefficients from its roots.

(x - 5) (x + 4) (x - 3) expands to: x³ -4x² -17x + 60
"""
# polynomial_from_roots([5, -4, 3]) → [1, -4, -17, 60]
factors = zip(repeat(1), map(operator.neg, roots))
return list(functools.reduce(convolve, factors, [1]))
factors = zip(repeat(1), map(neg, roots))
return list(reduce(convolve, factors, [1]))

def polynomial_eval(coefficients, x):
"""Evaluate a polynomial at a specific value.
Expand All @@ -1081,7 +1077,7 @@ The following recipes have a more mathematical flavor:
if not n:
return type(x)(0)
powers = map(pow, repeat(x), reversed(range(n)))
return math.sumprod(coefficients, powers)
return sumprod(coefficients, powers)

def polynomial_derivative(coefficients):
"""Compute the first derivative of a polynomial.
Expand All @@ -1092,15 +1088,15 @@ The following recipes have a more mathematical flavor:
# polynomial_derivative([1, -4, -17, 60]) → [3, -8, -17]
n = len(coefficients)
powers = reversed(range(1, n))
return list(map(operator.mul, coefficients, powers))
return list(map(mul, coefficients, powers))

def sieve(n):
"Primes less than n."
# sieve(30) → 2 3 5 7 11 13 17 19 23 29
if n > 2:
yield 2
data = bytearray((0, 1)) * (n // 2)
for p in iter_index(data, 1, start=3, stop=math.isqrt(n) + 1):
for p in iter_index(data, 1, start=3, stop=isqrt(n) + 1):
data[p*p : n : p+p] = bytes(len(range(p*p, n, p+p)))
yield from iter_index(data, 1, start=3)

Expand All @@ -1109,7 +1105,7 @@ The following recipes have a more mathematical flavor:
# factor(99) → 3 3 11
# factor(1_000_000_000_000_007) → 47 59 360620266859
# factor(1_000_000_000_000_403) → 1000000000000403
for prime in sieve(math.isqrt(n) + 1):
for prime in sieve(isqrt(n) + 1):
while not n % prime:
yield prime
n //= prime
Expand Down Expand Up @@ -1740,7 +1736,7 @@ The following recipes have a more mathematical flavor:

# Old recipes and their tests which are guaranteed to continue to work.

def sumprod(vec1, vec2):
def old_sumprod_recipe(vec1, vec2):
"Compute a sum of products."
return sum(starmap(operator.mul, zip(vec1, vec2, strict=True)))

Expand Down Expand Up @@ -1823,7 +1819,7 @@ The following recipes have a more mathematical flavor:
32


>>> sumprod([1,2,3], [4,5,6])
>>> old_sumprod_recipe([1,2,3], [4,5,6])
32


Expand Down

0 comments on commit 292067f

Please sign in to comment.