codekingpro/portable-devtools
114k
1import warnings2 3from collections import Counter, defaultdict, deque, abc4from collections.abc import Sequence5from functools import cached_property, partial, reduce, wraps6from heapq import heapify, heapreplace, heappop7from itertools import (8 chain,9 compress,10 count,11 cycle,12 dropwhile,13 groupby,14 islice,15 repeat,16 starmap,17 takewhile,18 tee,19 zip_longest,20 product,21)22from math import exp, factorial, floor, log, perm, comb23from queue import Empty, Queue24from random import random, randrange, uniform25from operator import itemgetter, mul, sub, gt, lt, ge, le26from sys import hexversion, maxsize27from time import monotonic28 29from .recipes import (30 _marker,31 _zip_equal,32 UnequalIterablesError,33 consume,34 flatten,35 pairwise,36 powerset,37 take,38 unique_everseen,39 all_equal,40 batched,41)42 43__all__ = [44 'AbortThread',45 'SequenceView',46 'UnequalIterablesError',47 'adjacent',48 'all_unique',49 'always_iterable',50 'always_reversible',51 'bucket',52 'callback_iter',53 'chunked',54 'chunked_even',55 'circular_shifts',56 'collapse',57 'combination_index',58 'combination_with_replacement_index',59 'consecutive_groups',60 'constrained_batches',61 'consumer',62 'count_cycle',63 'countable',64 'difference',65 'distinct_combinations',66 'distinct_permutations',67 'distribute',68 'divide',69 'duplicates_everseen',70 'duplicates_justseen',71 'classify_unique',72 'exactly_n',73 'filter_except',74 'filter_map',75 'first',76 'gray_product',77 'groupby_transform',78 'ichunked',79 'iequals',80 'ilen',81 'interleave',82 'interleave_evenly',83 'interleave_longest',84 'intersperse',85 'is_sorted',86 'islice_extended',87 'iterate',88 'iter_suppress',89 'last',90 'locate',91 'longest_common_prefix',92 'lstrip',93 'make_decorator',94 'map_except',95 'map_if',96 'map_reduce',97 'mark_ends',98 'minmax',99 'nth_or_last',100 'nth_permutation',101 'nth_product',102 'nth_combination_with_replacement',103 'numeric_range',104 'one',105 'only',106 'outer_product',107 'padded',108 'partial_product',109 'partitions',110 'peekable',111 'permutation_index',112 'product_index',113 'raise_',114 'repeat_each',115 'repeat_last',116 'replace',117 'rlocate',118 'rstrip',119 'run_length',120 'sample',121 'seekable',122 'set_partitions',123 'side_effect',124 'sliced',125 'sort_together',126 'split_after',127 'split_at',128 'split_before',129 'split_into',130 'split_when',131 'spy',132 'stagger',133 'strip',134 'strictly_n',135 'substrings',136 'substrings_indexes',137 'takewhile_inclusive',138 'time_limited',139 'unique_in_window',140 'unique_to_each',141 'unzip',142 'value_chain',143 'windowed',144 'windowed_complete',145 'with_iter',146 'zip_broadcast',147 'zip_equal',148 'zip_offset',149]150 151 152def chunked(iterable, n, strict=False):153 """Break *iterable* into lists of length *n*:154 155 >>> list(chunked([1, 2, 3, 4, 5, 6], 3))156 [[1, 2, 3], [4, 5, 6]]157 158 By the default, the last yielded list will have fewer than *n* elements159 if the length of *iterable* is not divisible by *n*:160 161 >>> list(chunked([1, 2, 3, 4, 5, 6, 7, 8], 3))162 [[1, 2, 3], [4, 5, 6], [7, 8]]163 164 To use a fill-in value instead, see the :func:`grouper` recipe.165 166 If the length of *iterable* is not divisible by *n* and *strict* is167 ``True``, then ``ValueError`` will be raised before the last168 list is yielded.169 170 """171 iterator = iter(partial(take, n, iter(iterable)), [])172 if strict:173 if n is None:174 raise ValueError('n must not be None when using strict mode.')175 176 def ret():177 for chunk in iterator:178 if len(chunk) != n:179 raise ValueError('iterable is not divisible by n.')180 yield chunk181 182 return iter(ret())183 else:184 return iterator185 186 187def first(iterable, default=_marker):188 """Return the first item of *iterable*, or *default* if *iterable* is189 empty.190 191 >>> first([0, 1, 2, 3])192 0193 >>> first([], 'some default')194 'some default'195 196 If *default* is not provided and there are no items in the iterable,197 raise ``ValueError``.198 199 :func:`first` is useful when you have a generator of expensive-to-retrieve200 values and want any arbitrary one. It is marginally shorter than201 ``next(iter(iterable), default)``.202 203 """204 for item in iterable:205 return item206 if default is _marker:207 raise ValueError(208 'first() was called on an empty iterable, and no '209 'default value was provided.'210 )211 return default212 213 214def last(iterable, default=_marker):215 """Return the last item of *iterable*, or *default* if *iterable* is216 empty.217 218 >>> last([0, 1, 2, 3])219 3220 >>> last([], 'some default')221 'some default'222 223 If *default* is not provided and there are no items in the iterable,224 raise ``ValueError``.225 """226 try:227 if isinstance(iterable, Sequence):228 return iterable[-1]229 # Work around https://bugs.python.org/issue38525230 elif hasattr(iterable, '__reversed__') and (hexversion != 0x030800F0):231 return next(reversed(iterable))232 else:233 return deque(iterable, maxlen=1)[-1]234 except (IndexError, TypeError, StopIteration):235 if default is _marker:236 raise ValueError(237 'last() was called on an empty iterable, and no default was '238 'provided.'239 )240 return default241 242 243def nth_or_last(iterable, n, default=_marker):244 """Return the nth or the last item of *iterable*,245 or *default* if *iterable* is empty.246 247 >>> nth_or_last([0, 1, 2, 3], 2)248 2249 >>> nth_or_last([0, 1], 2)250 1251 >>> nth_or_last([], 0, 'some default')252 'some default'253 254 If *default* is not provided and there are no items in the iterable,255 raise ``ValueError``.256 """257 return last(islice(iterable, n + 1), default=default)258 259 260class peekable:261 """Wrap an iterator to allow lookahead and prepending elements.262 263 Call :meth:`peek` on the result to get the value that will be returned264 by :func:`next`. This won't advance the iterator:265 266 >>> p = peekable(['a', 'b'])267 >>> p.peek()268 'a'269 >>> next(p)270 'a'271 272 Pass :meth:`peek` a default value to return that instead of raising273 ``StopIteration`` when the iterator is exhausted.274 275 >>> p = peekable([])276 >>> p.peek('hi')277 'hi'278 279 peekables also offer a :meth:`prepend` method, which "inserts" items280 at the head of the iterable:281 282 >>> p = peekable([1, 2, 3])283 >>> p.prepend(10, 11, 12)284 >>> next(p)285 10286 >>> p.peek()287 11288 >>> list(p)289 [11, 12, 1, 2, 3]290 291 peekables can be indexed. Index 0 is the item that will be returned by292 :func:`next`, index 1 is the item after that, and so on:293 The values up to the given index will be cached.294 295 >>> p = peekable(['a', 'b', 'c', 'd'])296 >>> p[0]297 'a'298 >>> p[1]299 'b'300 >>> next(p)301 'a'302 303 Negative indexes are supported, but be aware that they will cache the304 remaining items in the source iterator, which may require significant305 storage.306 307 To check whether a peekable is exhausted, check its truth value:308 309 >>> p = peekable(['a', 'b'])310 >>> if p: # peekable has items311 ... list(p)312 ['a', 'b']313 >>> if not p: # peekable is exhausted314 ... list(p)315 []316 317 """318 319 def __init__(self, iterable):320 self._it = iter(iterable)321 self._cache = deque()322 323 def __iter__(self):324 return self325 326 def __bool__(self):327 try:328 self.peek()329 except StopIteration:330 return False331 return True332 333 def peek(self, default=_marker):334 """Return the item that will be next returned from ``next()``.335 336 Return ``default`` if there are no items left. If ``default`` is not337 provided, raise ``StopIteration``.338 339 """340 if not self._cache:341 try:342 self._cache.append(next(self._it))343 except StopIteration:344 if default is _marker:345 raise346 return default347 return self._cache[0]348 349 def prepend(self, *items):350 """Stack up items to be the next ones returned from ``next()`` or351 ``self.peek()``. The items will be returned in352 first in, first out order::353 354 >>> p = peekable([1, 2, 3])355 >>> p.prepend(10, 11, 12)356 >>> next(p)357 10358 >>> list(p)359 [11, 12, 1, 2, 3]360 361 It is possible, by prepending items, to "resurrect" a peekable that362 previously raised ``StopIteration``.363 364 >>> p = peekable([])365 >>> next(p)366 Traceback (most recent call last):367 ...368 StopIteration369 >>> p.prepend(1)370 >>> next(p)371 1372 >>> next(p)373 Traceback (most recent call last):374 ...375 StopIteration376 377 """378 self._cache.extendleft(reversed(items))379 380 def __next__(self):381 if self._cache:382 return self._cache.popleft()383 384 return next(self._it)385 386 def _get_slice(self, index):387 # Normalize the slice's arguments388 step = 1 if (index.step is None) else index.step389 if step > 0:390 start = 0 if (index.start is None) else index.start391 stop = maxsize if (index.stop is None) else index.stop392 elif step < 0:393 start = -1 if (index.start is None) else index.start394 stop = (-maxsize - 1) if (index.stop is None) else index.stop395 else:396 raise ValueError('slice step cannot be zero')397 398 # If either the start or stop index is negative, we'll need to cache399 # the rest of the iterable in order to slice from the right side.400 if (start < 0) or (stop < 0):401 self._cache.extend(self._it)402 # Otherwise we'll need to find the rightmost index and cache to that403 # point.404 else:405 n = min(max(start, stop) + 1, maxsize)406 cache_len = len(self._cache)407 if n >= cache_len:408 self._cache.extend(islice(self._it, n - cache_len))409 410 return list(self._cache)[index]411 412 def __getitem__(self, index):413 if isinstance(index, slice):414 return self._get_slice(index)415 416 cache_len = len(self._cache)417 if index < 0:418 self._cache.extend(self._it)419 elif index >= cache_len:420 self._cache.extend(islice(self._it, index + 1 - cache_len))421 422 return self._cache[index]423 424 425def consumer(func):426 """Decorator that automatically advances a PEP-342-style "reverse iterator"427 to its first yield point so you don't have to call ``next()`` on it428 manually.429 430 >>> @consumer431 ... def tally():432 ... i = 0433 ... while True:434 ... print('Thing number %s is %s.' % (i, (yield)))435 ... i += 1436 ...437 >>> t = tally()438 >>> t.send('red')439 Thing number 0 is red.440 >>> t.send('fish')441 Thing number 1 is fish.442 443 Without the decorator, you would have to call ``next(t)`` before444 ``t.send()`` could be used.445 446 """447 448 @wraps(func)449 def wrapper(*args, **kwargs):450 gen = func(*args, **kwargs)451 next(gen)452 return gen453 454 return wrapper455 456 457def ilen(iterable):458 """Return the number of items in *iterable*.459 460 >>> ilen(x for x in range(1000000) if x % 3 == 0)461 333334462 463 This consumes the iterable, so handle with care.464 465 """466 # This approach was selected because benchmarks showed it's likely the467 # fastest of the known implementations at the time of writing.468 # See GitHub tracker: #236, #230.469 counter = count()470 deque(zip(iterable, counter), maxlen=0)471 return next(counter)472 473 474def iterate(func, start):475 """Return ``start``, ``func(start)``, ``func(func(start))``, ...476 477 >>> from itertools import islice478 >>> list(islice(iterate(lambda x: 2*x, 1), 10))479 [1, 2, 4, 8, 16, 32, 64, 128, 256, 512]480 481 """482 while True:483 yield start484 try:485 start = func(start)486 except StopIteration:487 break488 489 490def with_iter(context_manager):491 """Wrap an iterable in a ``with`` statement, so it closes once exhausted.492 493 For example, this will close the file when the iterator is exhausted::494 495 upper_lines = (line.upper() for line in with_iter(open('foo')))496 497 Any context manager which returns an iterable is a candidate for498 ``with_iter``.499 500 """501 with context_manager as iterable:502 yield from iterable503 504 505def one(iterable, too_short=None, too_long=None):506 """Return the first item from *iterable*, which is expected to contain only507 that item. Raise an exception if *iterable* is empty or has more than one508 item.509 510 :func:`one` is useful for ensuring that an iterable contains only one item.511 For example, it can be used to retrieve the result of a database query512 that is expected to return a single row.513 514 If *iterable* is empty, ``ValueError`` will be raised. You may specify a515 different exception with the *too_short* keyword:516 517 >>> it = []518 >>> one(it) # doctest: +IGNORE_EXCEPTION_DETAIL519 Traceback (most recent call last):520 ...521 ValueError: too many items in iterable (expected 1)'522 >>> too_short = IndexError('too few items')523 >>> one(it, too_short=too_short) # doctest: +IGNORE_EXCEPTION_DETAIL524 Traceback (most recent call last):525 ...526 IndexError: too few items527 528 Similarly, if *iterable* contains more than one item, ``ValueError`` will529 be raised. You may specify a different exception with the *too_long*530 keyword:531 532 >>> it = ['too', 'many']533 >>> one(it) # doctest: +IGNORE_EXCEPTION_DETAIL534 Traceback (most recent call last):535 ...536 ValueError: Expected exactly one item in iterable, but got 'too',537 'many', and perhaps more.538 >>> too_long = RuntimeError539 >>> one(it, too_long=too_long) # doctest: +IGNORE_EXCEPTION_DETAIL540 Traceback (most recent call last):541 ...542 RuntimeError543 544 Note that :func:`one` attempts to advance *iterable* twice to ensure there545 is only one item. See :func:`spy` or :func:`peekable` to check iterable546 contents less destructively.547 548 """549 it = iter(iterable)550 551 try:552 first_value = next(it)553 except StopIteration as e:554 raise (555 too_short or ValueError('too few items in iterable (expected 1)')556 ) from e557 558 try:559 second_value = next(it)560 except StopIteration:561 pass562 else:563 msg = (564 'Expected exactly one item in iterable, but got {!r}, {!r}, '565 'and perhaps more.'.format(first_value, second_value)566 )567 raise too_long or ValueError(msg)568 569 return first_value570 571 572def raise_(exception, *args):573 raise exception(*args)574 575 576def strictly_n(iterable, n, too_short=None, too_long=None):577 """Validate that *iterable* has exactly *n* items and return them if578 it does. If it has fewer than *n* items, call function *too_short*579 with those items. If it has more than *n* items, call function580 *too_long* with the first ``n + 1`` items.581 582 >>> iterable = ['a', 'b', 'c', 'd']583 >>> n = 4584 >>> list(strictly_n(iterable, n))585 ['a', 'b', 'c', 'd']586 587 Note that the returned iterable must be consumed in order for the check to588 be made.589 590 By default, *too_short* and *too_long* are functions that raise591 ``ValueError``.592 593 >>> list(strictly_n('ab', 3)) # doctest: +IGNORE_EXCEPTION_DETAIL594 Traceback (most recent call last):595 ...596 ValueError: too few items in iterable (got 2)597 598 >>> list(strictly_n('abc', 2)) # doctest: +IGNORE_EXCEPTION_DETAIL599 Traceback (most recent call last):600 ...601 ValueError: too many items in iterable (got at least 3)602 603 You can instead supply functions that do something else.604 *too_short* will be called with the number of items in *iterable*.605 *too_long* will be called with `n + 1`.606 607 >>> def too_short(item_count):608 ... raise RuntimeError609 >>> it = strictly_n('abcd', 6, too_short=too_short)610 >>> list(it) # doctest: +IGNORE_EXCEPTION_DETAIL611 Traceback (most recent call last):612 ...613 RuntimeError614 615 >>> def too_long(item_count):616 ... print('The boss is going to hear about this')617 >>> it = strictly_n('abcdef', 4, too_long=too_long)618 >>> list(it)619 The boss is going to hear about this620 ['a', 'b', 'c', 'd']621 622 """623 if too_short is None:624 too_short = lambda item_count: raise_(625 ValueError,626 'Too few items in iterable (got {})'.format(item_count),627 )628 629 if too_long is None:630 too_long = lambda item_count: raise_(631 ValueError,632 'Too many items in iterable (got at least {})'.format(item_count),633 )634 635 it = iter(iterable)636 for i in range(n):637 try:638 item = next(it)639 except StopIteration:640 too_short(i)641 return642 else:643 yield item644 645 try:646 next(it)647 except StopIteration:648 pass649 else:650 too_long(n + 1)651 652 653def distinct_permutations(iterable, r=None):654 """Yield successive distinct permutations of the elements in *iterable*.655 656 >>> sorted(distinct_permutations([1, 0, 1]))657 [(0, 1, 1), (1, 0, 1), (1, 1, 0)]658 659 Equivalent to ``set(permutations(iterable))``, except duplicates are not660 generated and thrown away. For larger input sequences this is much more661 efficient.662 663 Duplicate permutations arise when there are duplicated elements in the664 input iterable. The number of items returned is665 `n! / (x_1! * x_2! * ... * x_n!)`, where `n` is the total number of666 items input, and each `x_i` is the count of a distinct item in the input667 sequence.668 669 If *r* is given, only the *r*-length permutations are yielded.670 671 >>> sorted(distinct_permutations([1, 0, 1], r=2))672 [(0, 1), (1, 0), (1, 1)]673 >>> sorted(distinct_permutations(range(3), r=2))674 [(0, 1), (0, 2), (1, 0), (1, 2), (2, 0), (2, 1)]675 676 """677 678 # Algorithm: https://w.wiki/Qai679 def _full(A):680 while True:681 # Yield the permutation we have682 yield tuple(A)683 684 # Find the largest index i such that A[i] < A[i + 1]685 for i in range(size - 2, -1, -1):686 if A[i] < A[i + 1]:687 break688 # If no such index exists, this permutation is the last one689 else:690 return691 692 # Find the largest index j greater than j such that A[i] < A[j]693 for j in range(size - 1, i, -1):694 if A[i] < A[j]:695 break696 697 # Swap the value of A[i] with that of A[j], then reverse the698 # sequence from A[i + 1] to form the new permutation699 A[i], A[j] = A[j], A[i]700 A[i + 1 :] = A[: i - size : -1] # A[i + 1:][::-1]701 702 # Algorithm: modified from the above703 def _partial(A, r):704 # Split A into the first r items and the last r items705 head, tail = A[:r], A[r:]706 right_head_indexes = range(r - 1, -1, -1)707 left_tail_indexes = range(len(tail))708 709 while True:710 # Yield the permutation we have711 yield tuple(head)712 713 # Starting from the right, find the first index of the head with714 # value smaller than the maximum value of the tail - call it i.715 pivot = tail[-1]716 for i in right_head_indexes:717 if head[i] < pivot:718 break719 pivot = head[i]720 else:721 return722 723 # Starting from the left, find the first value of the tail724 # with a value greater than head[i] and swap.725 for j in left_tail_indexes:726 if tail[j] > head[i]:727 head[i], tail[j] = tail[j], head[i]728 break729 # If we didn't find one, start from the right and find the first730 # index of the head with a value greater than head[i] and swap.731 else:732 for j in right_head_indexes:733 if head[j] > head[i]:734 head[i], head[j] = head[j], head[i]735 break736 737 # Reverse head[i + 1:] and swap it with tail[:r - (i + 1)]738 tail += head[: i - r : -1] # head[i + 1:][::-1]739 i += 1740 head[i:], tail[:] = tail[: r - i], tail[r - i :]741 742 items = sorted(iterable)743 744 size = len(items)745 if r is None:746 r = size747 748 if 0 < r <= size:749 return _full(items) if (r == size) else _partial(items, r)750 751 return iter(() if r else ((),))752 753 754def intersperse(e, iterable, n=1):755 """Intersperse filler element *e* among the items in *iterable*, leaving756 *n* items between each filler element.757 758 >>> list(intersperse('!', [1, 2, 3, 4, 5]))759 [1, '!', 2, '!', 3, '!', 4, '!', 5]760 761 >>> list(intersperse(None, [1, 2, 3, 4, 5], n=2))762 [1, 2, None, 3, 4, None, 5]763 764 """765 if n == 0:766 raise ValueError('n must be > 0')767 elif n == 1:768 # interleave(repeat(e), iterable) -> e, x_0, e, x_1, e, x_2...769 # islice(..., 1, None) -> x_0, e, x_1, e, x_2...770 return islice(interleave(repeat(e), iterable), 1, None)771 else:772 # interleave(filler, chunks) -> [e], [x_0, x_1], [e], [x_2, x_3]...773 # islice(..., 1, None) -> [x_0, x_1], [e], [x_2, x_3]...774 # flatten(...) -> x_0, x_1, e, x_2, x_3...775 filler = repeat([e])776 chunks = chunked(iterable, n)777 return flatten(islice(interleave(filler, chunks), 1, None))778 779 780def unique_to_each(*iterables):781 """Return the elements from each of the input iterables that aren't in the782 other input iterables.783 784 For example, suppose you have a set of packages, each with a set of785 dependencies::786 787 {'pkg_1': {'A', 'B'}, 'pkg_2': {'B', 'C'}, 'pkg_3': {'B', 'D'}}788 789 If you remove one package, which dependencies can also be removed?790 791 If ``pkg_1`` is removed, then ``A`` is no longer necessary - it is not792 associated with ``pkg_2`` or ``pkg_3``. Similarly, ``C`` is only needed for793 ``pkg_2``, and ``D`` is only needed for ``pkg_3``::794 795 >>> unique_to_each({'A', 'B'}, {'B', 'C'}, {'B', 'D'})796 [['A'], ['C'], ['D']]797 798 If there are duplicates in one input iterable that aren't in the others799 they will be duplicated in the output. Input order is preserved::800 801 >>> unique_to_each("mississippi", "missouri")802 [['p', 'p'], ['o', 'u', 'r']]803 804 It is assumed that the elements of each iterable are hashable.805 806 """807 pool = [list(it) for it in iterables]808 counts = Counter(chain.from_iterable(map(set, pool)))809 uniques = {element for element in counts if counts[element] == 1}810 return [list(filter(uniques.__contains__, it)) for it in pool]811 812 813def windowed(seq, n, fillvalue=None, step=1):814 """Return a sliding window of width *n* over the given iterable.815 816 >>> all_windows = windowed([1, 2, 3, 4, 5], 3)817 >>> list(all_windows)818 [(1, 2, 3), (2, 3, 4), (3, 4, 5)]819 820 When the window is larger than the iterable, *fillvalue* is used in place821 of missing values:822 823 >>> list(windowed([1, 2, 3], 4))824 [(1, 2, 3, None)]825 826 Each window will advance in increments of *step*:827 828 >>> list(windowed([1, 2, 3, 4, 5, 6], 3, fillvalue='!', step=2))829 [(1, 2, 3), (3, 4, 5), (5, 6, '!')]830 831 To slide into the iterable's items, use :func:`chain` to add filler items832 to the left:833 834 >>> iterable = [1, 2, 3, 4]835 >>> n = 3836 >>> padding = [None] * (n - 1)837 >>> list(windowed(chain(padding, iterable), 3))838 [(None, None, 1), (None, 1, 2), (1, 2, 3), (2, 3, 4)]839 """840 if n < 0:841 raise ValueError('n must be >= 0')842 if n == 0:843 yield tuple()844 return845 if step < 1:846 raise ValueError('step must be >= 1')847 848 window = deque(maxlen=n)849 i = n850 for _ in map(window.append, seq):851 i -= 1852 if not i:853 i = step854 yield tuple(window)855 856 size = len(window)857 if size == 0:858 return859 elif size < n:860 yield tuple(chain(window, repeat(fillvalue, n - size)))861 elif 0 < i < min(step, n):862 window += (fillvalue,) * i863 yield tuple(window)864 865 866def substrings(iterable):867 """Yield all of the substrings of *iterable*.868 869 >>> [''.join(s) for s in substrings('more')]870 ['m', 'o', 'r', 'e', 'mo', 'or', 're', 'mor', 'ore', 'more']871 872 Note that non-string iterables can also be subdivided.873 874 >>> list(substrings([0, 1, 2]))875 [(0,), (1,), (2,), (0, 1), (1, 2), (0, 1, 2)]876 877 """878 # The length-1 substrings879 seq = []880 for item in iter(iterable):881 seq.append(item)882 yield (item,)883 seq = tuple(seq)884 item_count = len(seq)885 886 # And the rest887 for n in range(2, item_count + 1):888 for i in range(item_count - n + 1):889 yield seq[i : i + n]890 891 892def substrings_indexes(seq, reverse=False):893 """Yield all substrings and their positions in *seq*894 895 The items yielded will be a tuple of the form ``(substr, i, j)``, where896 ``substr == seq[i:j]``.897 898 This function only works for iterables that support slicing, such as899 ``str`` objects.900 901 >>> for item in substrings_indexes('more'):902 ... print(item)903 ('m', 0, 1)904 ('o', 1, 2)905 ('r', 2, 3)906 ('e', 3, 4)907 ('mo', 0, 2)908 ('or', 1, 3)909 ('re', 2, 4)910 ('mor', 0, 3)911 ('ore', 1, 4)912 ('more', 0, 4)913 914 Set *reverse* to ``True`` to yield the same items in the opposite order.915 916 917 """918 r = range(1, len(seq) + 1)919 if reverse:920 r = reversed(r)921 return (922 (seq[i : i + L], i, i + L) for L in r for i in range(len(seq) - L + 1)923 )924 925 926class bucket:927 """Wrap *iterable* and return an object that buckets the iterable into928 child iterables based on a *key* function.929 930 >>> iterable = ['a1', 'b1', 'c1', 'a2', 'b2', 'c2', 'b3']931 >>> s = bucket(iterable, key=lambda x: x[0]) # Bucket by 1st character932 >>> sorted(list(s)) # Get the keys933 ['a', 'b', 'c']934 >>> a_iterable = s['a']935 >>> next(a_iterable)936 'a1'937 >>> next(a_iterable)938 'a2'939 >>> list(s['b'])940 ['b1', 'b2', 'b3']941 942 The original iterable will be advanced and its items will be cached until943 they are used by the child iterables. This may require significant storage.944 945 By default, attempting to select a bucket to which no items belong will946 exhaust the iterable and cache all values.947 If you specify a *validator* function, selected buckets will instead be948 checked against it.949 950 >>> from itertools import count951 >>> it = count(1, 2) # Infinite sequence of odd numbers952 >>> key = lambda x: x % 10 # Bucket by last digit953 >>> validator = lambda x: x in {1, 3, 5, 7, 9} # Odd digits only954 >>> s = bucket(it, key=key, validator=validator)955 >>> 2 in s956 False957 >>> list(s[2])958 []959 960 """961 962 def __init__(self, iterable, key, validator=None):963 self._it = iter(iterable)964 self._key = key965 self._cache = defaultdict(deque)966 self._validator = validator or (lambda x: True)967 968 def __contains__(self, value):969 if not self._validator(value):970 return False971 972 try:973 item = next(self[value])974 except StopIteration:975 return False976 else:977 self._cache[value].appendleft(item)978 979 return True980 981 def _get_values(self, value):982 """983 Helper to yield items from the parent iterator that match *value*.984 Items that don't match are stored in the local cache as they985 are encountered.986 """987 while True:988 # If we've cached some items that match the target value, emit989 # the first one and evict it from the cache.990 if self._cache[value]:991 yield self._cache[value].popleft()992 # Otherwise we need to advance the parent iterator to search for993 # a matching item, caching the rest.994 else:995 while True:996 try:997 item = next(self._it)998 except StopIteration:999 return1000 item_value = self._key(item)1001 if item_value == value:1002 yield item1003 break1004 elif self._validator(item_value):1005 self._cache[item_value].append(item)1006 1007 def __iter__(self):1008 for item in self._it:1009 item_value = self._key(item)1010 if self._validator(item_value):1011 self._cache[item_value].append(item)1012 1013 yield from self._cache.keys()1014 1015 def __getitem__(self, value):1016 if not self._validator(value):1017 return iter(())1018 1019 return self._get_values(value)1020 1021 1022def spy(iterable, n=1):1023 """Return a 2-tuple with a list containing the first *n* elements of1024 *iterable*, and an iterator with the same items as *iterable*.1025 This allows you to "look ahead" at the items in the iterable without1026 advancing it.1027 1028 There is one item in the list by default:1029 1030 >>> iterable = 'abcdefg'1031 >>> head, iterable = spy(iterable)1032 >>> head1033 ['a']1034 >>> list(iterable)1035 ['a', 'b', 'c', 'd', 'e', 'f', 'g']1036 1037 You may use unpacking to retrieve items instead of lists:1038 1039 >>> (head,), iterable = spy('abcdefg')1040 >>> head1041 'a'1042 >>> (first, second), iterable = spy('abcdefg', 2)1043 >>> first1044 'a'1045 >>> second1046 'b'1047 1048 The number of items requested can be larger than the number of items in1049 the iterable:1050 1051 >>> iterable = [1, 2, 3, 4, 5]1052 >>> head, iterable = spy(iterable, 10)1053 >>> head1054 [1, 2, 3, 4, 5]1055 >>> list(iterable)1056 [1, 2, 3, 4, 5]1057 1058 """1059 it = iter(iterable)1060 head = take(n, it)1061 1062 return head.copy(), chain(head, it)1063 1064 1065def interleave(*iterables):1066 """Return a new iterable yielding from each iterable in turn,1067 until the shortest is exhausted.1068 1069 >>> list(interleave([1, 2, 3], [4, 5], [6, 7, 8]))1070 [1, 4, 6, 2, 5, 7]1071 1072 For a version that doesn't terminate after the shortest iterable is1073 exhausted, see :func:`interleave_longest`.1074 1075 """1076 return chain.from_iterable(zip(*iterables))1077 1078 1079def interleave_longest(*iterables):1080 """Return a new iterable yielding from each iterable in turn,1081 skipping any that are exhausted.1082 1083 >>> list(interleave_longest([1, 2, 3], [4, 5], [6, 7, 8]))1084 [1, 4, 6, 2, 5, 7, 3, 8]1085 1086 This function produces the same output as :func:`roundrobin`, but may1087 perform better for some inputs (in particular when the number of iterables1088 is large).1089 1090 """1091 i = chain.from_iterable(zip_longest(*iterables, fillvalue=_marker))1092 return (x for x in i if x is not _marker)1093 1094 1095def interleave_evenly(iterables, lengths=None):1096 """1097 Interleave multiple iterables so that their elements are evenly distributed1098 throughout the output sequence.1099 1100 >>> iterables = [1, 2, 3, 4, 5], ['a', 'b']1101 >>> list(interleave_evenly(iterables))1102 [1, 2, 'a', 3, 4, 'b', 5]1103 1104 >>> iterables = [[1, 2, 3], [4, 5], [6, 7, 8]]1105 >>> list(interleave_evenly(iterables))1106 [1, 6, 4, 2, 7, 3, 8, 5]1107 1108 This function requires iterables of known length. Iterables without1109 ``__len__()`` can be used by manually specifying lengths with *lengths*:1110 1111 >>> from itertools import combinations, repeat1112 >>> iterables = [combinations(range(4), 2), ['a', 'b', 'c']]1113 >>> lengths = [4 * (4 - 1) // 2, 3]1114 >>> list(interleave_evenly(iterables, lengths=lengths))1115 [(0, 1), (0, 2), 'a', (0, 3), (1, 2), 'b', (1, 3), (2, 3), 'c']1116 1117 Based on Bresenham's algorithm.1118 """1119 if lengths is None:1120 try:1121 lengths = [len(it) for it in iterables]1122 except TypeError:1123 raise ValueError(1124 'Iterable lengths could not be determined automatically. '1125 'Specify them with the lengths keyword.'1126 )1127 elif len(iterables) != len(lengths):1128 raise ValueError('Mismatching number of iterables and lengths.')1129 1130 dims = len(lengths)1131 1132 # sort iterables by length, descending1133 lengths_permute = sorted(1134 range(dims), key=lambda i: lengths[i], reverse=True1135 )1136 lengths_desc = [lengths[i] for i in lengths_permute]1137 iters_desc = [iter(iterables[i]) for i in lengths_permute]1138 1139 # the longest iterable is the primary one (Bresenham: the longest1140 # distance along an axis)1141 delta_primary, deltas_secondary = lengths_desc[0], lengths_desc[1:]1142 iter_primary, iters_secondary = iters_desc[0], iters_desc[1:]1143 errors = [delta_primary // dims] * len(deltas_secondary)1144 1145 to_yield = sum(lengths)1146 while to_yield:1147 yield next(iter_primary)1148 to_yield -= 11149 # update errors for each secondary iterable1150 errors = [e - delta for e, delta in zip(errors, deltas_secondary)]1151 1152 # those iterables for which the error is negative are yielded1153 # ("diagonal step" in Bresenham)1154 for i, e in enumerate(errors):1155 if e < 0:1156 yield next(iters_secondary[i])1157 to_yield -= 11158 errors[i] += delta_primary1159 1160 1161def collapse(iterable, base_type=None, levels=None):1162 """Flatten an iterable with multiple levels of nesting (e.g., a list of1163 lists of tuples) into non-iterable types.1164 1165 >>> iterable = [(1, 2), ([3, 4], [[5], [6]])]1166 >>> list(collapse(iterable))1167 [1, 2, 3, 4, 5, 6]1168 1169 Binary and text strings are not considered iterable and1170 will not be collapsed.1171 1172 To avoid collapsing other types, specify *base_type*:1173 1174 >>> iterable = ['ab', ('cd', 'ef'), ['gh', 'ij']]1175 >>> list(collapse(iterable, base_type=tuple))1176 ['ab', ('cd', 'ef'), 'gh', 'ij']1177 1178 Specify *levels* to stop flattening after a certain level:1179 1180 >>> iterable = [('a', ['b']), ('c', ['d'])]1181 >>> list(collapse(iterable)) # Fully flattened1182 ['a', 'b', 'c', 'd']1183 >>> list(collapse(iterable, levels=1)) # Only one level flattened1184 ['a', ['b'], 'c', ['d']]1185 1186 """1187 1188 def walk(node, level):1189 if (1190 ((levels is not None) and (level > levels))1191 or isinstance(node, (str, bytes))1192 or ((base_type is not None) and isinstance(node, base_type))1193 ):1194 yield node1195 return1196 1197 try:1198 tree = iter(node)1199 except TypeError:1200 yield node