back to Reference (Gold) summary
Reference (Gold): more-itertools
Pytest Summary for test tests
| status | count |
|---|---|
| passed | 662 |
| skipped | 1 |
| total | 663 |
| collected | 663 |
Failed pytests:
Patch diff
diff --git a/more_itertools/more.py b/more_itertools/more.py
index 4906f6d..3bf2c76 100755
--- a/more_itertools/more.py
+++ b/more_itertools/more.py
@@ -1,44 +1,161 @@
import math
import warnings
+
from collections import Counter, defaultdict, deque, abc
from collections.abc import Sequence
from contextlib import suppress
from functools import cached_property, partial, reduce, wraps
from heapq import heapify, heapreplace
-from itertools import chain, combinations, compress, count, cycle, dropwhile, groupby, islice, repeat, starmap, takewhile, tee, zip_longest, product
+from itertools import (
+ chain,
+ combinations,
+ compress,
+ count,
+ cycle,
+ dropwhile,
+ groupby,
+ islice,
+ repeat,
+ starmap,
+ takewhile,
+ tee,
+ zip_longest,
+ product,
+)
from math import comb, e, exp, factorial, floor, fsum, log, log1p, perm, tau
from queue import Empty, Queue
from random import random, randrange, shuffle, uniform
from operator import itemgetter, mul, sub, gt, lt, le
from sys import hexversion, maxsize
from time import monotonic
-from .recipes import _marker, _zip_equal, UnequalIterablesError, consume, flatten, powerset, take, unique_everseen, all_equal, batched
-__all__ = ['AbortThread', 'SequenceView', 'UnequalIterablesError',
- 'adjacent', 'all_unique', 'always_iterable', 'always_reversible',
- 'bucket', 'callback_iter', 'chunked', 'chunked_even', 'circular_shifts',
- 'collapse', 'combination_index', 'combination_with_replacement_index',
- 'consecutive_groups', 'constrained_batches', 'consumer', 'count_cycle',
- 'countable', 'dft', 'difference', 'distinct_combinations',
- 'distinct_permutations', 'distribute', 'divide', 'doublestarmap',
- 'duplicates_everseen', 'duplicates_justseen', 'classify_unique',
- 'exactly_n', 'filter_except', 'filter_map', 'first', 'gray_product',
- 'groupby_transform', 'ichunked', 'iequals', 'idft', 'ilen',
- 'interleave', 'interleave_evenly', 'interleave_longest', 'intersperse',
- 'is_sorted', 'islice_extended', 'iterate', 'iter_suppress',
- 'join_mappings', 'last', 'locate', 'longest_common_prefix', 'lstrip',
- 'make_decorator', 'map_except', 'map_if', 'map_reduce', 'mark_ends',
- 'minmax', 'nth_or_last', 'nth_permutation', 'nth_product',
- 'nth_combination_with_replacement', 'numeric_range', 'one', 'only',
- 'outer_product', 'padded', 'partial_product', 'partitions', 'peekable',
- 'permutation_index', 'powerset_of_sets', 'product_index', 'raise_',
- 'repeat_each', 'repeat_last', 'replace', 'rlocate', 'rstrip',
- 'run_length', 'sample', 'seekable', 'set_partitions', 'side_effect',
- 'sliced', 'sort_together', 'split_after', 'split_at', 'split_before',
- 'split_into', 'split_when', 'spy', 'stagger', 'strip', 'strictly_n',
- 'substrings', 'substrings_indexes', 'takewhile_inclusive',
- 'time_limited', 'unique_in_window', 'unique_to_each', 'unzip',
- 'value_chain', 'windowed', 'windowed_complete', 'with_iter',
- 'zip_broadcast', 'zip_equal', 'zip_offset']
+
+from .recipes import (
+ _marker,
+ _zip_equal,
+ UnequalIterablesError,
+ consume,
+ flatten,
+ powerset,
+ take,
+ unique_everseen,
+ all_equal,
+ batched,
+)
+
+__all__ = [
+ 'AbortThread',
+ 'SequenceView',
+ 'UnequalIterablesError',
+ 'adjacent',
+ 'all_unique',
+ 'always_iterable',
+ 'always_reversible',
+ 'bucket',
+ 'callback_iter',
+ 'chunked',
+ 'chunked_even',
+ 'circular_shifts',
+ 'collapse',
+ 'combination_index',
+ 'combination_with_replacement_index',
+ 'consecutive_groups',
+ 'constrained_batches',
+ 'consumer',
+ 'count_cycle',
+ 'countable',
+ 'dft',
+ 'difference',
+ 'distinct_combinations',
+ 'distinct_permutations',
+ 'distribute',
+ 'divide',
+ 'doublestarmap',
+ 'duplicates_everseen',
+ 'duplicates_justseen',
+ 'classify_unique',
+ 'exactly_n',
+ 'filter_except',
+ 'filter_map',
+ 'first',
+ 'gray_product',
+ 'groupby_transform',
+ 'ichunked',
+ 'iequals',
+ 'idft',
+ 'ilen',
+ 'interleave',
+ 'interleave_evenly',
+ 'interleave_longest',
+ 'intersperse',
+ 'is_sorted',
+ 'islice_extended',
+ 'iterate',
+ 'iter_suppress',
+ 'join_mappings',
+ 'last',
+ 'locate',
+ 'longest_common_prefix',
+ 'lstrip',
+ 'make_decorator',
+ 'map_except',
+ 'map_if',
+ 'map_reduce',
+ 'mark_ends',
+ 'minmax',
+ 'nth_or_last',
+ 'nth_permutation',
+ 'nth_product',
+ 'nth_combination_with_replacement',
+ 'numeric_range',
+ 'one',
+ 'only',
+ 'outer_product',
+ 'padded',
+ 'partial_product',
+ 'partitions',
+ 'peekable',
+ 'permutation_index',
+ 'powerset_of_sets',
+ 'product_index',
+ 'raise_',
+ 'repeat_each',
+ 'repeat_last',
+ 'replace',
+ 'rlocate',
+ 'rstrip',
+ 'run_length',
+ 'sample',
+ 'seekable',
+ 'set_partitions',
+ 'side_effect',
+ 'sliced',
+ 'sort_together',
+ 'split_after',
+ 'split_at',
+ 'split_before',
+ 'split_into',
+ 'split_when',
+ 'spy',
+ 'stagger',
+ 'strip',
+ 'strictly_n',
+ 'substrings',
+ 'substrings_indexes',
+ 'takewhile_inclusive',
+ 'time_limited',
+ 'unique_in_window',
+ 'unique_to_each',
+ 'unzip',
+ 'value_chain',
+ 'windowed',
+ 'windowed_complete',
+ 'with_iter',
+ 'zip_broadcast',
+ 'zip_equal',
+ 'zip_offset',
+]
+
+# math.sumprod is available for Python 3.12+
_fsumprod = getattr(math, 'sumprod', lambda x, y: fsum(map(mul, x, y)))
@@ -61,7 +178,20 @@ def chunked(iterable, n, strict=False):
list is yielded.
"""
- pass
+ iterator = iter(partial(take, n, iter(iterable)), [])
+ if strict:
+ if n is None:
+ raise ValueError('n must not be None when using strict mode.')
+
+ def ret():
+ for chunk in iterator:
+ if len(chunk) != n:
+ raise ValueError('iterable is not divisible by n.')
+ yield chunk
+
+ return iter(ret())
+ else:
+ return iterator
def first(iterable, default=_marker):
@@ -81,7 +211,14 @@ def first(iterable, default=_marker):
``next(iter(iterable), default)``.
"""
- pass
+ for item in iterable:
+ return item
+ if default is _marker:
+ raise ValueError(
+ 'first() was called on an empty iterable, and no '
+ 'default value was provided.'
+ )
+ return default
def last(iterable, default=_marker):
@@ -96,7 +233,21 @@ def last(iterable, default=_marker):
If *default* is not provided and there are no items in the iterable,
raise ``ValueError``.
"""
- pass
+ try:
+ if isinstance(iterable, Sequence):
+ return iterable[-1]
+ # Work around https://bugs.python.org/issue38525
+ elif hasattr(iterable, '__reversed__') and (hexversion != 0x030800F0):
+ return next(reversed(iterable))
+ else:
+ return deque(iterable, maxlen=1)[-1]
+ except (IndexError, TypeError, StopIteration):
+ if default is _marker:
+ raise ValueError(
+ 'last() was called on an empty iterable, and no default was '
+ 'provided.'
+ )
+ return default
def nth_or_last(iterable, n, default=_marker):
@@ -113,7 +264,7 @@ def nth_or_last(iterable, n, default=_marker):
If *default* is not provided and there are no items in the iterable,
raise ``ValueError``.
"""
- pass
+ return last(islice(iterable, n + 1), default=default)
class peekable:
@@ -196,7 +347,14 @@ class peekable:
provided, raise ``StopIteration``.
"""
- pass
+ if not self._cache:
+ try:
+ self._cache.append(next(self._it))
+ except StopIteration:
+ if default is _marker:
+ raise
+ return default
+ return self._cache[0]
def prepend(self, *items):
"""Stack up items to be the next ones returned from ``next()`` or
@@ -227,21 +385,50 @@ class peekable:
StopIteration
"""
- pass
+ self._cache.extendleft(reversed(items))
def __next__(self):
if self._cache:
return self._cache.popleft()
+
return next(self._it)
+ def _get_slice(self, index):
+ # Normalize the slice's arguments
+ step = 1 if (index.step is None) else index.step
+ if step > 0:
+ start = 0 if (index.start is None) else index.start
+ stop = maxsize if (index.stop is None) else index.stop
+ elif step < 0:
+ start = -1 if (index.start is None) else index.start
+ stop = (-maxsize - 1) if (index.stop is None) else index.stop
+ else:
+ raise ValueError('slice step cannot be zero')
+
+ # If either the start or stop index is negative, we'll need to cache
+ # the rest of the iterable in order to slice from the right side.
+ if (start < 0) or (stop < 0):
+ self._cache.extend(self._it)
+ # Otherwise we'll need to find the rightmost index and cache to that
+ # point.
+ else:
+ n = min(max(start, stop) + 1, maxsize)
+ cache_len = len(self._cache)
+ if n >= cache_len:
+ self._cache.extend(islice(self._it, n - cache_len))
+
+ return list(self._cache)[index]
+
def __getitem__(self, index):
if isinstance(index, slice):
return self._get_slice(index)
+
cache_len = len(self._cache)
if index < 0:
self._cache.extend(self._it)
elif index >= cache_len:
self._cache.extend(islice(self._it, index + 1 - cache_len))
+
return self._cache[index]
@@ -267,7 +454,14 @@ def consumer(func):
``t.send()`` could be used.
"""
- pass
+
+ @wraps(func)
+ def wrapper(*args, **kwargs):
+ gen = func(*args, **kwargs)
+ next(gen)
+ return gen
+
+ return wrapper
def ilen(iterable):
@@ -279,7 +473,10 @@ def ilen(iterable):
This consumes the iterable, so handle with care.
"""
- pass
+ # This is the "most beautiful of the fast variants" of this function.
+ # If you think you can improve on it, please ensure that your version
+ # is both 10x faster and 10x more beautiful.
+ return sum(compress(repeat(1), zip(iterable)))
def iterate(func, start):
@@ -290,7 +487,12 @@ def iterate(func, start):
[1, 2, 4, 8, 16, 32, 64, 128, 256, 512]
"""
- pass
+ while True:
+ yield start
+ try:
+ start = func(start)
+ except StopIteration:
+ break
def with_iter(context_manager):
@@ -304,7 +506,8 @@ def with_iter(context_manager):
``with_iter``.
"""
- pass
+ with context_manager as iterable:
+ yield from iterable
def one(iterable, too_short=None, too_long=None):
@@ -351,7 +554,31 @@ def one(iterable, too_short=None, too_long=None):
contents less destructively.
"""
- pass
+ it = iter(iterable)
+
+ try:
+ first_value = next(it)
+ except StopIteration as exc:
+ raise (
+ too_short or ValueError('too few items in iterable (expected 1)')
+ ) from exc
+
+ try:
+ second_value = next(it)
+ except StopIteration:
+ pass
+ else:
+ msg = (
+ 'Expected exactly one item in iterable, but got {!r}, {!r}, '
+ 'and perhaps more.'.format(first_value, second_value)
+ )
+ raise too_long or ValueError(msg)
+
+ return first_value
+
+
+def raise_(exception, *args):
+ raise exception(*args)
def strictly_n(iterable, n, too_short=None, too_long=None):
@@ -401,7 +628,34 @@ def strictly_n(iterable, n, too_short=None, too_long=None):
['a', 'b', 'c', 'd']
"""
- pass
+ if too_short is None:
+ too_short = lambda item_count: raise_(
+ ValueError,
+ 'Too few items in iterable (got {})'.format(item_count),
+ )
+
+ if too_long is None:
+ too_long = lambda item_count: raise_(
+ ValueError,
+ 'Too many items in iterable (got at least {})'.format(item_count),
+ )
+
+ it = iter(iterable)
+ for i in range(n):
+ try:
+ item = next(it)
+ except StopIteration:
+ too_short(i)
+ return
+ else:
+ yield item
+
+ try:
+ next(it)
+ except StopIteration:
+ pass
+ else:
+ too_long(n + 1)
def distinct_permutations(iterable, r=None):
@@ -447,7 +701,111 @@ def distinct_permutations(iterable, r=None):
('3', 2, 1)
]
"""
- pass
+
+ # Algorithm: https://w.wiki/Qai
+ def _full(A):
+ while True:
+ # Yield the permutation we have
+ yield tuple(A)
+
+ # Find the largest index i such that A[i] < A[i + 1]
+ for i in range(size - 2, -1, -1):
+ if A[i] < A[i + 1]:
+ break
+ # If no such index exists, this permutation is the last one
+ else:
+ return
+
+ # Find the largest index j greater than j such that A[i] < A[j]
+ for j in range(size - 1, i, -1):
+ if A[i] < A[j]:
+ break
+
+ # Swap the value of A[i] with that of A[j], then reverse the
+ # sequence from A[i + 1] to form the new permutation
+ A[i], A[j] = A[j], A[i]
+ A[i + 1 :] = A[: i - size : -1] # A[i + 1:][::-1]
+
+ # Algorithm: modified from the above
+ def _partial(A, r):
+ # Split A into the first r items and the last r items
+ head, tail = A[:r], A[r:]
+ right_head_indexes = range(r - 1, -1, -1)
+ left_tail_indexes = range(len(tail))
+
+ while True:
+ # Yield the permutation we have
+ yield tuple(head)
+
+ # Starting from the right, find the first index of the head with
+ # value smaller than the maximum value of the tail - call it i.
+ pivot = tail[-1]
+ for i in right_head_indexes:
+ if head[i] < pivot:
+ break
+ pivot = head[i]
+ else:
+ return
+
+ # Starting from the left, find the first value of the tail
+ # with a value greater than head[i] and swap.
+ for j in left_tail_indexes:
+ if tail[j] > head[i]:
+ head[i], tail[j] = tail[j], head[i]
+ break
+ # If we didn't find one, start from the right and find the first
+ # index of the head with a value greater than head[i] and swap.
+ else:
+ for j in right_head_indexes:
+ if head[j] > head[i]:
+ head[i], head[j] = head[j], head[i]
+ break
+
+ # Reverse head[i + 1:] and swap it with tail[:r - (i + 1)]
+ tail += head[: i - r : -1] # head[i + 1:][::-1]
+ i += 1
+ head[i:], tail[:] = tail[: r - i], tail[r - i :]
+
+ items = list(iterable)
+
+ try:
+ items.sort()
+ sortable = True
+ except TypeError:
+ sortable = False
+
+ indices_dict = defaultdict(list)
+
+ for item in items:
+ indices_dict[items.index(item)].append(item)
+
+ indices = [items.index(item) for item in items]
+ indices.sort()
+
+ equivalent_items = {k: cycle(v) for k, v in indices_dict.items()}
+
+ def permuted_items(permuted_indices):
+ return tuple(
+ next(equivalent_items[index]) for index in permuted_indices
+ )
+
+ size = len(items)
+ if r is None:
+ r = size
+
+ # functools.partial(_partial, ... )
+ algorithm = _full if (r == size) else partial(_partial, r=r)
+
+ if 0 < r <= size:
+ if sortable:
+ return algorithm(items)
+ else:
+ return (
+ permuted_items(permuted_indices)
+ for permuted_indices in algorithm(indices)
+ )
+
+ return iter(() if r else ((),))
def intersperse(e, iterable, n=1):
@@ -461,7 +819,19 @@ def intersperse(e, iterable, n=1):
[1, 2, None, 3, 4, None, 5]
"""
- pass
+ if n == 0:
+ raise ValueError('n must be > 0')
+ elif n == 1:
+ # interleave(repeat(e), iterable) -> e, x_0, e, x_1, e, x_2...
+ # islice(..., 1, None) -> x_0, e, x_1, e, x_2...
+ return islice(interleave(repeat(e), iterable), 1, None)
+ else:
+ # interleave(filler, chunks) -> [e], [x_0, x_1], [e], [x_2, x_3]...
+ # islice(..., 1, None) -> [x_0, x_1], [e], [x_2, x_3]...
+ # flatten(...) -> x_0, x_1, e, x_2, x_3...
+ filler = repeat([e])
+ chunks = chunked(iterable, n)
+ return flatten(islice(interleave(filler, chunks), 1, None))
def unique_to_each(*iterables):
@@ -491,7 +861,10 @@ def unique_to_each(*iterables):
It is assumed that the elements of each iterable are hashable.
"""
- pass
+ pool = [list(it) for it in iterables]
+ counts = Counter(chain.from_iterable(map(set, pool)))
+ uniques = {element for element in counts if counts[element] == 1}
+ return [list(filter(uniques.__contains__, it)) for it in pool]
def windowed(seq, n, fillvalue=None, step=1):
@@ -521,7 +894,35 @@ def windowed(seq, n, fillvalue=None, step=1):
>>> list(windowed(chain(padding, iterable), 3))
[(None, None, 1), (None, 1, 2), (1, 2, 3), (2, 3, 4)]
"""
- pass
+ if n < 0:
+ raise ValueError('n must be >= 0')
+ if n == 0:
+ yield ()
+ return
+ if step < 1:
+ raise ValueError('step must be >= 1')
+
+ iterable = iter(seq)
+
+ # Generate first window
+ window = deque(islice(iterable, n), maxlen=n)
+
+ # Deal with the first window not being full
+ if not window:
+ return
+ if len(window) < n:
+ yield tuple(window) + ((fillvalue,) * (n - len(window)))
+ return
+ yield tuple(window)
+
+ # Create the filler for the next windows. The padding ensures
+ # we have just enough elements to fill the last window.
+ padding = (fillvalue,) * (n - 1 if step >= n else step - 1)
+ filler = map(window.append, chain(iterable, padding))
+
+ # Generate the rest of the windows
+ for _ in islice(filler, step - 1, None, step):
+ yield tuple(window)
def substrings(iterable):
@@ -536,7 +937,18 @@ def substrings(iterable):
[(0,), (1,), (2,), (0, 1), (1, 2), (0, 1, 2)]
"""
- pass
+ # The length-1 substrings
+ seq = []
+ for item in iter(iterable):
+ seq.append(item)
+ yield (item,)
+ seq = tuple(seq)
+ item_count = len(seq)
+
+ # And the rest
+ for n in range(2, item_count + 1):
+ for i in range(item_count - n + 1):
+ yield seq[i : i + n]
def substrings_indexes(seq, reverse=False):
@@ -565,7 +977,12 @@ def substrings_indexes(seq, reverse=False):
"""
- pass
+ r = range(1, len(seq) + 1)
+ if reverse:
+ r = reversed(r)
+ return (
+ (seq[i : i + L], i, i + L) for L in r for i in range(len(seq) - L + 1)
+ )
class bucket:
@@ -613,12 +1030,14 @@ class bucket:
def __contains__(self, value):
if not self._validator(value):
return False
+
try:
item = next(self[value])
except StopIteration:
return False
else:
self._cache[value].appendleft(item)
+
return True
def _get_values(self, value):
@@ -627,18 +1046,38 @@ class bucket:
Items that don't match are stored in the local cache as they
are encountered.
"""
- pass
+ while True:
+ # If we've cached some items that match the target value, emit
+ # the first one and evict it from the cache.
+ if self._cache[value]:
+ yield self._cache[value].popleft()
+ # Otherwise we need to advance the parent iterator to search for
+ # a matching item, caching the rest.
+ else:
+ while True:
+ try:
+ item = next(self._it)
+ except StopIteration:
+ return
+ item_value = self._key(item)
+ if item_value == value:
+ yield item
+ break
+ elif self._validator(item_value):
+ self._cache[item_value].append(item)
def __iter__(self):
for item in self._it:
item_value = self._key(item)
if self._validator(item_value):
self._cache[item_value].append(item)
+
yield from self._cache.keys()
def __getitem__(self, value):
if not self._validator(value):
return iter(())
+
return self._get_values(value)
@@ -679,7 +1118,10 @@ def spy(iterable, n=1):
[1, 2, 3, 4, 5]
"""
- pass
+ it = iter(iterable)
+ head = take(n, it)
+
+ return head.copy(), chain(head, it)
def interleave(*iterables):
@@ -693,7 +1135,7 @@ def interleave(*iterables):
exhausted, see :func:`interleave_longest`.
"""
- pass
+ return chain.from_iterable(zip(*iterables))
def interleave_longest(*iterables):
@@ -708,7 +1150,8 @@ def interleave_longest(*iterables):
is large).
"""
- pass
+ i = chain.from_iterable(zip_longest(*iterables, fillvalue=_marker))
+ return (x for x in i if x is not _marker)
def interleave_evenly(iterables, lengths=None):
@@ -735,7 +1178,46 @@ def interleave_evenly(iterables, lengths=None):
Based on Bresenham's algorithm.
"""
- pass
+ if lengths is None:
+ try:
+ lengths = [len(it) for it in iterables]
+ except TypeError:
+ raise ValueError(
+ 'Iterable lengths could not be determined automatically. '
+ 'Specify them with the lengths keyword.'
+ )
+ elif len(iterables) != len(lengths):
+ raise ValueError('Mismatching number of iterables and lengths.')
+
+ dims = len(lengths)
+
+ # sort iterables by length, descending
+ lengths_permute = sorted(
+ range(dims), key=lambda i: lengths[i], reverse=True
+ )
+ lengths_desc = [lengths[i] for i in lengths_permute]
+ iters_desc = [iter(iterables[i]) for i in lengths_permute]
+
+ # the longest iterable is the primary one (Bresenham: the longest
+ # distance along an axis)
+ delta_primary, deltas_secondary = lengths_desc[0], lengths_desc[1:]
+ iter_primary, iters_secondary = iters_desc[0], iters_desc[1:]
+ errors = [delta_primary // dims] * len(deltas_secondary)
+
+ to_yield = sum(lengths)
+ while to_yield:
+ yield next(iter_primary)
+ to_yield -= 1
+ # update errors for each secondary iterable
+ errors = [e - delta for e, delta in zip(errors, deltas_secondary)]
+
+ # those iterables for which the error is negative are yielded
+ # ("diagonal step" in Bresenham)
+ for i, e_ in enumerate(errors):
+ if e_ < 0:
+ yield next(iters_secondary[i])
+ to_yield -= 1
+ errors[i] += delta_primary
def collapse(iterable, base_type=None, levels=None):
@@ -764,7 +1246,38 @@ def collapse(iterable, base_type=None, levels=None):
['a', ['b'], 'c', ['d']]
"""
- pass
+ stack = deque()
+ # Add our first node group, treat the iterable as a single node
+ stack.appendleft((0, repeat(iterable, 1)))
+
+ while stack:
+ node_group = stack.popleft()
+ level, nodes = node_group
+
+ # Check if beyond max level
+ if levels is not None and level > levels:
+ yield from nodes
+ continue
+
+ for node in nodes:
+ # Check if done iterating
+ if isinstance(node, (str, bytes)) or (
+ (base_type is not None) and isinstance(node, base_type)
+ ):
+ yield node
+ # Otherwise try to create child nodes
+ else:
+ try:
+ tree = iter(node)
+ except TypeError:
+ yield node
+ else:
+ # Save our current location
+ stack.appendleft(node_group)
+ # Append the new child node
+ stack.appendleft((level + 1, tree))
+ # Break to process child node
+ break
def side_effect(func, iterable, chunk_size=None, before=None, after=None):
@@ -811,7 +1324,21 @@ def side_effect(func, iterable, chunk_size=None, before=None, after=None):
True
"""
- pass
+ try:
+ if before is not None:
+ before()
+
+ if chunk_size is None:
+ for item in iterable:
+ func(item)
+ yield item
+ else:
+ for chunk in chunked(iterable, chunk_size):
+ func(chunk)
+ yield from chunk
+ finally:
+ if after is not None:
+ after()
def sliced(seq, n, strict=False):
@@ -834,7 +1361,18 @@ def sliced(seq, n, strict=False):
For non-sliceable iterables, see :func:`chunked`.
"""
- pass
+ iterator = takewhile(len, (seq[i : i + n] for i in count(0, n)))
+ if strict:
+
+ def ret():
+ for _slice in iterator:
+ if len(_slice) != n:
+ raise ValueError("seq is not divisible by n.")
+ yield _slice
+
+ return iter(ret())
+ else:
+ return iterator
def split_at(iterable, pred, maxsplit=-1, keep_separator=False):
@@ -860,7 +1398,25 @@ def split_at(iterable, pred, maxsplit=-1, keep_separator=False):
[['a'], ['b'], ['c', 'd', 'c'], ['b'], ['a']]
"""
- pass
+ if maxsplit == 0:
+ yield list(iterable)
+ return
+
+ buf = []
+ it = iter(iterable)
+ for item in it:
+ if pred(item):
+ yield buf
+ if keep_separator:
+ yield [item]
+ if maxsplit == 1:
+ yield list(it)
+ return
+ buf = []
+ maxsplit -= 1
+ else:
+ buf.append(item)
+ yield buf
def split_before(iterable, pred, maxsplit=-1):
@@ -879,7 +1435,23 @@ def split_before(iterable, pred, maxsplit=-1):
>>> list(split_before(range(10), lambda n: n % 3 == 0, maxsplit=2))
[[0, 1, 2], [3, 4, 5], [6, 7, 8, 9]]
"""
- pass
+ if maxsplit == 0:
+ yield list(iterable)
+ return
+
+ buf = []
+ it = iter(iterable)
+ for item in it:
+ if pred(item) and buf:
+ yield buf
+ if maxsplit == 1:
+ yield [item] + list(it)
+ return
+ buf = []
+ maxsplit -= 1
+ buf.append(item)
+ if buf:
+ yield buf
def split_after(iterable, pred, maxsplit=-1):
@@ -899,7 +1471,25 @@ def split_after(iterable, pred, maxsplit=-1):
[[0], [1, 2, 3], [4, 5, 6, 7, 8, 9]]
"""
- pass
+ if maxsplit == 0:
+ yield list(iterable)
+ return
+
+ buf = []
+ it = iter(iterable)
+ for item in it:
+ buf.append(item)
+ if pred(item) and buf:
+ yield buf
+ if maxsplit == 1:
+ buf = list(it)
+ if buf:
+ yield buf
+ return
+ buf = []
+ maxsplit -= 1
+ if buf:
+ yield buf
def split_when(iterable, pred, maxsplit=-1):
@@ -921,7 +1511,30 @@ def split_when(iterable, pred, maxsplit=-1):
[[1, 2, 3, 3], [2, 5], [2, 4, 2]]
"""
- pass
+ if maxsplit == 0:
+ yield list(iterable)
+ return
+
+ it = iter(iterable)
+ try:
+ cur_item = next(it)
+ except StopIteration:
+ return
+
+ buf = [cur_item]
+ for next_item in it:
+ if pred(cur_item, next_item):
+ yield buf
+ if maxsplit == 1:
+ yield [next_item] + list(it)
+ return
+ buf = []
+ maxsplit -= 1
+
+ buf.append(next_item)
+ cur_item = next_item
+
+ yield buf
def split_into(iterable, sizes):
@@ -957,7 +1570,16 @@ def split_into(iterable, sizes):
(e.g. a point represented by x,y,z) but, the format is not the same for
all columns.
"""
- pass
+ # convert the iterable argument into an iterator so its contents can
+ # be consumed by islice in case it is a generator
+ it = iter(iterable)
+
+ for size in sizes:
+ if size is None:
+ yield list(it)
+ return
+ else:
+ yield list(islice(it, size))
def padded(iterable, fillvalue=None, n=None, next_multiple=False):
@@ -984,7 +1606,25 @@ def padded(iterable, fillvalue=None, n=None, next_multiple=False):
[1, 2, 3, 4, 5]
"""
- pass
+ iterable = iter(iterable)
+ iterable_with_repeat = chain(iterable, repeat(fillvalue))
+
+ if n is None:
+ return iterable_with_repeat
+ elif n < 1:
+ raise ValueError('n must be at least 1')
+ elif next_multiple:
+
+ def slice_generator():
+ for first in iterable:
+ yield (first,)
+ yield islice(iterable_with_repeat, n - 1)
+
+ # While elements exist produce slices of size n
+ return chain.from_iterable(slice_generator())
+ else:
+ # Ensure the first batch is at least size n then iterate
+ return chain(islice(iterable_with_repeat, n), iterable)
def repeat_each(iterable, n=2):
@@ -993,7 +1633,7 @@ def repeat_each(iterable, n=2):
>>> list(repeat_each('ABC', 3))
['A', 'A', 'A', 'B', 'B', 'B', 'C', 'C', 'C']
"""
- pass
+ return chain.from_iterable(map(repeat, iterable, repeat(n)))
def repeat_last(iterable, default=None):
@@ -1008,7 +1648,11 @@ def repeat_last(iterable, default=None):
[42, 42, 42, 42, 42]
"""
- pass
+ item = _marker
+ for item in iterable:
+ yield item
+ final = default if item is _marker else item
+ yield from repeat(final)
def distribute(n, iterable):
@@ -1041,7 +1685,11 @@ def distribute(n, iterable):
original iterable, see :func:`divide`.
"""
- pass
+ if n < 1:
+ raise ValueError('n must be at least 1')
+
+ children = tee(iterable, n)
+ return [islice(it, index, None, n) for index, it in enumerate(children)]
def stagger(iterable, offsets=(-1, 0, 1), longest=False, fillvalue=None):
@@ -1065,7 +1713,11 @@ def stagger(iterable, offsets=(-1, 0, 1), longest=False, fillvalue=None):
sequence. Specify *fillvalue* to use some other value.
"""
- pass
+ children = tee(iterable, len(offsets))
+
+ return zip_offset(
+ *children, offsets=offsets, longest=longest, fillvalue=fillvalue
+ )
def zip_equal(*iterables):
@@ -1086,7 +1738,17 @@ def zip_equal(*iterables):
lengths
"""
- pass
+ if hexversion >= 0x30A00A6:
+ warnings.warn(
+ (
+ 'zip_equal will be removed in a future version of '
+ 'more-itertools. Use the builtin zip function with '
+ 'strict=True instead.'
+ ),
+ DeprecationWarning,
+ )
+
+ return _zip_equal(*iterables)
def zip_offset(*iterables, offsets, longest=False, fillvalue=None):
@@ -1110,11 +1772,27 @@ def zip_offset(*iterables, offsets, longest=False, fillvalue=None):
sequence. Specify *fillvalue* to use some other value.
"""
- pass
+ if len(iterables) != len(offsets):
+ raise ValueError("Number of iterables and offsets didn't match")
+
+ staggered = []
+ for it, n in zip(iterables, offsets):
+ if n < 0:
+ staggered.append(chain(repeat(fillvalue, -n), it))
+ elif n > 0:
+ staggered.append(islice(it, n, None))
+ else:
+ staggered.append(it)
+
+ if longest:
+ return zip_longest(*staggered, fillvalue=fillvalue)
+ return zip(*staggered)
-def sort_together(iterables, key_list=(0,), key=None, reverse=False, strict
- =False):
+
+def sort_together(
+ iterables, key_list=(0,), key=None, reverse=False, strict=False
+):
"""Return the input iterables sorted together, with *key_list* as the
priority for sorting. All iterables are trimmed to the length of the
shortest one.
@@ -1158,7 +1836,31 @@ def sort_together(iterables, key_list=(0,), key=None, reverse=False, strict
different lengths.
"""
- pass
+ if key is None:
+ # if there is no key function, the key argument to sorted is an
+ # itemgetter
+ key_argument = itemgetter(*key_list)
+ else:
+ # if there is a key function, call it with the items at the offsets
+ # specified by the key function as arguments
+ key_list = list(key_list)
+ if len(key_list) == 1:
+ # if key_list contains a single item, pass the item at that offset
+ # as the only argument to the key function
+ key_offset = key_list[0]
+ key_argument = lambda zipped_items: key(zipped_items[key_offset])
+ else:
+ # if key_list contains multiple items, use itemgetter to return a
+ # tuple of items, which we pass as *args to the key function
+ get_key_items = itemgetter(*key_list)
+ key_argument = lambda zipped_items: key(
+ *get_key_items(zipped_items)
+ )
+
+ zipper = zip_equal if strict else zip
+ return list(
+ zipper(*sorted(zipper(*iterables), key=key_argument, reverse=reverse))
+ )
def unzip(iterable):
@@ -1181,7 +1883,33 @@ def unzip(iterable):
:func:`itertools.tee` and thus may require significant storage.
"""
- pass
+ head, iterable = spy(iter(iterable))
+ if not head:
+ # empty iterable, e.g. zip([], [], [])
+ return ()
+ # spy returns a one-length iterable as head
+ head = head[0]
+ iterables = tee(iterable, len(head))
+
+ def itemgetter(i):
+ def getter(obj):
+ try:
+ return obj[i]
+ except IndexError:
+ # basically if we have an iterable like
+ # iter([(1, 2, 3), (4, 5), (6,)])
+ # the second unzipped iterable would fail at the third tuple
+ # since it would try to access tup[1]
+ # same with the third unzipped iterable and the second tuple
+ # to support these "improperly zipped" iterables,
+ # we create a custom itemgetter
+ # which just stops the unzipped iterables
+ # at first length mismatch
+ raise StopIteration
+
+ return getter
+
+ return tuple(map(itemgetter(i), it) for i, it in enumerate(iterables))
def divide(n, iterable):
@@ -1213,7 +1941,26 @@ def divide(n, iterable):
pull the iterable into memory.
"""
- pass
+ if n < 1:
+ raise ValueError('n must be at least 1')
+
+ try:
+ iterable[:0]
+ except TypeError:
+ seq = tuple(iterable)
+ else:
+ seq = iterable
+
+ q, r = divmod(len(seq), n)
+
+ ret = []
+ stop = 0
+ for i in range(1, n + 1):
+ start = stop
+ stop += q + 1 if i <= r else q
+ ret.append(iter(seq[start:stop]))
+
+ return ret
def always_iterable(obj, base_type=(str, bytes)):
@@ -1257,7 +2004,16 @@ def always_iterable(obj, base_type=(str, bytes)):
>>> list(always_iterable(obj, base_type=None))
['f', 'o', 'o']
"""
- pass
+ if obj is None:
+ return iter(())
+
+ if (base_type is not None) and isinstance(obj, base_type):
+ return iter((obj,))
+
+ try:
+ return iter(obj)
+ except TypeError:
+ return iter((obj,))
def adjacent(predicate, iterable, distance=1):
@@ -1288,7 +2044,15 @@ def adjacent(predicate, iterable, distance=1):
to group ranges of items with the same `bool` value.
"""
- pass
+ # Allow distance=0 mainly for testing that it reproduces results with map()
+ if distance < 0:
+ raise ValueError('distance must be at least 0')
+
+ i1, i2 = tee(iterable)
+ padding = [False] * distance
+ selected = chain(padding, map(predicate, i1), padding)
+ adjacent_to_selected = map(any, windowed(selected, 2 * distance + 1))
+ return zip(adjacent_to_selected, i2)
def groupby_transform(iterable, keyfunc=None, valuefunc=None, reducefunc=None):
@@ -1327,7 +2091,13 @@ def groupby_transform(iterable, keyfunc=None, valuefunc=None, reducefunc=None):
duplicate groups, you should sort the iterable by the key function.
"""
- pass
+ ret = groupby(iterable, keyfunc)
+ if valuefunc:
+ ret = ((k, map(valuefunc, g)) for k, g in ret)
+ if reducefunc:
+ ret = ((k, reducefunc(g)) for k, g in ret)
+
+ return ret
class numeric_range(abc.Sequence, abc.Hashable):
@@ -1381,12 +2151,13 @@ class numeric_range(abc.Sequence, abc.Hashable):
datetime.datetime(2019, 1, 2, 0, 0)
"""
+
_EMPTY_HASH = hash(range(0, 0))
def __init__(self, *args):
argc = len(args)
if argc == 1:
- self._stop, = args
+ (self._stop,) = args
self._start = type(self._stop)(0)
self._step = type(self._stop - self._start)(1)
elif argc == 2:
@@ -1396,12 +2167,15 @@ class numeric_range(abc.Sequence, abc.Hashable):
self._start, self._stop, self._step = args
elif argc == 0:
raise TypeError(
- 'numeric_range expected at least 1 argument, got {}'.format
- (argc))
+ 'numeric_range expected at least '
+ '1 argument, got {}'.format(argc)
+ )
else:
raise TypeError(
- 'numeric_range expected at most 3 arguments, got {}'.format
- (argc))
+ 'numeric_range expected at most '
+ '3 arguments, got {}'.format(argc)
+ )
+
self._zero = type(self._step)(0)
if self._step == self._zero:
raise ValueError('numeric_range() arg 3 must not be zero')
@@ -1417,8 +2191,10 @@ class numeric_range(abc.Sequence, abc.Hashable):
if self._growing:
if self._start <= elem < self._stop:
return (elem - self._start) % self._step == self._zero
- elif self._start >= elem > self._stop:
- return (self._start - elem) % -self._step == self._zero
+ else:
+ if self._start >= elem > self._stop:
+ return (self._start - elem) % (-self._step) == self._zero
+
return False
def __eq__(self, other):
@@ -1426,11 +2202,13 @@ class numeric_range(abc.Sequence, abc.Hashable):
empty_self = not bool(self)
empty_other = not bool(other)
if empty_self or empty_other:
- return empty_self and empty_other
+ return empty_self and empty_other # True if both empty
else:
- return (self._start == other._start and self._step == other
- ._step and self._get_by_index(-1) == other.
- _get_by_index(-1))
+ return (
+ self._start == other._start
+ and self._step == other._step
+ and self._get_by_index(-1) == other._get_by_index(-1)
+ )
else:
return False
@@ -1439,23 +2217,27 @@ class numeric_range(abc.Sequence, abc.Hashable):
return self._get_by_index(key)
elif isinstance(key, slice):
step = self._step if key.step is None else key.step * self._step
+
if key.start is None or key.start <= -self._len:
start = self._start
elif key.start >= self._len:
start = self._stop
- else:
+ else: # -self._len < key.start < self._len
start = self._get_by_index(key.start)
+
if key.stop is None or key.stop >= self._len:
stop = self._stop
elif key.stop <= -self._len:
stop = self._start
- else:
+ else: # -self._len < key.stop < self._len
stop = self._get_by_index(key.stop)
+
return numeric_range(start, stop, step)
else:
raise TypeError(
- 'numeric range indices must be integers or slices, not {}'.
- format(type(key).__name__))
+ 'numeric range indices must be '
+ 'integers or slices, not {}'.format(type(key).__name__)
+ )
def __hash__(self):
if self:
@@ -1464,7 +2246,7 @@ class numeric_range(abc.Sequence, abc.Hashable):
return self._EMPTY_HASH
def __iter__(self):
- values = (self._start + n * self._step for n in count())
+ values = (self._start + (n * self._step) for n in count())
if self._growing:
return takewhile(partial(gt, self._stop), values)
else:
@@ -1473,20 +2255,66 @@ class numeric_range(abc.Sequence, abc.Hashable):
def __len__(self):
return self._len
+ @cached_property
+ def _len(self):
+ if self._growing:
+ start = self._start
+ stop = self._stop
+ step = self._step
+ else:
+ start = self._stop
+ stop = self._start
+ step = -self._step
+ distance = stop - start
+ if distance <= self._zero:
+ return 0
+ else: # distance > 0 and step > 0: regular euclidean division
+ q, r = divmod(distance, step)
+ return int(q) + int(r != self._zero)
+
def __reduce__(self):
return numeric_range, (self._start, self._stop, self._step)
def __repr__(self):
if self._step == 1:
- return 'numeric_range({}, {})'.format(repr(self._start), repr(
- self._stop))
+ return "numeric_range({}, {})".format(
+ repr(self._start), repr(self._stop)
+ )
else:
- return 'numeric_range({}, {}, {})'.format(repr(self._start),
- repr(self._stop), repr(self._step))
+ return "numeric_range({}, {}, {})".format(
+ repr(self._start), repr(self._stop), repr(self._step)
+ )
def __reversed__(self):
- return iter(numeric_range(self._get_by_index(-1), self._start -
- self._step, -self._step))
+ return iter(
+ numeric_range(
+ self._get_by_index(-1), self._start - self._step, -self._step
+ )
+ )
+
+ def count(self, value):
+ return int(value in self)
+
+ def index(self, value):
+ if self._growing:
+ if self._start <= value < self._stop:
+ q, r = divmod(value - self._start, self._step)
+ if r == self._zero:
+ return int(q)
+ else:
+ if self._start >= value > self._stop:
+ q, r = divmod(self._start - value, -self._step)
+ if r == self._zero:
+ return int(q)
+
+ raise ValueError("{} is not in numeric range".format(value))
+
+ def _get_by_index(self, i):
+ if i < 0:
+ i += self._len
+ if i < 0 or i >= self._len:
+ raise IndexError("numeric range object index out of range")
+ return self._start + i * self._step
def count_cycle(iterable, n=None):
@@ -1498,7 +2326,11 @@ def count_cycle(iterable, n=None):
[(0, 'A'), (0, 'B'), (1, 'A'), (1, 'B'), (2, 'A'), (2, 'B')]
"""
- pass
+ iterable = tuple(iterable)
+ if not iterable:
+ return iter(())
+ counter = count() if n is None else range(n)
+ return ((i, item) for i in counter for item in iterable)
def mark_ends(iterable):
@@ -1521,7 +2353,21 @@ def mark_ends(iterable):
>>> print(total)
300
"""
- pass
+ it = iter(iterable)
+
+ try:
+ b = next(it)
+ except StopIteration:
+ return
+
+ try:
+ for i in count():
+ a = b
+ b = next(it)
+ yield i == 0, False, a
+
+ except StopIteration:
+ yield i == 0, True, a
def locate(iterable, pred=bool, window_size=None):
@@ -1562,7 +2408,14 @@ def locate(iterable, pred=bool, window_size=None):
106
"""
- pass
+ if window_size is None:
+ return compress(count(), map(pred, iterable))
+
+ if window_size < 1:
+ raise ValueError('window size must be at least 1')
+
+ it = windowed(iterable, window_size, fillvalue=_marker)
+ return compress(count(), starmap(pred, it))
def longest_common_prefix(iterables):
@@ -1572,7 +2425,7 @@ def longest_common_prefix(iterables):
'ab'
"""
- pass
+ return (c[0] for c in takewhile(all_equal, zip(*iterables)))
def lstrip(iterable, pred):
@@ -1590,7 +2443,7 @@ def lstrip(iterable, pred):
an wrapper for :func:`itertools.dropwhile`.
"""
- pass
+ return dropwhile(pred, iterable)
def rstrip(iterable, pred):
@@ -1607,7 +2460,16 @@ def rstrip(iterable, pred):
This function is analogous to :func:`str.rstrip`.
"""
- pass
+ cache = []
+ cache_append = cache.append
+ cache_clear = cache.clear
+ for x in iterable:
+ if pred(x):
+ cache_append(x)
+ else:
+ yield from cache
+ cache_clear()
+ yield x
def strip(iterable, pred):
@@ -1624,7 +2486,7 @@ def strip(iterable, pred):
This function is analogous to :func:`str.strip`.
"""
- pass
+ return rstrip(lstrip(iterable, pred), pred)
class islice_extended:
@@ -1669,9 +2531,106 @@ class islice_extended:
def __getitem__(self, key):
if isinstance(key, slice):
return islice_extended(_islice_helper(self._iterable, key))
+
raise TypeError('islice_extended.__getitem__ argument must be a slice')
+def _islice_helper(it, s):
+ start = s.start
+ stop = s.stop
+ if s.step == 0:
+ raise ValueError('step argument must be a non-zero integer or None.')
+ step = s.step or 1
+
+ if step > 0:
+ start = 0 if (start is None) else start
+
+ if start < 0:
+ # Consume all but the last -start items
+ cache = deque(enumerate(it, 1), maxlen=-start)
+ len_iter = cache[-1][0] if cache else 0
+
+ # Adjust start to be positive
+ i = max(len_iter + start, 0)
+
+ # Adjust stop to be positive
+ if stop is None:
+ j = len_iter
+ elif stop >= 0:
+ j = min(stop, len_iter)
+ else:
+ j = max(len_iter + stop, 0)
+
+ # Slice the cache
+ n = j - i
+ if n <= 0:
+ return
+
+ for index, item in islice(cache, 0, n, step):
+ yield item
+ elif (stop is not None) and (stop < 0):
+ # Advance to the start position
+ next(islice(it, start, start), None)
+
+ # When stop is negative, we have to carry -stop items while
+ # iterating
+ cache = deque(islice(it, -stop), maxlen=-stop)
+
+ for index, item in enumerate(it):
+ cached_item = cache.popleft()
+ if index % step == 0:
+ yield cached_item
+ cache.append(item)
+ else:
+ # When both start and stop are positive we have the normal case
+ yield from islice(it, start, stop, step)
+ else:
+ start = -1 if (start is None) else start
+
+ if (stop is not None) and (stop < 0):
+ # Consume all but the last items
+ n = -stop - 1
+ cache = deque(enumerate(it, 1), maxlen=n)
+ len_iter = cache[-1][0] if cache else 0
+
+ # If start and stop are both negative they are comparable and
+ # we can just slice. Otherwise we can adjust start to be negative
+ # and then slice.
+ if start < 0:
+ i, j = start, stop
+ else:
+ i, j = min(start - len_iter, -1), None
+
+ for index, item in list(cache)[i:j:step]:
+ yield item
+ else:
+ # Advance to the stop position
+ if stop is not None:
+ m = stop + 1
+ next(islice(it, m, m), None)
+
+ # stop is positive, so if start is negative they are not comparable
+ # and we need the rest of the items.
+ if start < 0:
+ i = start
+ n = None
+ # stop is None and start is positive, so we just need items up to
+ # the start index.
+ elif stop is None:
+ i = None
+ n = start + 1
+ # Both stop and start are positive, so they are comparable.
+ else:
+ i = None
+ n = start - stop
+ if n <= 0:
+ return
+
+ cache = list(islice(it, n))
+
+ yield from cache[i::step]
+
+
def always_reversible(iterable):
"""An extension of :func:`reversed` that supports all iterables, not
just those which implement the ``Reversible`` or ``Sequence`` protocols.
@@ -1684,7 +2643,10 @@ def always_reversible(iterable):
this function will cache the remaining items in the iterable and
yield them in reverse order, which may require significant storage.
"""
- pass
+ try:
+ return reversed(iterable)
+ except TypeError:
+ return reversed(list(iterable))
def consecutive_groups(iterable, ordering=lambda x: x):
@@ -1729,7 +2691,10 @@ def consecutive_groups(iterable, ordering=lambda x: x):
[[1, 2], [11, 12], [21, 22]]
"""
- pass
+ for k, g in groupby(
+ enumerate(iterable), key=lambda x: x[0] - ordering(x[1])
+ ):
+ yield map(itemgetter(1), g)
def difference(iterable, func=sub, *, initial=None):
@@ -1762,7 +2727,16 @@ def difference(iterable, func=sub, *, initial=None):
[1, 2, 3]
"""
- pass
+ a, b = tee(iterable)
+ try:
+ first = [next(b)]
+ except StopIteration:
+ return iter([])
+
+ if initial is not None:
+ first = []
+
+ return chain(first, map(func, b, a))
class SequenceView(Sequence):
@@ -1931,6 +2905,7 @@ class seekable:
else:
self._index += 1
return item
+
item = next(self._source)
self._cache.append(item)
return item
@@ -1942,6 +2917,33 @@ class seekable:
return False
return True
+ def peek(self, default=_marker):
+ try:
+ peeked = next(self)
+ except StopIteration:
+ if default is _marker:
+ raise
+ return default
+ if self._index is None:
+ self._index = len(self._cache)
+ self._index -= 1
+ return peeked
+
+ def elements(self):
+ return SequenceView(self._cache)
+
+ def seek(self, index):
+ self._index = index
+ remainder = index - len(self._cache)
+ if remainder > 0:
+ consume(self, remainder)
+
+ def relative_seek(self, count):
+ if self._index is None:
+ self._index = len(self._cache)
+
+ self.seek(max(self._index + count, 0))
+
class run_length:
"""
@@ -1963,6 +2965,14 @@ class run_length:
"""
+ @staticmethod
+ def encode(iterable):
+ return ((k, ilen(g)) for k, g in groupby(iterable))
+
+ @staticmethod
+ def decode(iterable):
+ return chain.from_iterable(starmap(repeat, iterable))
+
def exactly_n(iterable, n, predicate=bool):
"""Return ``True`` if exactly ``n`` items in the iterable are ``True``
@@ -1979,7 +2989,7 @@ def exactly_n(iterable, n, predicate=bool):
so avoid calling it on infinite iterables.
"""
- pass
+ return len(take(n + 1, filter(predicate, iterable))) == n
def circular_shifts(iterable, steps=1):
@@ -1998,7 +3008,18 @@ def circular_shifts(iterable, steps=1):
[(0, 1, 2, 3), (3, 0, 1, 2), (2, 3, 0, 1), (1, 2, 3, 0)]
"""
- pass
+ buffer = deque(iterable)
+ if steps == 0:
+ raise ValueError('Steps should be a non-zero integer')
+
+ buffer.rotate(steps)
+ steps = -steps
+ n = len(buffer)
+ n //= math.gcd(n, steps)
+
+ for __ in repeat(None, n):
+ buffer.rotate(steps)
+ yield tuple(buffer)
def make_decorator(wrapping_func, result_index=0):
@@ -2049,7 +3070,22 @@ def make_decorator(wrapping_func, result_index=0):
'7'
"""
- pass
+
+ # See https://sites.google.com/site/bbayles/index/decorator_factory for
+ # notes on how this works.
+ def decorator(*wrapping_args, **wrapping_kwargs):
+ def outer_wrapper(f):
+ def inner_wrapper(*args, **kwargs):
+ result = f(*args, **kwargs)
+ wrapping_args_ = list(wrapping_args)
+ wrapping_args_.insert(result_index, result)
+ return wrapping_func(*wrapping_args_, **wrapping_kwargs)
+
+ return inner_wrapper
+
+ return outer_wrapper
+
+ return decorator
def map_reduce(iterable, keyfunc, valuefunc=None, reducefunc=None):
@@ -2103,7 +3139,20 @@ def map_reduce(iterable, keyfunc, valuefunc=None, reducefunc=None):
dictionary.
"""
- pass
+ valuefunc = (lambda x: x) if (valuefunc is None) else valuefunc
+
+ ret = defaultdict(list)
+ for item in iterable:
+ key = keyfunc(item)
+ value = valuefunc(item)
+ ret[key].append(value)
+
+ if reducefunc is not None:
+ for key, value_list in ret.items():
+ ret[key] = reducefunc(value_list)
+
+ ret.default_factory = None
+ return ret
def rlocate(iterable, pred=bool, window_size=None):
@@ -2139,7 +3188,14 @@ def rlocate(iterable, pred=bool, window_size=None):
See :func:`locate` to for other example applications.
"""
- pass
+ if window_size is None:
+ try:
+ len_iter = len(iterable)
+ return (len_iter - i - 1 for i in locate(reversed(iterable), pred))
+ except TypeError:
+ pass
+
+ return reversed(list(locate(iterable, pred, window_size)))
def replace(iterable, pred, substitutes, count=None, window_size=1):
@@ -2171,7 +3227,36 @@ def replace(iterable, pred, substitutes, count=None, window_size=1):
[3, 4, 5, 3, 4, 5]
"""
- pass
+ if window_size < 1:
+ raise ValueError('window_size must be at least 1')
+
+ # Save the substitutes iterable, since it's used more than once
+ substitutes = tuple(substitutes)
+
+ # Add padding such that the number of windows matches the length of the
+ # iterable
+ it = chain(iterable, [_marker] * (window_size - 1))
+ windows = windowed(it, window_size)
+
+ n = 0
+ for w in windows:
+ # If the current window matches our predicate (and we haven't hit
+ # our maximum number of replacements), splice in the substitutes
+ # and then consume the following windows that overlap with this one.
+ # For example, if the iterable is (0, 1, 2, 3, 4...)
+ # and the window size is 2, we have (0, 1), (1, 2), (2, 3)...
+ # If the predicate matches on (0, 1), we need to zap (0, 1) and (1, 2)
+ if pred(*w):
+ if (count is None) or (n < count):
+ n += 1
+ yield from substitutes
+ consume(windows, window_size - 1)
+ continue
+
+ # If there was no match (or we've reached the replacement limit),
+ # yield the first item from the window.
+ if w and (w[0] is not _marker):
+ yield w[0]
def partitions(iterable):
@@ -2188,7 +3273,10 @@ def partitions(iterable):
This is unrelated to :func:`partition`.
"""
- pass
+ sequence = list(iterable)
+ n = len(sequence)
+ for i in powerset(range(1, n)):
+ yield [sequence[i:j] for i, j in zip((0,) + i, i + (n,))]
def set_partitions(iterable, k=None, min_size=None, max_size=None):
@@ -2230,7 +3318,46 @@ def set_partitions(iterable, k=None, min_size=None, max_size=None):
['a', 'b', 'c']
"""
- pass
+ L = list(iterable)
+ n = len(L)
+ if k is not None:
+ if k < 1:
+ raise ValueError(
+ "Can't partition in a negative or zero number of groups"
+ )
+ elif k > n:
+ return
+
+ min_size = min_size if min_size is not None else 0
+ max_size = max_size if max_size is not None else n
+ if min_size > max_size:
+ return
+
+ def set_partitions_helper(L, k):
+ n = len(L)
+ if k == 1:
+ yield [L]
+ elif n == k:
+ yield [[s] for s in L]
+ else:
+ e, *M = L
+ for p in set_partitions_helper(M, k - 1):
+ yield [[e], *p]
+ for p in set_partitions_helper(M, k):
+ for i in range(len(p)):
+ yield p[:i] + [[e] + p[i]] + p[i + 1 :]
+
+ if k is None:
+ for k in range(1, n + 1):
+ yield from filter(
+ lambda z: all(min_size <= len(bk) <= max_size for bk in z),
+ set_partitions_helper(L, k),
+ )
+ else:
+ yield from filter(
+ lambda z: all(min_size <= len(bk) <= max_size for bk in z),
+ set_partitions_helper(L, k),
+ )
class time_limited:
@@ -2279,6 +3406,7 @@ class time_limited:
if monotonic() - self._start_time > self.limit_seconds:
self.timed_out = True
raise StopIteration
+
return item
@@ -2306,7 +3434,55 @@ def only(iterable, default=None, too_long=None):
is only one item. See :func:`spy` or :func:`peekable` to check
iterable contents less destructively.
"""
- pass
+ it = iter(iterable)
+ first_value = next(it, default)
+
+ try:
+ second_value = next(it)
+ except StopIteration:
+ pass
+ else:
+ msg = (
+ 'Expected exactly one item in iterable, but got {!r}, {!r}, '
+ 'and perhaps more.'.format(first_value, second_value)
+ )
+ raise too_long or ValueError(msg)
+
+ return first_value
+
+
+def _ichunk(iterable, n):
+ cache = deque()
+ chunk = islice(iterable, n)
+
+ def generator():
+ while True:
+ if cache:
+ yield cache.popleft()
+ else:
+ try:
+ item = next(chunk)
+ except StopIteration:
+ return
+ else:
+ yield item
+
+ def materialize_next(n=1):
+ # if n not specified materialize everything
+ if n is None:
+ cache.extend(chunk)
+ return len(cache)
+
+ to_cache = n - len(cache)
+
+ # materialize up to n
+ if to_cache > 0:
+ cache.extend(islice(chunk, to_cache))
+
+ # return number materialized up to n
+ return min(n, len(cache))
+
+ return (generator(), materialize_next)
def ichunked(iterable, n):
@@ -2330,7 +3506,19 @@ def ichunked(iterable, n):
[8, 9, 10, 11]
"""
- pass
+ iterable = iter(iterable)
+ while True:
+ # Create new chunk
+ chunk, materialize_next = _ichunk(iterable, n)
+
+ # Check to see whether we're at the end of the source iterable
+ if not materialize_next():
+ return
+
+ yield chunk
+
+ # Fill previous chunk's cache
+ materialize_next(None)
def iequals(*iterables):
@@ -2350,7 +3538,7 @@ def iequals(*iterables):
elements of iterable are equal to each other.
"""
- pass
+ return all(map(all_equal, zip_longest(*iterables, fillvalue=object())))
def distinct_combinations(iterable, r):
@@ -2364,7 +3552,33 @@ def distinct_combinations(iterable, r):
efficient.
"""
- pass
+ if r < 0:
+ raise ValueError('r must be non-negative')
+ elif r == 0:
+ yield ()
+ return
+ pool = tuple(iterable)
+ generators = [unique_everseen(enumerate(pool), key=itemgetter(1))]
+ current_combo = [None] * r
+ level = 0
+ while generators:
+ try:
+ cur_idx, p = next(generators[-1])
+ except StopIteration:
+ generators.pop()
+ level -= 1
+ continue
+ current_combo[level] = p
+ if level + 1 == r:
+ yield tuple(current_combo)
+ else:
+ generators.append(
+ unique_everseen(
+ enumerate(pool[cur_idx + 1 :], cur_idx + 1),
+ key=itemgetter(1),
+ )
+ )
+ level += 1
def filter_except(validator, iterable, *exceptions):
@@ -2382,7 +3596,13 @@ def filter_except(validator, iterable, *exceptions):
If an exception other than one given by *exceptions* is raised by
*validator*, it is raised like normal.
"""
- pass
+ for item in iterable:
+ try:
+ validator(item)
+ except exceptions:
+ pass
+ else:
+ yield item
def map_except(function, iterable, *exceptions):
@@ -2399,7 +3619,11 @@ def map_except(function, iterable, *exceptions):
If an exception other than one given by *exceptions* is raised by
*function*, it is raised like normal.
"""
- pass
+ for item in iterable:
+ try:
+ yield function(item)
+ except exceptions:
+ pass
def map_if(iterable, pred, func, func_else=lambda x: x):
@@ -2420,7 +3644,100 @@ def map_if(iterable, pred, func, func_else=lambda x: x):
... lambda x: f'{sqrt(x):.2f}', lambda x: None))
[None, None, None, None, None, '0.00', '1.00', '1.41', '1.73', '2.00']
"""
- pass
+ for item in iterable:
+ yield func(item) if pred(item) else func_else(item)
+
+
+def _sample_unweighted(iterator, k, strict):
+ # Algorithm L in the 1994 paper by Kim-Hung Li:
+ # "Reservoir-Sampling Algorithms of Time Complexity O(n(1+log(N/n)))".
+
+ reservoir = list(islice(iterator, k))
+ if strict and len(reservoir) < k:
+ raise ValueError('Sample larger than population')
+ W = 1.0
+
+ with suppress(StopIteration):
+ while True:
+ W *= exp(log(random()) / k)
+ skip = floor(log(random()) / log1p(-W))
+ element = next(islice(iterator, skip, None))
+ reservoir[randrange(k)] = element
+
+ shuffle(reservoir)
+ return reservoir
+
+
+def _sample_weighted(iterator, k, weights, strict):
+ # Implementation of "A-ExpJ" from the 2006 paper by Efraimidis et al. :
+ # "Weighted random sampling with a reservoir".
+
+ # Log-transform for numerical stability for weights that are small/large
+ weight_keys = (log(random()) / weight for weight in weights)
+
+ # Fill up the reservoir (collection of samples) with the first `k`
+ # weight-keys and elements, then heapify the list.
+ reservoir = take(k, zip(weight_keys, iterator))
+ if strict and len(reservoir) < k:
+ raise ValueError('Sample larger than population')
+
+ heapify(reservoir)
+
+ # The number of jumps before changing the reservoir is a random variable
+ # with an exponential distribution. Sample it using random() and logs.
+ smallest_weight_key, _ = reservoir[0]
+ weights_to_skip = log(random()) / smallest_weight_key
+
+ for weight, element in zip(weights, iterator):
+ if weight >= weights_to_skip:
+ # The notation here is consistent with the paper, but we store
+ # the weight-keys in log-space for better numerical stability.
+ smallest_weight_key, _ = reservoir[0]
+ t_w = exp(weight * smallest_weight_key)
+ r_2 = uniform(t_w, 1) # generate U(t_w, 1)
+ weight_key = log(r_2) / weight
+ heapreplace(reservoir, (weight_key, element))
+ smallest_weight_key, _ = reservoir[0]
+ weights_to_skip = log(random()) / smallest_weight_key
+ else:
+ weights_to_skip -= weight
+
+ ret = [element for weight_key, element in reservoir]
+ shuffle(ret)
+ return ret
+
+
+def _sample_counted(population, k, counts, strict):
+ element = None
+ remaining = 0
+
+ def feed(i):
+ # Advance *i* steps ahead and consume an element
+ nonlocal element, remaining
+
+ while i + 1 > remaining:
+ i = i - remaining
+ element = next(population)
+ remaining = next(counts)
+ remaining -= i + 1
+ return element
+
+ with suppress(StopIteration):
+ reservoir = []
+ for _ in range(k):
+ reservoir.append(feed(0))
+ if strict and len(reservoir) < k:
+ raise ValueError('Sample larger than population')
+
+ W = 1.0
+ while True:
+ W *= exp(log(random()) / k)
+ skip = floor(log(random()) / log1p(-W))
+ element = feed(skip)
+ reservoir[randrange(k)] = element
+
+ shuffle(reservoir)
+ return reservoir
def sample(iterable, k, weights=None, *, counts=None, strict=False):
@@ -2461,7 +3778,27 @@ def sample(iterable, k, weights=None, *, counts=None, strict=False):
technique is used. When *weights* are provided,
`Algorithm A-ExpJ <https://w.wiki/ANrS>`__ is used.
"""
- pass
+ iterator = iter(iterable)
+
+ if k < 0:
+ raise ValueError('k must be non-negative')
+
+ if k == 0:
+ return []
+
+ if weights is not None and counts is not None:
+ raise TypeError('weights and counts are mutally exclusive')
+
+ elif weights is not None:
+ weights = iter(weights)
+ return _sample_weighted(iterator, k, weights, strict)
+
+ elif counts is not None:
+ counts = iter(counts)
+ return _sample_counted(iterator, k, counts, strict)
+
+ else:
+ return _sample_unweighted(iterator, k, strict)
def is_sorted(iterable, key=None, reverse=False, strict=False):
@@ -2487,7 +3824,12 @@ def is_sorted(iterable, key=None, reverse=False, strict=False):
:func:`sorted` function for objects with unusual comparison dynamics.
If there are no out-of-order items, the iterable is exhausted.
"""
- pass
+ compare = le if strict else lt
+ it = iterable if (key is None) else map(key, iterable)
+ it_1, it_2 = tee(it)
+ next(it_2 if reverse else it_1, None)
+
+ return not any(map(compare, it_1, it_2))
class AbortThread(BaseException):
@@ -2551,8 +3893,10 @@ class callback_iter:
self._aborted = False
self._future = None
self._wait_seconds = wait_seconds
- self._executor = __import__('concurrent.futures'
- ).futures.ThreadPoolExecutor(max_workers=1)
+ # Lazily import concurrent.future
+ self._executor = __import__(
+ 'concurrent.futures'
+ ).futures.ThreadPoolExecutor(max_workers=1)
self._iterator = self._reader()
def __enter__(self):
@@ -2568,6 +3912,56 @@ class callback_iter:
def __next__(self):
return next(self._iterator)
+ @property
+ def done(self):
+ if self._future is None:
+ return False
+ return self._future.done()
+
+ @property
+ def result(self):
+ if not self.done:
+ raise RuntimeError('Function has not yet completed')
+
+ return self._future.result()
+
+ def _reader(self):
+ q = Queue()
+
+ def callback(*args, **kwargs):
+ if self._aborted:
+ raise AbortThread('canceled by user')
+
+ q.put((args, kwargs))
+
+ self._future = self._executor.submit(
+ self._func, **{self._callback_kwd: callback}
+ )
+
+ while True:
+ try:
+ item = q.get(timeout=self._wait_seconds)
+ except Empty:
+ pass
+ else:
+ q.task_done()
+ yield item
+
+ if self._future.done():
+ break
+
+ remaining = []
+ while True:
+ try:
+ item = q.get_nowait()
+ except Empty:
+ break
+ else:
+ q.task_done()
+ remaining.append(item)
+ q.join()
+ yield from remaining
+
def windowed_complete(iterable, n):
"""
@@ -2593,7 +3987,20 @@ def windowed_complete(iterable, n):
This function will exhaust the iterable and may require significant
storage.
"""
- pass
+ if n < 0:
+ raise ValueError('n must be >= 0')
+
+ seq = tuple(iterable)
+ size = len(seq)
+
+ if n > size:
+ raise ValueError('n must be <= len(seq)')
+
+ for i in range(size - n + 1):
+ beginning = seq[:i]
+ middle = seq[i : i + n]
+ end = seq[i + n :]
+ yield beginning, middle, end
def all_unique(iterable, key=None):
@@ -2615,7 +4022,20 @@ def all_unique(iterable, key=None):
encountered. Iterables with a mix of hashable and unhashable items can
be used, but the function will be slower for unhashable items.
"""
- pass
+ seenset = set()
+ seenset_add = seenset.add
+ seenlist = []
+ seenlist_add = seenlist.append
+ for element in map(key, iterable) if key else iterable:
+ try:
+ if element in seenset:
+ return False
+ seenset_add(element)
+ except TypeError:
+ if element in seenlist:
+ return False
+ seenlist_add(element)
+ return True
def nth_product(index, *args):
@@ -2630,7 +4050,23 @@ def nth_product(index, *args):
``IndexError`` will be raised if the given *index* is invalid.
"""
- pass
+ pools = list(map(tuple, reversed(args)))
+ ns = list(map(len, pools))
+
+ c = reduce(mul, ns)
+
+ if index < 0:
+ index += c
+
+ if not 0 <= index < c:
+ raise IndexError
+
+ result = []
+ for pool, n in zip(pools, ns):
+ result.append(pool[index % n])
+ index //= n
+
+ return tuple(reversed(result))
def nth_permutation(iterable, r, index):
@@ -2648,7 +4084,33 @@ def nth_permutation(iterable, r, index):
of *iterable*.
``IndexError`` will be raised if the given *index* is invalid.
"""
- pass
+ pool = list(iterable)
+ n = len(pool)
+
+ if r is None or r == n:
+ r, c = n, factorial(n)
+ elif not 0 <= r < n:
+ raise ValueError
+ else:
+ c = perm(n, r)
+ assert c > 0 # factortial(n)>0, and r<n so perm(n,r) is never zero
+
+ if index < 0:
+ index += c
+
+ if not 0 <= index < c:
+ raise IndexError
+
+ result = [0] * r
+ q = index * factorial(n) // c if r < n else index
+ for d in range(1, n + 1):
+ q, i = divmod(q, d)
+ if 0 <= n - d < r:
+ result[n - d] = i
+ if q == 0:
+ break
+
+ return tuple(map(pool.pop, result))
def nth_combination_with_replacement(iterable, r, index):
@@ -2668,7 +4130,33 @@ def nth_combination_with_replacement(iterable, r, index):
of *iterable*.
``IndexError`` will be raised if the given *index* is invalid.
"""
- pass
+ pool = tuple(iterable)
+ n = len(pool)
+ if (r < 0) or (r > n):
+ raise ValueError
+
+ c = comb(n + r - 1, r)
+
+ if index < 0:
+ index += c
+
+ if (index < 0) or (index >= c):
+ raise IndexError
+
+ result = []
+ i = 0
+ while r:
+ r -= 1
+ while n >= 0:
+ num_combs = comb(n + r - 1, r)
+ if index < num_combs:
+ break
+ n -= 1
+ i += 1
+ index -= num_combs
+ result.append(pool[i])
+
+ return tuple(result)
def value_chain(*args):
@@ -2695,7 +4183,14 @@ def value_chain(*args):
Multiple levels of nesting are not flattened.
"""
- pass
+ for value in args:
+ if isinstance(value, (str, bytes)):
+ yield value
+ continue
+ try:
+ yield from value
+ except TypeError:
+ yield value
def product_index(element, *args):
@@ -2711,7 +4206,16 @@ def product_index(element, *args):
``ValueError`` will be raised if the given *element* isn't in the product
of *args*.
"""
- pass
+ index = 0
+
+ for x, pool in zip_longest(element, args, fillvalue=_marker):
+ if x is _marker or pool is _marker:
+ raise ValueError('element is not a product of args')
+
+ pool = tuple(pool)
+ index = index * len(pool) + pool.index(x)
+
+ return index
def combination_index(element, iterable):
@@ -2727,7 +4231,34 @@ def combination_index(element, iterable):
``ValueError`` will be raised if the given *element* isn't one of the
combinations of *iterable*.
"""
- pass
+ element = enumerate(element)
+ k, y = next(element, (None, None))
+ if k is None:
+ return 0
+
+ indexes = []
+ pool = enumerate(iterable)
+ for n, x in pool:
+ if x == y:
+ indexes.append(n)
+ tmp, y = next(element, (None, None))
+ if tmp is None:
+ break
+ else:
+ k = tmp
+ else:
+ raise ValueError('element is not a combination of iterable')
+
+ n, _ = last(pool, default=(n, None))
+
+ # Python versions below 3.8 don't have math.comb
+ index = 1
+ for i, j in enumerate(reversed(indexes), start=1):
+ j = n - j
+ if i <= j:
+ index += comb(j, i)
+
+ return comb(n + 1, k + 1) - index
def combination_with_replacement_index(element, iterable):
@@ -2745,7 +4276,46 @@ def combination_with_replacement_index(element, iterable):
``ValueError`` will be raised if the given *element* isn't one of the
combinations with replacement of *iterable*.
"""
- pass
+ element = tuple(element)
+ l = len(element)
+ element = enumerate(element)
+
+ k, y = next(element, (None, None))
+ if k is None:
+ return 0
+
+ indexes = []
+ pool = tuple(iterable)
+ for n, x in enumerate(pool):
+ while x == y:
+ indexes.append(n)
+ tmp, y = next(element, (None, None))
+ if tmp is None:
+ break
+ else:
+ k = tmp
+ if y is None:
+ break
+ else:
+ raise ValueError(
+ 'element is not a combination with replacement of iterable'
+ )
+
+ n = len(pool)
+ occupations = [0] * n
+ for p in indexes:
+ occupations[p] += 1
+
+ index = 0
+ cumulative_sum = 0
+ for k in range(1, n):
+ cumulative_sum += occupations[k - 1]
+ j = l + n - 1 - k - cumulative_sum
+ i = n - k
+ if i <= j:
+ index += comb(j, i)
+
+ return index
def permutation_index(element, iterable):
@@ -2762,7 +4332,14 @@ def permutation_index(element, iterable):
``ValueError`` will be raised if the given *element* isn't one of the
permutations of *iterable*.
"""
- pass
+ index = 0
+ pool = list(iterable)
+ for i, x in zip(range(len(pool), -1, -1), element):
+ r = pool.index(x)
+ index = index * i + r
+ del pool[r]
+
+ return index
class countable:
@@ -2793,6 +4370,7 @@ class countable:
def __next__(self):
item = next(self._it)
self.items_seen += 1
+
return item
@@ -2809,7 +4387,41 @@ def chunked_even(iterable, n):
[[1, 2, 3], [4, 5, 6], [7]]
"""
- pass
+ iterable = iter(iterable)
+
+ # Initialize a buffer to process the chunks while keeping
+ # some back to fill any underfilled chunks
+ min_buffer = (n - 1) * (n - 2)
+ buffer = list(islice(iterable, min_buffer))
+
+ # Append items until we have a completed chunk
+ for _ in islice(map(buffer.append, iterable), n, None, n):
+ yield buffer[:n]
+ del buffer[:n]
+
+ # Check if any chunks need addition processing
+ if not buffer:
+ return
+ length = len(buffer)
+
+ # Chunks are either size `full_size <= n` or `partial_size = full_size - 1`
+ q, r = divmod(length, n)
+ num_lists = q + (1 if r > 0 else 0)
+ q, r = divmod(length, num_lists)
+ full_size = q + (1 if r > 0 else 0)
+ partial_size = full_size - 1
+ num_full = length - partial_size * num_lists
+
+ # Yield chunks of full size
+ partial_start_idx = num_full * full_size
+ if full_size > 0:
+ for i in range(0, partial_start_idx, full_size):
+ yield buffer[i : i + full_size]
+
+ # Yield chunks of partial size
+ if partial_size > 0:
+ for i in range(partial_start_idx, length, partial_size):
+ yield buffer[i : i + partial_size]
def zip_broadcast(*objects, scalar_types=(str, bytes), strict=False):
@@ -2833,7 +4445,39 @@ def zip_broadcast(*objects, scalar_types=(str, bytes), strict=False):
``UnequalIterablesError`` will be raised if any of the iterables have
different lengths.
"""
- pass
+
+ def is_scalar(obj):
+ if scalar_types and isinstance(obj, scalar_types):
+ return True
+ try:
+ iter(obj)
+ except TypeError:
+ return True
+ else:
+ return False
+
+ size = len(objects)
+ if not size:
+ return
+
+ new_item = [None] * size
+ iterables, iterable_positions = [], []
+ for i, obj in enumerate(objects):
+ if is_scalar(obj):
+ new_item[i] = obj
+ else:
+ iterables.append(iter(obj))
+ iterable_positions.append(i)
+
+ if not iterables:
+ yield tuple(objects)
+ return
+
+ zipper = _zip_equal if strict else zip
+ for item in zipper(*iterables):
+ for i, new_item[i] in zip(iterable_positions, item):
+ pass
+ yield tuple(new_item)
def unique_in_window(iterable, n, key=None):
@@ -2853,7 +4497,26 @@ def unique_in_window(iterable, n, key=None):
The items in *iterable* must be hashable.
"""
- pass
+ if n <= 0:
+ raise ValueError('n must be greater than 0')
+
+ window = deque(maxlen=n)
+ counts = defaultdict(int)
+ use_key = key is not None
+
+ for item in iterable:
+ if len(window) == n:
+ to_discard = window[0]
+ if counts[to_discard] == 1:
+ del counts[to_discard]
+ else:
+ counts[to_discard] -= 1
+
+ k = key(item) if use_key else item
+ if k not in counts:
+ yield item
+ counts[k] += 1
+ window.append(k)
def duplicates_everseen(iterable, key=None):
@@ -2868,7 +4531,22 @@ def duplicates_everseen(iterable, key=None):
the same performance considerations.
"""
- pass
+ seen_set = set()
+ seen_list = []
+ use_key = key is not None
+
+ for element in iterable:
+ k = key(element) if use_key else element
+ try:
+ if k not in seen_set:
+ seen_set.add(k)
+ else:
+ yield element
+ except TypeError:
+ if k not in seen_list:
+ seen_list.append(k)
+ else:
+ yield element
def duplicates_justseen(iterable, key=None):
@@ -2882,7 +4560,7 @@ def duplicates_justseen(iterable, key=None):
This function is analogous to :func:`unique_justseen`.
"""
- pass
+ return flatten(g for _, g in groupby(iterable, key) for _ in g)
def classify_unique(iterable, key=None):
@@ -2906,7 +4584,25 @@ def classify_unique(iterable, key=None):
the same performance considerations.
"""
- pass
+ seen_set = set()
+ seen_list = []
+ use_key = key is not None
+ previous = None
+
+ for i, element in enumerate(iterable):
+ k = key(element) if use_key else element
+ is_unique_justseen = not i or previous != k
+ previous = k
+ is_unique_everseen = False
+ try:
+ if k not in seen_set:
+ seen_set.add(k)
+ is_unique_everseen = True
+ except TypeError:
+ if k not in seen_list:
+ seen_list.append(k)
+ is_unique_everseen = True
+ yield element, is_unique_justseen, is_unique_everseen
def minmax(iterable_or_value, *others, key=None, default=_marker):
@@ -2938,11 +4634,51 @@ def minmax(iterable_or_value, *others, key=None, default=_marker):
Raymond Hettinger and takes care to minimize the number of comparisons
performed.
"""
- pass
-
-
-def constrained_batches(iterable, max_size, max_count=None, get_len=len,
- strict=True):
+ iterable = (iterable_or_value, *others) if others else iterable_or_value
+
+ it = iter(iterable)
+
+ try:
+ lo = hi = next(it)
+ except StopIteration as exc:
+ if default is _marker:
+ raise ValueError(
+ '`minmax()` argument is an empty iterable. '
+ 'Provide a `default` value to suppress this error.'
+ ) from exc
+ return default
+
+ # Different branches depending on the presence of key. This saves a lot
+ # of unimportant copies which would slow the "key=None" branch
+ # significantly down.
+ if key is None:
+ for x, y in zip_longest(it, it, fillvalue=lo):
+ if y < x:
+ x, y = y, x
+ if x < lo:
+ lo = x
+ if hi < y:
+ hi = y
+
+ else:
+ lo_key = hi_key = key(lo)
+
+ for x, y in zip_longest(it, it, fillvalue=lo):
+ x_key, y_key = key(x), key(y)
+
+ if y_key < x_key:
+ x, y, x_key, y_key = y, x, y_key, x_key
+ if x_key < lo_key:
+ lo, lo_key = x, x_key
+ if hi_key < y_key:
+ hi, hi_key = y, y_key
+
+ return lo, hi
+
+
+def constrained_batches(
+ iterable, max_size, max_count=None, get_len=len, strict=True
+):
"""Yield batches of items from *iterable* with a combined size limited by
*max_size*.
@@ -2963,7 +4699,31 @@ def constrained_batches(iterable, max_size, max_count=None, get_len=len,
If *strict* is ``True``, raise ``ValueError`` if any single item is bigger
than *max_size*. Otherwise, allow single items to exceed *max_size*.
"""
- pass
+ if max_size <= 0:
+ raise ValueError('maximum size must be greater than zero')
+
+ batch = []
+ batch_size = 0
+ batch_count = 0
+ for item in iterable:
+ item_len = get_len(item)
+ if strict and item_len > max_size:
+ raise ValueError('item size exceeds maximum size')
+
+ reached_count = batch_count == max_count
+ reached_size = item_len + batch_size > max_size
+ if batch_count and (reached_size or reached_count):
+ yield tuple(batch)
+ batch.clear()
+ batch_size = 0
+ batch_count = 0
+
+ batch.append(item)
+ batch_size += item_len
+ batch_count += 1
+
+ if batch:
+ yield tuple(batch)
def gray_product(*iterables):
@@ -2982,7 +4742,30 @@ def gray_product(*iterables):
`this section <https://www-cs-faculty.stanford.edu/~knuth/fasc2a.ps.gz>`__
of Donald Knuth's *The Art of Computer Programming*.
"""
- pass
+ all_iterables = tuple(tuple(x) for x in iterables)
+ iterable_count = len(all_iterables)
+ for iterable in all_iterables:
+ if len(iterable) < 2:
+ raise ValueError("each iterable must have two or more items")
+
+ # This is based on "Algorithm H" from section 7.2.1.1, page 20.
+ # a holds the indexes of the source iterables for the n-tuple to be yielded
+ # f is the array of "focus pointers"
+ # o is the array of "directions"
+ a = [0] * iterable_count
+ f = list(range(iterable_count + 1))
+ o = [1] * iterable_count
+ while True:
+ yield tuple(all_iterables[i][a[i]] for i in range(iterable_count))
+ j = f[0]
+ f[0] = 0
+ if j == iterable_count:
+ break
+ a[j] = a[j] + o[j]
+ if a[j] == 0 or a[j] == len(all_iterables[j]) - 1:
+ o[j] = -o[j]
+ f[j] = f[j + 1]
+ f[j + 1] = j + 1
def partial_product(*iterables):
@@ -2996,7 +4779,18 @@ def partial_product(*iterables):
>>> list(partial_product('AB', 'C', 'DEF'))
[('A', 'C', 'D'), ('B', 'C', 'D'), ('B', 'C', 'E'), ('B', 'C', 'F')]
"""
- pass
+
+ iterators = list(map(iter, iterables))
+
+ try:
+ prod = [next(it) for it in iterators]
+ except StopIteration:
+ return
+ yield tuple(prod)
+
+ for i, it in enumerate(iterators):
+ for prod[i] in it:
+ yield tuple(prod)
def takewhile_inclusive(predicate, iterable):
@@ -3007,7 +4801,10 @@ def takewhile_inclusive(predicate, iterable):
:func:`takewhile` would return ``[1, 4]``.
"""
- pass
+ for x in iterable:
+ yield x
+ if not predicate(x):
+ break
def outer_product(func, xs, ys, *args, **kwargs):
@@ -3036,7 +4833,11 @@ def outer_product(func, xs, ys, *args, **kwargs):
>>> list(outer_product(min, animals, animals, key=len))
[('cat', 'cat', 'cat'), ('cat', 'wolf', 'wolf'), ('cat', 'wolf', 'mouse')]
"""
- pass
+ ys = tuple(ys)
+ return batched(
+ starmap(lambda x, y: func(x, y, *args, **kwargs), product(xs, ys)),
+ n=len(ys),
+ )
def iter_suppress(iterable, *exceptions):
@@ -3056,7 +4857,10 @@ def iter_suppress(iterable, *exceptions):
>>> list(chain(it_1, it_2))
[1, 2, 3, 4, 2, 3, 4]
"""
- pass
+ try:
+ yield from iterable
+ except exceptions:
+ return
def filter_map(func, iterable):
@@ -3067,7 +4871,10 @@ def filter_map(func, iterable):
>>> list(filter_map(lambda s: int(s) if s.isnumeric() else None, elems))
[1, 2, 3]
"""
- pass
+ for x in iterable:
+ y = func(x)
+ if y is not None:
+ yield y
def powerset_of_sets(iterable):
@@ -3081,7 +4888,9 @@ def powerset_of_sets(iterable):
:func:`powerset_of_sets` takes care to minimize the number
of hash operations performed.
"""
- pass
+ sets = tuple(map(set, dict.fromkeys(map(frozenset, zip(iterable)))))
+ for r in range(len(sets) + 1):
+ yield from starmap(set().union, combinations(sets, r))
def join_mappings(**field_to_map):
@@ -3093,14 +4902,25 @@ def join_mappings(**field_to_map):
>>> join_mappings(score=user_scores, time=user_times)
{'elliot': {'score': 50, 'time': 30}, 'claris': {'score': 60, 'time': 40}}
"""
- pass
+ ret = defaultdict(dict)
+
+ for field_name, mapping in field_to_map.items():
+ for key, value in mapping.items():
+ ret[key][field_name] = value
+
+ return dict(ret)
def _complex_sumprod(v1, v2):
"""High precision sumprod() for complex numbers.
Used by :func:`dft` and :func:`idft`.
"""
- pass
+
+ r1 = chain((p.real for p in v1), (-p.imag for p in v1))
+ r2 = chain((q.real for q in v2), (q.imag for q in v2))
+ i1 = chain((p.real for p in v1), (p.imag for p in v1))
+ i2 = chain((q.imag for q in v2), (q.real for q in v2))
+ return complex(_fsumprod(r1, r2), _fsumprod(i1, i2))
def dft(xarr):
@@ -3115,7 +4935,11 @@ def dft(xarr):
See :func:`idft` for the inverse Discrete Fourier Transform.
"""
- pass
+ N = len(xarr)
+ roots_of_unity = [e ** (n / N * tau * -1j) for n in range(N)]
+ for k in range(N):
+ coeffs = [roots_of_unity[k * n % N] for n in range(N)]
+ yield _complex_sumprod(xarr, coeffs)
def idft(Xarr):
@@ -3131,7 +4955,11 @@ def idft(Xarr):
See :func:`dft` for the Discrete Fourier Transform.
"""
- pass
+ N = len(Xarr)
+ roots_of_unity = [e ** (n / N * tau * 1j) for n in range(N)]
+ for k in range(N):
+ coeffs = [roots_of_unity[k * n % N] for n in range(N)]
+ yield _complex_sumprod(Xarr, coeffs) / N
def doublestarmap(func, iterable):
@@ -3148,4 +4976,5 @@ def doublestarmap(func, iterable):
``TypeError`` will be raised if *func*'s signature doesn't match the
mapping contained in *iterable* or if *iterable* does not contain mappings.
"""
- pass
+ for item in iterable:
+ yield func(**item)
diff --git a/more_itertools/recipes.py b/more_itertools/recipes.py
index 9c47a54..a21a1f5 100644
--- a/more_itertools/recipes.py
+++ b/more_itertools/recipes.py
@@ -7,32 +7,91 @@ Some backward-compatible usability improvements have been made.
.. [1] http://docs.python.org/library/itertools.html#recipes
"""
+
import math
import operator
+
from collections import deque
from collections.abc import Sized
from functools import partial, reduce
-from itertools import chain, combinations, compress, count, cycle, groupby, islice, product, repeat, starmap, tee, zip_longest
+from itertools import (
+ chain,
+ combinations,
+ compress,
+ count,
+ cycle,
+ groupby,
+ islice,
+ product,
+ repeat,
+ starmap,
+ tee,
+ zip_longest,
+)
from random import randrange, sample, choice
from sys import hexversion
-__all__ = ['all_equal', 'batched', 'before_and_after', 'consume',
- 'convolve', 'dotproduct', 'first_true', 'factor', 'flatten', 'grouper',
- 'iter_except', 'iter_index', 'matmul', 'ncycles', 'nth',
- 'nth_combination', 'padnone', 'pad_none', 'pairwise', 'partition',
- 'polynomial_eval', 'polynomial_from_roots', 'polynomial_derivative',
- 'powerset', 'prepend', 'quantify', 'reshape',
- 'random_combination_with_replacement', 'random_combination',
- 'random_permutation', 'random_product', 'repeatfunc', 'roundrobin',
- 'sieve', 'sliding_window', 'subslices', 'sum_of_squares', 'tabulate',
- 'tail', 'take', 'totient', 'transpose', 'triplewise', 'unique',
- 'unique_everseen', 'unique_justseen']
+
+__all__ = [
+ 'all_equal',
+ 'batched',
+ 'before_and_after',
+ 'consume',
+ 'convolve',
+ 'dotproduct',
+ 'first_true',
+ 'factor',
+ 'flatten',
+ 'grouper',
+ 'iter_except',
+ 'iter_index',
+ 'matmul',
+ 'ncycles',
+ 'nth',
+ 'nth_combination',
+ 'padnone',
+ 'pad_none',
+ 'pairwise',
+ 'partition',
+ 'polynomial_eval',
+ 'polynomial_from_roots',
+ 'polynomial_derivative',
+ 'powerset',
+ 'prepend',
+ 'quantify',
+ 'reshape',
+ 'random_combination_with_replacement',
+ 'random_combination',
+ 'random_permutation',
+ 'random_product',
+ 'repeatfunc',
+ 'roundrobin',
+ 'sieve',
+ 'sliding_window',
+ 'subslices',
+ 'sum_of_squares',
+ 'tabulate',
+ 'tail',
+ 'take',
+ 'totient',
+ 'transpose',
+ 'triplewise',
+ 'unique',
+ 'unique_everseen',
+ 'unique_justseen',
+]
+
_marker = object()
+
+
+# zip with strict is available for Python 3.10+
try:
zip(strict=True)
except TypeError:
_zip_strict = zip
else:
_zip_strict = partial(zip, strict=True)
+
+# math.sumprod is available for Python 3.12+
_sumprod = getattr(math, 'sumprod', lambda x, y: dotproduct(x, y))
@@ -49,7 +108,7 @@ def take(n, iterable):
[0, 1, 2]
"""
- pass
+ return list(islice(iterable, n))
def tabulate(function, start=0):
@@ -67,7 +126,7 @@ def tabulate(function, start=0):
[9, 4, 1, 0]
"""
- pass
+ return map(function, count(start))
def tail(n, iterable):
@@ -78,7 +137,14 @@ def tail(n, iterable):
['E', 'F', 'G']
"""
- pass
+ # If the given iterable has a length, then we can use islice to get its
+ # final elements. Note that if the iterable is not actually Iterable,
+ # either islice or deque will throw a TypeError. This is why we don't
+ # check if it is Iterable.
+ if isinstance(iterable, Sized):
+ yield from islice(iterable, max(0, len(iterable) - n), None)
+ else:
+ yield from iter(deque(iterable, maxlen=n))
def consume(iterator, n=None):
@@ -112,7 +178,13 @@ def consume(iterator, n=None):
StopIteration
"""
- pass
+ # Use functions that consume iterators at C speed.
+ if n is None:
+ # feed the entire iterator into a zero-length deque
+ deque(iterator, maxlen=0)
+ else:
+ # advance to the empty slice starting at position n
+ next(islice(iterator, n, n), None)
def nth(iterable, n, default=None):
@@ -125,7 +197,7 @@ def nth(iterable, n, default=None):
'zebra'
"""
- pass
+ return next(islice(iterable, n, None), default)
def all_equal(iterable, key=None):
@@ -146,7 +218,7 @@ def all_equal(iterable, key=None):
True
"""
- pass
+ return len(list(islice(groupby(iterable, key), 2))) <= 1
def quantify(iterable, pred=bool):
@@ -156,7 +228,7 @@ def quantify(iterable, pred=bool):
2
"""
- pass
+ return sum(map(pred, iterable))
def pad_none(iterable):
@@ -170,7 +242,7 @@ def pad_none(iterable):
See also :func:`padded`.
"""
- pass
+ return chain(iterable, repeat(None))
padnone = pad_none
@@ -183,7 +255,7 @@ def ncycles(iterable, n):
['a', 'b', 'a', 'b', 'a', 'b']
"""
- pass
+ return chain.from_iterable(repeat(tuple(iterable), n))
def dotproduct(vec1, vec2):
@@ -193,7 +265,7 @@ def dotproduct(vec1, vec2):
400
"""
- pass
+ return sum(map(operator.mul, vec1, vec2))
def flatten(listOfLists):
@@ -205,7 +277,7 @@ def flatten(listOfLists):
See also :func:`collapse`, which can flatten multiple levels of nesting.
"""
- pass
+ return chain.from_iterable(listOfLists)
def repeatfunc(func, times=None, *args):
@@ -230,7 +302,9 @@ def repeatfunc(func, times=None, *args):
[2, 4, 8, 1, 8, 4]
"""
- pass
+ if times is None:
+ return starmap(func, repeat(args))
+ return starmap(func, repeat(args, times))
def _pairwise(iterable):
@@ -242,7 +316,9 @@ def _pairwise(iterable):
On Python 3.10 and above, this is an alias for :func:`itertools.pairwise`.
"""
- pass
+ a, b = tee(iterable)
+ next(b, None)
+ return zip(a, b)
try:
@@ -250,19 +326,48 @@ try:
except ImportError:
pairwise = _pairwise
else:
+
+ def pairwise(iterable):
+ return itertools_pairwise(iterable)
+
pairwise.__doc__ = _pairwise.__doc__
class UnequalIterablesError(ValueError):
-
def __init__(self, details=None):
msg = 'Iterables have different lengths'
if details is not None:
- msg += ': index 0 has length {}; index {} has length {}'.format(*
- details)
+ msg += (': index 0 has length {}; index {} has length {}').format(
+ *details
+ )
+
super().__init__(msg)
+def _zip_equal_generator(iterables):
+ for combo in zip_longest(*iterables, fillvalue=_marker):
+ for val in combo:
+ if val is _marker:
+ raise UnequalIterablesError()
+ yield combo
+
+
+def _zip_equal(*iterables):
+ # Check whether the iterables are all the same size.
+ try:
+ first_size = len(iterables[0])
+ for i, it in enumerate(iterables[1:], 1):
+ size = len(it)
+ if size != first_size:
+ raise UnequalIterablesError(details=(first_size, i, size))
+ # All sizes are equal, we can use the built-in zip.
+ return zip(*iterables)
+ # If any one of the iterables didn't have a length, start reading
+ # them until one runs out.
+ except TypeError:
+ return _zip_equal_generator(iterables)
+
+
def grouper(iterable, n, incomplete='fill', fillvalue=None):
"""Group elements from *iterable* into fixed-length groups of length *n*.
@@ -292,7 +397,15 @@ def grouper(iterable, n, incomplete='fill', fillvalue=None):
UnequalIterablesError
"""
- pass
+ args = [iter(iterable)] * n
+ if incomplete == 'fill':
+ return zip_longest(*args, fillvalue=fillvalue)
+ if incomplete == 'strict':
+ return _zip_equal(*args)
+ if incomplete == 'ignore':
+ return zip(*args)
+ else:
+ raise ValueError('Expected fill, strict, or ignore')
def roundrobin(*iterables):
@@ -306,7 +419,11 @@ def roundrobin(*iterables):
iterables is small).
"""
- pass
+ # Algorithm credited to George Sakkis
+ iterators = map(iter, iterables)
+ for num_active in range(len(iterables), 0, -1):
+ iterators = cycle(islice(iterators, num_active))
+ yield from map(next, iterators)
def partition(pred, iterable):
@@ -329,7 +446,12 @@ def partition(pred, iterable):
([0, False, ''], [1, True, ' '])
"""
- pass
+ if pred is None:
+ pred = bool
+
+ t1, t2, p = tee(iterable, 3)
+ p1, p2 = tee(map(pred, p))
+ return (compress(t1, map(operator.not_, p1)), compress(t2, p2))
def powerset(iterable):
@@ -349,7 +471,8 @@ def powerset(iterable):
For a variant that efficiently yields actual :class:`set` instances, see
:func:`powerset_of_sets`.
"""
- pass
+ s = list(iterable)
+ return chain.from_iterable(combinations(s, r) for r in range(len(s) + 1))
def unique_everseen(iterable, key=None):
@@ -379,7 +502,22 @@ def unique_everseen(iterable, key=None):
``key=lambda x: frozenset(x.items())`` can be used.
"""
- pass
+ seenset = set()
+ seenset_add = seenset.add
+ seenlist = []
+ seenlist_add = seenlist.append
+ use_key = key is not None
+
+ for element in iterable:
+ k = key(element) if use_key else element
+ try:
+ if k not in seenset:
+ seenset_add(k)
+ yield element
+ except TypeError:
+ if k not in seenlist:
+ seenlist_add(k)
+ yield element
def unique_justseen(iterable, key=None):
@@ -391,7 +529,10 @@ def unique_justseen(iterable, key=None):
['A', 'B', 'C', 'A', 'D']
"""
- pass
+ if key is None:
+ return map(operator.itemgetter(0), groupby(iterable))
+
+ return map(next, map(operator.itemgetter(1), groupby(iterable, key)))
def unique(iterable, key=None, reverse=False):
@@ -410,7 +551,7 @@ def unique(iterable, key=None, reverse=False):
The elements in *iterable* need not be hashable, but they must be
comparable for sorting to work.
"""
- pass
+ return unique_justseen(sorted(iterable, key=key, reverse=reverse), key=key)
def iter_except(func, exception, first=None):
@@ -435,7 +576,13 @@ def iter_except(func, exception, first=None):
[]
"""
- pass
+ try:
+ if first is not None:
+ yield first()
+ while 1:
+ yield func()
+ except exception:
+ pass
def first_true(iterable, default=None, pred=None):
@@ -455,7 +602,7 @@ def first_true(iterable, default=None, pred=None):
'missing'
"""
- pass
+ return next(filter(pred, iterable), default)
def random_product(*args, repeat=1):
@@ -474,7 +621,8 @@ def random_product(*args, repeat=1):
``itertools.product(*args, **kwarg)``.
"""
- pass
+ pools = [tuple(pool) for pool in args] * repeat
+ return tuple(choice(pool) for pool in pools)
def random_permutation(iterable, r=None):
@@ -490,7 +638,9 @@ def random_permutation(iterable, r=None):
``itertools.permutations(iterable, r)``.
"""
- pass
+ pool = tuple(iterable)
+ r = len(pool) if r is None else r
+ return tuple(sample(pool, r))
def random_combination(iterable, r):
@@ -503,7 +653,10 @@ def random_combination(iterable, r):
``itertools.combinations(iterable, r)``.
"""
- pass
+ pool = tuple(iterable)
+ n = len(pool)
+ indices = sorted(sample(range(n), r))
+ return tuple(pool[i] for i in indices)
def random_combination_with_replacement(iterable, r):
@@ -517,7 +670,10 @@ def random_combination_with_replacement(iterable, r):
``itertools.combinations_with_replacement(iterable, r)``.
"""
- pass
+ pool = tuple(iterable)
+ n = len(pool)
+ indices = sorted(randrange(n) for i in range(r))
+ return tuple(pool[i] for i in indices)
def nth_combination(iterable, r, index):
@@ -535,7 +691,31 @@ def nth_combination(iterable, r, index):
of *iterable*.
``IndexError`` will be raised if the given *index* is invalid.
"""
- pass
+ pool = tuple(iterable)
+ n = len(pool)
+ if (r < 0) or (r > n):
+ raise ValueError
+
+ c = 1
+ k = min(r, n - r)
+ for i in range(1, k + 1):
+ c = c * (n - k + i) // i
+
+ if index < 0:
+ index += c
+
+ if (index < 0) or (index >= c):
+ raise IndexError
+
+ result = []
+ while r:
+ c, n, r = c * r // n, n - 1, r - 1
+ while index >= c:
+ index -= c
+ c, n = c * (n - r) // n, n - 1
+ result.append(pool[-1 - n])
+
+ return tuple(result)
def prepend(value, iterator):
@@ -550,7 +730,7 @@ def prepend(value, iterator):
or :func:`value_chain`.
"""
- pass
+ return chain([value], iterator)
def convolve(signal, kernel):
@@ -565,7 +745,14 @@ def convolve(signal, kernel):
is immediately consumed and stored.
"""
- pass
+ # This implementation intentionally doesn't match the one in the itertools
+ # documentation.
+ kernel = tuple(kernel)[::-1]
+ n = len(kernel)
+ window = deque([0], maxlen=n) * n
+ for x in chain(signal, repeat(0, n - 1)):
+ window.append(x)
+ yield _sumprod(kernel, window)
def before_and_after(predicate, it):
@@ -582,7 +769,23 @@ def before_and_after(predicate, it):
Note that the first iterator must be fully consumed before the second
iterator can generate valid results.
"""
- pass
+ it = iter(it)
+ transition = []
+
+ def true_iterator():
+ for elem in it:
+ if predicate(elem):
+ yield elem
+ else:
+ transition.append(elem)
+ return
+
+ # Note: this is different from itertools recipes to allow nesting
+ # before_and_after remainders into before_and_after again. See tests
+ # for an example.
+ remainder_iterator = chain(transition, it)
+
+ return true_iterator(), remainder_iterator
def triplewise(iterable):
@@ -592,7 +795,30 @@ def triplewise(iterable):
[('A', 'B', 'C'), ('B', 'C', 'D'), ('C', 'D', 'E')]
"""
- pass
+ # This deviates from the itertools documentation reciple - see
+ # https://github.com/more-itertools/more-itertools/issues/889
+ t1, t2, t3 = tee(iterable, 3)
+ next(t3, None)
+ next(t3, None)
+ next(t2, None)
+ return zip(t1, t2, t3)
+
+
+def _sliding_window_islice(iterable, n):
+ # Fast path for small, non-zero values of n.
+ iterators = tee(iterable, n)
+ for i, iterator in enumerate(iterators):
+ next(islice(iterator, i, i), None)
+ return zip(*iterators)
+
+
+def _sliding_window_deque(iterable, n):
+ # Normal path for other values of n.
+ it = iter(iterable)
+ window = deque(islice(it, n - 1), maxlen=n)
+ for x in it:
+ window.append(x)
+ yield tuple(window)
def sliding_window(iterable, n):
@@ -608,7 +834,16 @@ def sliding_window(iterable, n):
For a variant with more features, see :func:`windowed`.
"""
- pass
+ if n > 20:
+ return _sliding_window_deque(iterable, n)
+ elif n > 2:
+ return _sliding_window_islice(iterable, n)
+ elif n == 2:
+ return pairwise(iterable)
+ elif n == 1:
+ return zip(iterable)
+ else:
+ raise ValueError(f'n should be at least one, not {n}')
def subslices(iterable):
@@ -620,7 +855,9 @@ def subslices(iterable):
This is similar to :func:`substrings`, but emits items in a different
order.
"""
- pass
+ seq = list(iterable)
+ slices = starmap(slice, combinations(range(len(seq) + 1), 2))
+ return map(operator.getitem, repeat(seq), slices)
def polynomial_from_roots(roots):
@@ -630,7 +867,8 @@ def polynomial_from_roots(roots):
>>> polynomial_from_roots(roots) # x^3 - 4 * x^2 - 17 * x + 60
[1, -4, -17, 60]
"""
- pass
+ factors = zip(repeat(1), map(operator.neg, roots))
+ return list(reduce(convolve, factors, [1]))
def iter_index(iterable, value, start=0, stop=None):
@@ -658,7 +896,22 @@ def iter_index(iterable, value, start=0, stop=None):
associated with particular values.
"""
- pass
+ seq_index = getattr(iterable, 'index', None)
+ if seq_index is None:
+ # Slow path for general iterables
+ it = islice(iterable, start, stop)
+ for i, element in enumerate(it, start):
+ if element is value or element == value:
+ yield i
+ else:
+ # Fast path for sequences
+ stop = len(iterable) if stop is None else stop
+ i = start - 1
+ try:
+ while True:
+ yield (i := seq_index(value, i + 1, stop))
+ except ValueError:
+ pass
def sieve(n):
@@ -667,7 +920,16 @@ def sieve(n):
>>> list(sieve(30))
[2, 3, 5, 7, 11, 13, 17, 19, 23, 29]
"""
- pass
+ if n > 2:
+ yield 2
+ start = 3
+ data = bytearray((0, 1)) * (n // 2)
+ limit = math.isqrt(n) + 1
+ for p in iter_index(data, 1, start, limit):
+ yield from iter_index(data, 1, start, p * p)
+ data[p * p : n : p + p] = bytes(len(range(p * p, n, p + p)))
+ start = p * p
+ yield from iter_index(data, 1, start)
def _batched(iterable, n, *, strict=False):
@@ -681,13 +943,24 @@ def _batched(iterable, n, *, strict=False):
On Python 3.13 and above, this is an alias for :func:`itertools.batched`.
"""
- pass
+ if n < 1:
+ raise ValueError('n must be at least one')
+ it = iter(iterable)
+ while batch := tuple(islice(it, n)):
+ if strict and len(batch) != n:
+ raise ValueError('batched(): incomplete batch')
+ yield batch
-if hexversion >= 51183778:
+if hexversion >= 0x30D00A2:
from itertools import batched as itertools_batched
+
+ def batched(iterable, n, *, strict=False):
+ return itertools_batched(iterable, n, strict=strict)
+
else:
batched = _batched
+
batched.__doc__ = _batched.__doc__
@@ -700,7 +973,7 @@ def transpose(it):
The caller should ensure that the dimensions of the input are compatible.
If the input is empty, no output will be produced.
"""
- pass
+ return _zip_strict(*it)
def reshape(matrix, cols):
@@ -711,7 +984,7 @@ def reshape(matrix, cols):
>>> list(reshape(matrix, cols))
[(0, 1, 2), (3, 4, 5)]
"""
- pass
+ return batched(chain.from_iterable(matrix), cols)
def matmul(m1, m2):
@@ -723,7 +996,8 @@ def matmul(m1, m2):
The caller should ensure that the dimensions of the input matrices are
compatible with each other.
"""
- pass
+ n = len(m2[0])
+ return batched(starmap(_sumprod, product(m1, transpose(m2))), n)
def factor(n):
@@ -732,7 +1006,14 @@ def factor(n):
>>> list(factor(360))
[2, 2, 2, 3, 3, 5]
"""
- pass
+ for prime in sieve(math.isqrt(n) + 1):
+ while not n % prime:
+ yield prime
+ n //= prime
+ if n == 1:
+ return
+ if n > 1:
+ yield n
def polynomial_eval(coefficients, x):
@@ -745,7 +1026,11 @@ def polynomial_eval(coefficients, x):
>>> polynomial_eval(coefficients, x)
8.125
"""
- pass
+ n = len(coefficients)
+ if n == 0:
+ return x * 0 # coerce zero to the type of x
+ powers = map(pow, repeat(x), reversed(range(n)))
+ return _sumprod(coefficients, powers)
def sum_of_squares(it):
@@ -754,7 +1039,7 @@ def sum_of_squares(it):
>>> sum_of_squares([10, 20, 30])
1400
"""
- pass
+ return _sumprod(*tee(it))
def polynomial_derivative(coefficients):
@@ -767,7 +1052,9 @@ def polynomial_derivative(coefficients):
>>> derivative_coefficients
[3, -8, -17]
"""
- pass
+ n = len(coefficients)
+ powers = reversed(range(1, n))
+ return list(map(operator.mul, coefficients, powers))
def totient(n):
@@ -778,4 +1065,6 @@ def totient(n):
>>> totient(12)
4
"""
- pass
+ for prime in set(factor(n)):
+ n -= n // prime
+ return n