Coverage for /pythoncovmergedfiles/medio/medio/usr/local/lib/python3.11/site-packages/more_itertools/more.py: 19%
Shortcuts on this page
r m x toggle line displays
j k next/prev highlighted chunk
0 (zero) top of page
1 (one) first highlighted chunk
1__lazy_modules__ = frozenset({'queue', 'threading'})
3import math
4import types
6from collections import Counter, defaultdict, deque
7from collections.abc import Sequence
8from contextlib import suppress
9from functools import cached_property, partial, wraps
10from heapq import heapify, heapreplace
11from itertools import (
12 chain,
13 combinations,
14 compress,
15 count,
16 cycle,
17 dropwhile,
18 groupby,
19 islice,
20 permutations,
21 repeat,
22 starmap,
23 takewhile,
24 tee,
25 zip_longest,
26 product,
27)
28from math import comb, e, exp, floor, fsum, log, log1p, perm, tau
29from math import ceil, prod
30from queue import Empty, Queue
31from random import random, randrange, shuffle, uniform
32from operator import (
33 attrgetter,
34 getitem,
35 is_not,
36 itemgetter,
37 lt,
38 neg,
39 sub,
40 gt,
41)
42from sys import maxsize
43from time import monotonic
44from threading import Lock
46from .recipes import (
47 _marker,
48 consume,
49 first_true,
50 flatten,
51 is_prime,
52 nth,
53 powerset,
54 sieve,
55 take,
56 unique_everseen,
57 all_equal,
58 batched,
59)
61__all__ = [
62 'AbortThread',
63 'SequenceView',
64 'adjacent',
65 'all_unique',
66 'always_iterable',
67 'always_reversible',
68 'argmax',
69 'argmin',
70 'bucket',
71 'callback_iter',
72 'chunked',
73 'chunked_even',
74 'circular_shifts',
75 'classify_unique',
76 'collapse',
77 'combination_index',
78 'combination_with_replacement_index',
79 'concurrent_tee',
80 'consecutive_groups',
81 'constrained_batches',
82 'consumer',
83 'count_cycle',
84 'countable',
85 'derangements',
86 'dft',
87 'difference',
88 'distinct_combinations',
89 'distinct_permutations',
90 'distribute',
91 'divide',
92 'doublestarmap',
93 'duplicates_everseen',
94 'duplicates_justseen',
95 'exactly_n',
96 'extract',
97 'filter_except',
98 'filter_map',
99 'first',
100 'gray_product',
101 'groupby_transform',
102 'ichunked',
103 'idft',
104 'iequals',
105 'ilen',
106 'interleave',
107 'interleave_evenly',
108 'interleave_longest',
109 'interleave_randomly',
110 'intersperse',
111 'is_sorted',
112 'islice_extended',
113 'iter_suppress',
114 'iterate',
115 'join_mappings',
116 'last',
117 'locate',
118 'longest_common_prefix',
119 'lstrip',
120 'make_decorator',
121 'map_except',
122 'map_if',
123 'map_reduce',
124 'mark_ends',
125 'minmax',
126 'nth_combination_with_replacement',
127 'nth_or_last',
128 'nth_permutation',
129 'nth_prime',
130 'nth_product',
131 'numeric_range',
132 'one',
133 'only',
134 'outer_product',
135 'padded',
136 'partial_product',
137 'partitions',
138 'peekable',
139 'permutation_index',
140 'powerset_of_sets',
141 'product_index',
142 'raise_',
143 'repeat_each',
144 'repeat_last',
145 'replace',
146 'rlocate',
147 'rstrip',
148 'run_length',
149 'sample',
150 'seekable',
151 'serialize',
152 'set_partitions',
153 'side_effect',
154 'sized_iterator',
155 'sliced',
156 'sort_together',
157 'split_after',
158 'split_at',
159 'split_before',
160 'split_into',
161 'split_when',
162 'spy',
163 'stagger',
164 'strictly_n',
165 'strip',
166 'subfactorial',
167 'substrings',
168 'substrings_indexes',
169 'synchronized',
170 'takewhile_inclusive',
171 'time_limited',
172 'unique_in_window',
173 'unique_to_each',
174 'unzip',
175 'value_chain',
176 'windowed',
177 'windowed_complete',
178 'with_iter',
179 'zip_broadcast',
180 'zip_offset',
181]
183# math.sumprod is available for Python 3.12+
184try:
185 from math import sumprod as _fsumprod
187except ImportError: # pragma: no cover
188 # Extended precision algorithms from T. J. Dekker,
189 # "A Floating-Point Technique for Extending the Available Precision"
190 # https://csclub.uwaterloo.ca/~pbarfuss/dekker1971.pdf
191 # Formulas: (5.5) (5.6) and (5.8). Code: mul12()
193 def dl_split(x: float):
194 "Split a float into two half-precision components."
195 t = x * 134217729.0 # Veltkamp constant = 2.0 ** 27 + 1
196 hi = t - (t - x)
197 lo = x - hi
198 return hi, lo
200 def dl_mul(x, y):
201 "Lossless multiplication."
202 xx_hi, xx_lo = dl_split(x)
203 yy_hi, yy_lo = dl_split(y)
204 p = xx_hi * yy_hi
205 q = xx_hi * yy_lo + xx_lo * yy_hi
206 z = p + q
207 zz = p - z + q + xx_lo * yy_lo
208 return z, zz
210 def _fsumprod(p, q):
211 return fsum(chain.from_iterable(map(dl_mul, p, q)))
214def chunked(iterable, n, strict=False):
215 """Break *iterable* into lists of length *n*:
217 >>> list(chunked([1, 2, 3, 4, 5, 6], 3))
218 [[1, 2, 3], [4, 5, 6]]
220 By the default, the last yielded list will have fewer than *n* elements
221 if the length of *iterable* is not divisible by *n*:
223 >>> list(chunked([1, 2, 3, 4, 5, 6, 7, 8], 3))
224 [[1, 2, 3], [4, 5, 6], [7, 8]]
226 To use a fill-in value instead, see the :func:`grouper` recipe.
228 If the length of *iterable* is not divisible by *n* and *strict* is
229 ``True``, then ``ValueError`` will be raised before the last
230 list is yielded.
232 """
233 if n is not None and n < 0:
234 raise ValueError('n must be at least 0')
236 iterator = iter(partial(take, n, iter(iterable)), [])
237 if strict:
238 if n is None:
239 raise ValueError('n must not be None when using strict mode.')
241 def ret():
242 for chunk in iterator:
243 if len(chunk) != n:
244 raise ValueError('iterable is not divisible by n.')
245 yield chunk
247 return ret()
248 else:
249 return iterator
252def first(iterable, default=_marker):
253 """Return the first item of *iterable*, or *default* if *iterable* is
254 empty.
256 >>> first([0, 1, 2, 3])
257 0
258 >>> first([], 'some default')
259 'some default'
261 If *default* is not provided and there are no items in the iterable,
262 raise ``ValueError``.
264 :func:`first` is useful when you have a generator of expensive-to-retrieve
265 values and want any arbitrary one. It is marginally shorter than
266 ``next(iter(iterable), default)``.
268 """
269 for item in iterable:
270 return item
271 if default is _marker:
272 raise ValueError(
273 'first() was called on an empty iterable, '
274 'and no default value was provided.'
275 )
276 return default
279def last(iterable, default=_marker):
280 """Return the last item of *iterable*, or *default* if *iterable* is
281 empty.
283 >>> last([0, 1, 2, 3])
284 3
285 >>> last([], 'some default')
286 'some default'
288 If *default* is not provided and there are no items in the iterable,
289 raise ``ValueError``.
290 """
291 try:
292 if getattr(iterable, '__reversed__', None):
293 return next(reversed(iterable))
294 return deque(iterable, maxlen=1)[-1]
295 except (IndexError, StopIteration):
296 if default is _marker:
297 raise ValueError(
298 'last() was called on an empty iterable, '
299 'and no default value was provided.'
300 )
301 return default
304def nth_or_last(iterable, n, default=_marker):
305 """Return the nth or the last item of *iterable*,
306 or *default* if *iterable* is empty.
308 >>> nth_or_last([0, 1, 2, 3], 2)
309 2
310 >>> nth_or_last([0, 1], 2)
311 1
312 >>> nth_or_last([], 0, 'some default')
313 'some default'
315 If *default* is not provided and there are no items in the iterable,
316 raise ``ValueError``.
317 """
318 return last(islice(iterable, n + 1), default=default)
321class peekable:
322 """Wrap an iterator to allow lookahead and prepending elements.
324 Call :meth:`peek` on the result to get the value that will be returned
325 by :func:`next`. This won't advance the iterator:
327 >>> p = peekable(['a', 'b'])
328 >>> p.peek()
329 'a'
330 >>> next(p)
331 'a'
333 Pass :meth:`peek` a default value to return that instead of raising
334 ``StopIteration`` when the iterator is exhausted.
336 >>> p = peekable([])
337 >>> p.peek('hi')
338 'hi'
340 peekables also offer a :meth:`prepend` method, which "inserts" items
341 at the head of the iterable:
343 >>> p = peekable([1, 2, 3])
344 >>> p.prepend(10, 11, 12)
345 >>> next(p)
346 10
347 >>> p.peek()
348 11
349 >>> list(p)
350 [11, 12, 1, 2, 3]
352 peekables can be indexed. Index 0 is the item that will be returned by
353 :func:`next`, index 1 is the item after that, and so on:
354 The values up to the given index will be cached.
356 >>> p = peekable(['a', 'b', 'c', 'd'])
357 >>> p[0]
358 'a'
359 >>> p[1]
360 'b'
361 >>> next(p)
362 'a'
364 Negative indexes are supported, but be aware that they will cache the
365 remaining items in the source iterator, which may require significant
366 storage.
368 To check whether a peekable is exhausted, check its truth value:
370 >>> p = peekable(['a', 'b'])
371 >>> if p: # peekable has items
372 ... list(p)
373 ['a', 'b']
374 >>> if not p: # peekable is exhausted
375 ... list(p)
376 []
378 """
380 def __init__(self, iterable):
381 self._it = iter(iterable)
382 self._cache = deque()
384 def __iter__(self):
385 return self
387 def __bool__(self):
388 try:
389 self.peek()
390 except StopIteration:
391 return False
392 return True
394 def peek(self, default=_marker):
395 """Return the item that will be next returned from ``next()``.
397 Return ``default`` if there are no items left. If ``default`` is not
398 provided, raise ``StopIteration``.
400 """
401 if not self._cache:
402 try:
403 self._cache.append(next(self._it))
404 except StopIteration:
405 if default is _marker:
406 raise
407 return default
408 return self._cache[0]
410 def prepend(self, *items):
411 """Stack up items to be the next ones returned from ``next()`` or
412 ``self.peek()``. The items will be returned in
413 first in, first out order::
415 >>> p = peekable([1, 2, 3])
416 >>> p.prepend(10, 11, 12)
417 >>> next(p)
418 10
419 >>> list(p)
420 [11, 12, 1, 2, 3]
422 It is possible, by prepending items, to "resurrect" a peekable that
423 previously raised ``StopIteration``.
425 >>> p = peekable([])
426 >>> next(p)
427 Traceback (most recent call last):
428 ...
429 StopIteration
430 >>> p.prepend(1)
431 >>> next(p)
432 1
433 >>> next(p)
434 Traceback (most recent call last):
435 ...
436 StopIteration
438 """
439 self._cache.extendleft(reversed(items))
441 __class_getitem__ = classmethod(types.GenericAlias)
443 def __next__(self):
444 if self._cache:
445 return self._cache.popleft()
447 return next(self._it)
449 def _get_slice(self, index):
450 # Normalize the slice's arguments
451 step = 1 if (index.step is None) else index.step
452 if step > 0:
453 start = 0 if (index.start is None) else index.start
454 stop = maxsize if (index.stop is None) else index.stop
455 elif step < 0:
456 start = -1 if (index.start is None) else index.start
457 stop = (-maxsize - 1) if (index.stop is None) else index.stop
458 else:
459 raise ValueError('slice step cannot be zero')
461 # If either the start or stop index is negative, we'll need to cache
462 # the rest of the iterable in order to slice from the right side.
463 if (start < 0) or (stop < 0):
464 self._cache.extend(self._it)
465 # Otherwise we'll need to find the rightmost index and cache to that
466 # point.
467 else:
468 n = min(max(start, stop) + 1, maxsize)
469 cache_len = len(self._cache)
470 if n >= cache_len:
471 self._cache.extend(islice(self._it, n - cache_len))
473 return list(self._cache)[index]
475 def __getitem__(self, index):
476 if isinstance(index, slice):
477 return self._get_slice(index)
479 cache_len = len(self._cache)
480 if index < 0:
481 self._cache.extend(self._it)
482 elif index >= cache_len:
483 self._cache.extend(islice(self._it, index + 1 - cache_len))
485 return self._cache[index]
488def consumer(func):
489 """Decorator that automatically advances a PEP-342-style "reverse iterator"
490 to its first yield point so you don't have to call ``next()`` on it
491 manually.
493 >>> @consumer
494 ... def tally():
495 ... i = 0
496 ... while True:
497 ... print('Thing number %s is %s.' % (i, (yield)))
498 ... i += 1
499 ...
500 >>> t = tally()
501 >>> t.send('red')
502 Thing number 0 is red.
503 >>> t.send('fish')
504 Thing number 1 is fish.
506 Without the decorator, you would have to call ``next(t)`` before
507 ``t.send()`` could be used.
509 """
511 @wraps(func)
512 def wrapper(*args, **kwargs):
513 gen = func(*args, **kwargs)
514 next(gen)
515 return gen
517 return wrapper
520def ilen(iterable):
521 """Return the number of items in *iterable*.
523 For example, there are 168 prime numbers below 1,000:
525 >>> ilen(sieve(1000))
526 168
528 Equivalent to, but faster than::
530 def ilen(iterable):
531 count = 0
532 for _ in iterable:
533 count += 1
534 return count
536 This fully consumes the iterable, so handle with care.
538 """
539 # This is the "most beautiful of the fast variants" of this function.
540 # If you think you can improve on it, please ensure that your version
541 # is both 10x faster and 10x more beautiful.
542 return sum(compress(repeat(1), zip(iterable)))
545def iterate(func, start):
546 """Return ``start``, ``func(start)``, ``func(func(start))``, ...
548 Produces an infinite iterator. To add a stopping condition,
549 use :func:`take`, ``takewhile``, or :func:`takewhile_inclusive`:.
551 >>> take(10, iterate(lambda x: 2*x, 1))
552 [1, 2, 4, 8, 16, 32, 64, 128, 256, 512]
554 >>> collatz = lambda x: 3*x + 1 if x%2==1 else x // 2
555 >>> list(takewhile_inclusive(lambda x: x!=1, iterate(collatz, 10)))
556 [10, 5, 16, 8, 4, 2, 1]
558 """
559 with suppress(StopIteration):
560 while True:
561 yield start
562 start = func(start)
565def with_iter(context_manager):
566 """Wrap an iterable in a ``with`` statement, so it closes once exhausted.
568 For example, this will close the file when the iterator is exhausted::
570 upper_lines = (line.upper() for line in with_iter(open('foo')))
572 Note that you have to actually exhaust the iterator for opened files to be closed.
574 Any context manager which returns an iterable is a candidate for
575 ``with_iter``.
577 """
578 with context_manager as iterable:
579 yield from iterable
582class sized_iterator:
583 """Wrapper for *iterable* that implements ``__len__``.
585 >>> it = map(str, range(5))
586 >>> sized_it = sized_iterator(it, 5)
587 >>> len(sized_it)
588 5
589 >>> list(sized_it)
590 ['0', '1', '2', '3', '4']
592 This is useful for tools that use :func:`len`, like
593 `tqdm <https://pypi.org/project/tqdm/>`__ .
595 The wrapper doesn't validate the provided *length*, so be sure to choose
596 a value that reflects reality.
597 """
599 def __init__(self, iterable, length):
600 self._iterator = iter(iterable)
601 self._length = length
603 def __next__(self):
604 return next(self._iterator)
606 def __iter__(self):
607 return self
609 def __len__(self):
610 return self._length
613def one(iterable, too_short=None, too_long=None):
614 """Return the first item from *iterable*, which is expected to contain only
615 that item. Raise an exception if *iterable* is empty or has more than one
616 item.
618 :func:`one` is useful for ensuring that an iterable contains only one item.
619 For example, it can be used to retrieve the result of a database query
620 that is expected to return a single row.
622 If *iterable* is empty, ``ValueError`` will be raised. You may specify a
623 different exception with the *too_short* keyword:
625 >>> it = []
626 >>> one(it) # doctest: +IGNORE_EXCEPTION_DETAIL
627 Traceback (most recent call last):
628 ...
629 ValueError: too few items in iterable (expected 1)'
630 >>> too_short = IndexError('too few items')
631 >>> one(it, too_short=too_short) # doctest: +IGNORE_EXCEPTION_DETAIL
632 Traceback (most recent call last):
633 ...
634 IndexError: too few items
636 Similarly, if *iterable* contains more than one item, ``ValueError`` will
637 be raised. You may specify a different exception with the *too_long*
638 keyword:
640 >>> it = ['too', 'many']
641 >>> one(it) # doctest: +IGNORE_EXCEPTION_DETAIL
642 Traceback (most recent call last):
643 ...
644 ValueError: Expected exactly one item in iterable, but got 'too',
645 'many', and perhaps more.
646 >>> too_long = RuntimeError
647 >>> one(it, too_long=too_long) # doctest: +IGNORE_EXCEPTION_DETAIL
648 Traceback (most recent call last):
649 ...
650 RuntimeError
652 Note that :func:`one` attempts to advance *iterable* twice to ensure there
653 is only one item. See :func:`spy` or :func:`peekable` to check iterable
654 contents less destructively.
656 """
657 iterator = iter(iterable)
658 for first in iterator:
659 for second in iterator:
660 msg = (
661 f'Expected exactly one item in iterable, but got {first!r}, '
662 f'{second!r}, and perhaps more.'
663 )
664 raise too_long or ValueError(msg)
665 return first
666 raise too_short or ValueError('too few items in iterable (expected 1)')
669def raise_(exception, *args):
670 raise exception(*args)
673def strictly_n(iterable, n, too_short=None, too_long=None):
674 """Validate that *iterable* has exactly *n* items and return them if
675 it does. If it has fewer than *n* items, call function *too_short*
676 with the actual number of items. If it has more than *n* items, call function
677 *too_long* with the number ``n + 1``.
679 >>> iterable = ['a', 'b', 'c', 'd']
680 >>> n = 4
681 >>> list(strictly_n(iterable, n))
682 ['a', 'b', 'c', 'd']
684 Note that the returned iterable must be consumed in order for the check to
685 be made.
687 By default, *too_short* and *too_long* are functions that raise
688 ``ValueError``.
690 >>> list(strictly_n('ab', 3)) # doctest: +IGNORE_EXCEPTION_DETAIL
691 Traceback (most recent call last):
692 ...
693 ValueError: too few items in iterable (got 2)
695 >>> list(strictly_n('abc', 2)) # doctest: +IGNORE_EXCEPTION_DETAIL
696 Traceback (most recent call last):
697 ...
698 ValueError: too many items in iterable (got at least 3)
700 You can instead supply functions that do something else.
701 *too_short* will be called with the number of items in *iterable*.
702 *too_long* will be called with `n + 1`.
704 >>> def too_short(item_count):
705 ... raise RuntimeError
706 >>> it = strictly_n('abcd', 6, too_short=too_short)
707 >>> list(it) # doctest: +IGNORE_EXCEPTION_DETAIL
708 Traceback (most recent call last):
709 ...
710 RuntimeError
712 >>> def too_long(item_count):
713 ... print('The boss is going to hear about this')
714 >>> it = strictly_n('abcdef', 4, too_long=too_long)
715 >>> list(it)
716 The boss is going to hear about this
717 ['a', 'b', 'c', 'd']
719 """
720 if too_short is None:
721 too_short = lambda item_count: raise_(
722 ValueError,
723 f'Too few items in iterable (got {item_count})',
724 )
726 if too_long is None:
727 too_long = lambda item_count: raise_(
728 ValueError,
729 f'Too many items in iterable (got at least {item_count})',
730 )
732 it = iter(iterable)
734 sent = 0
735 for item in islice(it, n):
736 yield item
737 sent += 1
739 if sent < n:
740 too_short(sent)
741 return
743 for item in it:
744 too_long(n + 1)
745 return
748def distinct_permutations(iterable, r=None):
749 """Yield successive distinct permutations of the elements in *iterable*.
751 >>> sorted(distinct_permutations([1, 0, 1]))
752 [(0, 1, 1), (1, 0, 1), (1, 1, 0)]
754 Equivalent to yielding from ``set(permutations(iterable))``, except
755 duplicates are not generated and thrown away. For larger input sequences
756 this is much more efficient.
758 If the elements of the input iterable are sortable, the output tuples are
759 produced in sorted order.
761 Duplicate permutations arise when there are duplicated elements in the
762 input iterable. The number of items returned is
763 `n! / (x_1! * x_2! * ... * x_n!)`, where `n` is the total number of
764 items input, and each `x_i` is the count of a distinct item in the input
765 sequence. The function :func:`multinomial` computes this directly.
767 If *r* is given, only the *r*-length permutations are yielded.
769 >>> sorted(distinct_permutations([1, 0, 1], r=2))
770 [(0, 1), (1, 0), (1, 1)]
771 >>> sorted(distinct_permutations(range(3), r=2))
772 [(0, 1), (0, 2), (1, 0), (1, 2), (2, 0), (2, 1)]
774 *iterable* need not be sortable, but note that using equal (``x == y``)
775 but non-identical (``id(x) != id(y)``) elements may produce surprising
776 behavior. For example, ``1`` and ``True`` are equal but non-identical:
778 >>> list(distinct_permutations([1, True, '3'])) # doctest: +SKIP
779 [
780 (1, True, '3'),
781 (1, '3', True),
782 ('3', 1, True)
783 ]
784 >>> list(distinct_permutations([1, 2, '3'])) # doctest: +SKIP
785 [
786 (1, 2, '3'),
787 (1, '3', 2),
788 (2, 1, '3'),
789 (2, '3', 1),
790 ('3', 1, 2),
791 ('3', 2, 1)
792 ]
793 """
795 # Algorithm: https://w.wiki/Qai
796 def _full(A):
797 while True:
798 # Yield the permutation we have
799 yield tuple(A)
801 # Find the largest index i such that A[i] < A[i + 1]
802 for i in range(size - 2, -1, -1):
803 if A[i] < A[i + 1]:
804 break
805 # If no such index exists, this permutation is the last one
806 else:
807 return
809 # Find the largest index j greater than j such that A[i] < A[j]
810 for j in range(size - 1, i, -1):
811 if A[i] < A[j]:
812 break
814 # Swap the value of A[i] with that of A[j], then reverse the
815 # sequence from A[i + 1] to form the new permutation
816 A[i], A[j] = A[j], A[i]
817 A[i + 1 :] = A[: i - size : -1] # A[i + 1:][::-1]
819 # Algorithm: modified from the above
820 def _partial(A, r):
821 # Split A into the first r items and the last r items
822 head, tail = A[:r], A[r:]
823 right_head_indexes = range(r - 1, -1, -1)
824 left_tail_indexes = range(len(tail))
826 while True:
827 # Yield the permutation we have
828 yield tuple(head)
830 # Starting from the right, find the first index of the head with
831 # value smaller than the maximum value of the tail - call it i.
832 pivot = tail[-1]
833 for i in right_head_indexes:
834 if head[i] < pivot:
835 break
836 pivot = head[i]
837 else:
838 return
840 # Starting from the left, find the first value of the tail
841 # with a value greater than head[i] and swap.
842 for j in left_tail_indexes:
843 if tail[j] > head[i]:
844 head[i], tail[j] = tail[j], head[i]
845 break
846 # If we didn't find one, start from the right and find the first
847 # index of the head with a value greater than head[i] and swap.
848 else:
849 for j in right_head_indexes:
850 if head[j] > head[i]:
851 head[i], head[j] = head[j], head[i]
852 break
854 # Reverse head[i + 1:] and swap it with tail[:r - (i + 1)]
855 tail += head[: i - r : -1] # head[i + 1:][::-1]
856 i += 1
857 head[i:], tail[:] = tail[: r - i], tail[r - i :]
859 items = list(iterable)
861 try:
862 items.sort()
863 sortable = True
864 except TypeError:
865 sortable = False
867 indices_dict = defaultdict(list)
869 for item in items:
870 indices_dict[items.index(item)].append(item)
872 indices = [items.index(item) for item in items]
873 indices.sort()
875 equivalent_items = {k: cycle(v) for k, v in indices_dict.items()}
877 def permuted_items(permuted_indices):
878 return tuple(
879 next(equivalent_items[index]) for index in permuted_indices
880 )
882 size = len(items)
883 if r is None:
884 r = size
886 # functools.partial(_partial, ... )
887 algorithm = _full if (r == size) else partial(_partial, r=r)
889 if 0 < r <= size:
890 if sortable:
891 return algorithm(items)
892 else:
893 return (
894 permuted_items(permuted_indices)
895 for permuted_indices in algorithm(indices)
896 )
898 return iter(() if r else ((),))
901def derangements(iterable, r=None):
902 """Yield successive derangements of the elements in *iterable*.
904 A derangement is a permutation in which no element appears at its original
905 index. In other words, a derangement is a permutation that has no fixed points.
907 Suppose Alice, Bob, Carol, and Dave are playing Secret Santa.
908 The code below outputs all of the different ways to assign gift recipients
909 such that nobody is assigned to himself or herself:
911 >>> for d in derangements(['Alice', 'Bob', 'Carol', 'Dave']):
912 ... print(', '.join(d))
913 Bob, Alice, Dave, Carol
914 Bob, Carol, Dave, Alice
915 Bob, Dave, Alice, Carol
916 Carol, Alice, Dave, Bob
917 Carol, Dave, Alice, Bob
918 Carol, Dave, Bob, Alice
919 Dave, Alice, Bob, Carol
920 Dave, Carol, Alice, Bob
921 Dave, Carol, Bob, Alice
923 If *r* is given, only the *r*-length derangements are yielded.
925 >>> sorted(derangements(range(3), 2))
926 [(1, 0), (1, 2), (2, 0)]
927 >>> sorted(derangements([0, 2, 3], 2))
928 [(2, 0), (2, 3), (3, 0)]
930 Elements are treated as unique based on their position, not on their value.
932 Consider the Secret Santa example with two *different* people who have
933 the *same* name. Then there are two valid gift assignments even though
934 it might appear that a person is assigned to themselves:
936 >>> names = ['Alice', 'Bob', 'Bob']
937 >>> list(derangements(names))
938 [('Bob', 'Bob', 'Alice'), ('Bob', 'Alice', 'Bob')]
940 To avoid confusion, make the inputs distinct:
942 >>> deduped = [f'{name}{index}' for index, name in enumerate(names)]
943 >>> list(derangements(deduped))
944 [('Bob1', 'Bob2', 'Alice0'), ('Bob2', 'Alice0', 'Bob1')]
946 The number of derangements of a set of size *n* is known as the
947 "subfactorial of n". For n > 0, the subfactorial is:
948 ``round(math.factorial(n) / math.e)``. The more-itertools function
949 :func:`subfactorial` computes this directly.
951 References:
953 * Article: https://www.numberanalytics.com/blog/ultimate-guide-to-derangements-in-combinatorics
954 * Sizes: https://oeis.org/A000166
955 """
956 xs = tuple(iterable)
957 ys = tuple(range(len(xs)))
958 return compress(
959 permutations(xs, r=r),
960 map(all, map(map, repeat(is_not), repeat(ys), permutations(ys, r=r))),
961 )
964def intersperse(e, iterable, n=1):
965 """Intersperse filler element *e* among the items in *iterable*, leaving
966 *n* items between each filler element.
968 >>> list(intersperse('!', [1, 2, 3, 4, 5]))
969 [1, '!', 2, '!', 3, '!', 4, '!', 5]
971 >>> list(intersperse(None, [1, 2, 3, 4, 5], n=2))
972 [1, 2, None, 3, 4, None, 5]
974 """
975 if n == 0:
976 raise ValueError('n must be > 0')
977 elif n == 1:
978 # interleave(repeat(e), iterable) -> e, x_0, e, x_1, e, x_2...
979 # islice(..., 1, None) -> x_0, e, x_1, e, x_2...
980 return islice(interleave(repeat(e), iterable), 1, None)
981 else:
982 # interleave(filler, chunks) -> [e], [x_0, x_1], [e], [x_2, x_3]...
983 # islice(..., 1, None) -> [x_0, x_1], [e], [x_2, x_3]...
984 # flatten(...) -> x_0, x_1, e, x_2, x_3...
985 filler = repeat([e])
986 chunks = chunked(iterable, n)
987 return flatten(islice(interleave(filler, chunks), 1, None))
990def unique_to_each(*iterables):
991 """Return the elements from each of the input iterables that aren't in the
992 other input iterables.
994 For example, suppose you have a set of packages, each with a set of
995 dependencies::
997 {'pkg_1': {'A', 'B'}, 'pkg_2': {'B', 'C'}, 'pkg_3': {'B', 'D'}}
999 If you remove one package, which dependencies can also be removed?
1001 If ``pkg_1`` is removed, then ``A`` is no longer necessary - it is not
1002 associated with ``pkg_2`` or ``pkg_3``. Similarly, ``C`` is only needed for
1003 ``pkg_2``, and ``D`` is only needed for ``pkg_3``::
1005 >>> unique_to_each({'A', 'B'}, {'B', 'C'}, {'B', 'D'})
1006 [['A'], ['C'], ['D']]
1008 If there are duplicates in one input iterable that aren't in the others
1009 they will be duplicated in the output. Input order is preserved::
1011 >>> unique_to_each("mississippi", "missouri")
1012 [['p', 'p'], ['o', 'u', 'r']]
1014 It is assumed that the elements of each iterable are hashable.
1016 """
1017 pool = [list(it) for it in iterables]
1018 counts = Counter(chain.from_iterable(map(set, pool)))
1019 uniques = {element for element in counts if counts[element] == 1}
1020 return [list(filter(uniques.__contains__, it)) for it in pool]
1023def windowed(seq, n, fillvalue=None, step=1):
1024 """Return a sliding window of width *n* over the given iterable.
1026 >>> all_windows = windowed([1, 2, 3, 4, 5], 3)
1027 >>> list(all_windows)
1028 [(1, 2, 3), (2, 3, 4), (3, 4, 5)]
1030 When the window is larger than the iterable, *fillvalue* is used in place
1031 of missing values:
1033 >>> list(windowed([1, 2, 3], 4))
1034 [(1, 2, 3, None)]
1036 Each window will advance in increments of *step*:
1038 >>> list(windowed([1, 2, 3, 4, 5, 6], 3, fillvalue='!', step=2))
1039 [(1, 2, 3), (3, 4, 5), (5, 6, '!')]
1041 To slide into the iterable's items, use :func:`chain` to add filler items
1042 to the left:
1044 >>> iterable = [1, 2, 3, 4]
1045 >>> n = 3
1046 >>> padding = [None] * (n - 1)
1047 >>> list(windowed(chain(padding, iterable), 3))
1048 [(None, None, 1), (None, 1, 2), (1, 2, 3), (2, 3, 4)]
1049 """
1050 if n <= 0:
1051 raise ValueError('n must be > 0')
1052 if step < 1:
1053 raise ValueError('step must be >= 1')
1055 iterator = iter(seq)
1057 # Generate first window
1058 window = deque(islice(iterator, n), maxlen=n)
1060 # Deal with the first window not being full
1061 if not window:
1062 return
1063 if len(window) < n:
1064 yield tuple(window) + ((fillvalue,) * (n - len(window)))
1065 return
1066 yield tuple(window)
1068 # Create the filler for the next windows. The padding ensures
1069 # we have just enough elements to fill the last window.
1070 padding = (fillvalue,) * (n - 1 if step >= n else step - 1)
1071 filler = map(window.append, chain(iterator, padding))
1073 # Generate the rest of the windows
1074 for _ in islice(filler, step - 1, None, step):
1075 yield tuple(window)
1078def substrings(iterable):
1079 """Yield all of the substrings of *iterable*.
1081 >>> [''.join(s) for s in substrings('more')]
1082 ['m', 'o', 'r', 'e', 'mo', 'or', 're', 'mor', 'ore', 'more']
1084 Note that non-string iterables can also be subdivided.
1086 >>> list(substrings([0, 1, 2]))
1087 [(0,), (1,), (2,), (0, 1), (1, 2), (0, 1, 2)]
1089 Like subslices() but returns tuples instead of lists
1090 and returns the shortest substrings first.
1092 """
1093 seq = tuple(iterable)
1094 item_count = len(seq)
1095 for n in range(1, item_count + 1):
1096 slices = map(slice, range(item_count), range(n, item_count + 1))
1097 yield from map(getitem, repeat(seq), slices)
1100def substrings_indexes(seq, reverse=False):
1101 """Yield all substrings and their positions in *seq*
1103 The items yielded will be a tuple of the form ``(substr, i, j)``, where
1104 ``substr == seq[i:j]``.
1106 This function only works for iterables that support slicing, such as
1107 ``str`` objects.
1109 >>> for item in substrings_indexes('more'):
1110 ... print(item)
1111 ('m', 0, 1)
1112 ('o', 1, 2)
1113 ('r', 2, 3)
1114 ('e', 3, 4)
1115 ('mo', 0, 2)
1116 ('or', 1, 3)
1117 ('re', 2, 4)
1118 ('mor', 0, 3)
1119 ('ore', 1, 4)
1120 ('more', 0, 4)
1122 Set *reverse* to ``True`` to yield the same items in the opposite order.
1125 """
1126 r = range(1, len(seq) + 1)
1127 if reverse:
1128 r = reversed(r)
1129 return (
1130 (seq[i : i + L], i, i + L) for L in r for i in range(len(seq) - L + 1)
1131 )
1134class bucket:
1135 """Wrap *iterable* and return an object that buckets the iterable into
1136 child iterables based on a *key* function.
1138 >>> iterable = ['a1', 'b1', 'c1', 'a2', 'b2', 'c2', 'b3']
1139 >>> s = bucket(iterable, key=lambda x: x[0]) # Bucket by 1st character
1140 >>> sorted(list(s)) # Get the keys
1141 ['a', 'b', 'c']
1142 >>> a_iterable = s['a']
1143 >>> next(a_iterable)
1144 'a1'
1145 >>> next(a_iterable)
1146 'a2'
1147 >>> list(s['b'])
1148 ['b1', 'b2', 'b3']
1150 The original iterable will be advanced and its items will be cached until
1151 they are used by the child iterables. This may require significant storage.
1153 By default, attempting to select a bucket to which no items belong will
1154 exhaust the iterable and cache all values.
1155 If you specify a *validator* function, selected buckets will instead be
1156 checked against it.
1158 >>> from itertools import count
1159 >>> it = count(1, 2) # Infinite sequence of odd numbers
1160 >>> key = lambda x: x % 10 # Bucket by last digit
1161 >>> validator = lambda x: x in {1, 3, 5, 7, 9} # Odd digits only
1162 >>> s = bucket(it, key=key, validator=validator)
1163 >>> 2 in s
1164 False
1165 >>> list(s[2])
1166 []
1168 .. seealso:: :func:`map_reduce`, :func:`groupby_transform`
1170 If storage is not a concern, :func:`map_reduce` returns a Python
1171 dictionary, which is generally easier to work with. If the elements
1172 with the same key are already adjacent, :func:`groupby_transform`
1173 or :func:`itertools.groupby` can be used without any caching overhead.
1175 """
1177 def __init__(self, iterable, key, validator=None):
1178 self._it = iter(iterable)
1179 self._key = key
1180 self._cache = defaultdict(deque)
1181 self._validator = validator or (lambda x: True)
1183 def __contains__(self, value):
1184 if not self._validator(value):
1185 return False
1187 try:
1188 item = next(self[value])
1189 except StopIteration:
1190 return False
1191 else:
1192 self._cache[value].appendleft(item)
1194 return True
1196 def _get_values(self, value):
1197 """
1198 Helper to yield items from the parent iterator that match *value*.
1199 Items that don't match are stored in the local cache as they
1200 are encountered.
1201 """
1202 while True:
1203 # If we've cached some items that match the target value, emit
1204 # the first one and evict it from the cache.
1205 if self._cache[value]:
1206 yield self._cache[value].popleft()
1207 # Otherwise we need to advance the parent iterator to search for
1208 # a matching item, caching the rest.
1209 else:
1210 while True:
1211 try:
1212 item = next(self._it)
1213 except StopIteration:
1214 return
1215 item_value = self._key(item)
1216 if item_value == value:
1217 yield item
1218 break
1219 elif self._validator(item_value):
1220 self._cache[item_value].append(item)
1222 def __iter__(self):
1223 for item in self._it:
1224 item_value = self._key(item)
1225 if self._validator(item_value):
1226 self._cache[item_value].append(item)
1228 return iter(self._cache)
1230 def __getitem__(self, value):
1231 if not self._validator(value):
1232 return iter(())
1234 return self._get_values(value)
1237def spy(iterable, n=1):
1238 """Return a 2-tuple with a list containing the first *n* elements of
1239 *iterable*, and an iterator with the same items as *iterable*.
1240 This allows you to "look ahead" at the items in the iterable without
1241 advancing it.
1243 There is one item in the list by default:
1245 >>> iterable = 'abcdefg'
1246 >>> head, iterable = spy(iterable)
1247 >>> head
1248 ['a']
1249 >>> list(iterable)
1250 ['a', 'b', 'c', 'd', 'e', 'f', 'g']
1252 You may use unpacking to retrieve items instead of lists:
1254 >>> (head,), iterable = spy('abcdefg')
1255 >>> head
1256 'a'
1257 >>> (first, second), iterable = spy('abcdefg', 2)
1258 >>> first
1259 'a'
1260 >>> second
1261 'b'
1263 The number of items requested can be larger than the number of items in
1264 the iterable:
1266 >>> iterable = [1, 2, 3, 4, 5]
1267 >>> head, iterable = spy(iterable, 10)
1268 >>> head
1269 [1, 2, 3, 4, 5]
1270 >>> list(iterable)
1271 [1, 2, 3, 4, 5]
1273 """
1274 p, q = tee(iterable)
1275 return take(n, q), p
1278def interleave(*iterables):
1279 """Return a new iterable yielding from each iterable in turn,
1280 until the shortest is exhausted.
1282 >>> list(interleave([1, 2, 3], [4, 5], [6, 7, 8]))
1283 [1, 4, 6, 2, 5, 7]
1285 For a version that doesn't terminate after the shortest iterable is
1286 exhausted, see :func:`interleave_longest`.
1288 """
1289 return chain.from_iterable(zip(*iterables))
1292def interleave_longest(*iterables):
1293 """Return a new iterable yielding from each iterable in turn,
1294 skipping any that are exhausted.
1296 >>> list(interleave_longest([1, 2, 3], [4, 5], [6, 7, 8]))
1297 [1, 4, 6, 2, 5, 7, 3, 8]
1299 This function produces the same output as :func:`roundrobin`, but may
1300 perform better for some inputs (in particular when the number of iterables
1301 is large).
1303 """
1304 for xs in zip_longest(*iterables, fillvalue=_marker):
1305 for x in xs:
1306 if x is not _marker:
1307 yield x
1310def interleave_evenly(iterables, lengths=None):
1311 """
1312 Interleave multiple iterables so that their elements are evenly distributed
1313 throughout the output sequence.
1315 >>> iterables = [1, 2, 3, 4, 5], ['a', 'b']
1316 >>> list(interleave_evenly(iterables))
1317 [1, 2, 'a', 3, 4, 'b', 5]
1319 >>> iterables = [[1, 2, 3], [4, 5], [6, 7, 8]]
1320 >>> list(interleave_evenly(iterables))
1321 [1, 6, 4, 2, 7, 3, 8, 5]
1323 This function requires iterables of known length. Iterables without
1324 ``__len__()`` can be used by manually specifying lengths with *lengths*:
1326 >>> from itertools import combinations, repeat
1327 >>> iterables = [combinations(range(4), 2), ['a', 'b', 'c']]
1328 >>> lengths = [4 * (4 - 1) // 2, 3]
1329 >>> list(interleave_evenly(iterables, lengths=lengths))
1330 [(0, 1), (0, 2), 'a', (0, 3), (1, 2), 'b', (1, 3), (2, 3), 'c']
1332 Based on Bresenham's algorithm.
1333 """
1334 if lengths is None:
1335 try:
1336 lengths = [len(it) for it in iterables]
1337 except TypeError:
1338 raise ValueError(
1339 'Iterable lengths could not be determined automatically. '
1340 'Specify them with the lengths keyword.'
1341 )
1342 elif len(iterables) != len(lengths):
1343 raise ValueError('Mismatching number of iterables and lengths.')
1345 dims = len(lengths)
1347 if not dims:
1348 return
1350 # sort iterables by length, descending
1351 lengths_permute = sorted(
1352 range(dims), key=lambda i: lengths[i], reverse=True
1353 )
1354 lengths_desc = [lengths[i] for i in lengths_permute]
1355 iters_desc = [iter(iterables[i]) for i in lengths_permute]
1357 # the longest iterable is the primary one (Bresenham: the longest
1358 # distance along an axis)
1359 delta_primary, deltas_secondary = lengths_desc[0], lengths_desc[1:]
1360 iter_primary, iters_secondary = iters_desc[0], iters_desc[1:]
1361 errors = [delta_primary // dims] * len(deltas_secondary)
1363 to_yield = sum(lengths)
1364 while to_yield:
1365 yield next(iter_primary)
1366 to_yield -= 1
1367 # update errors for each secondary iterable
1368 errors = [e - delta for e, delta in zip(errors, deltas_secondary)]
1370 # those iterables for which the error is negative are yielded
1371 # ("diagonal step" in Bresenham)
1372 for i, e_ in enumerate(errors):
1373 if e_ < 0:
1374 yield next(iters_secondary[i])
1375 to_yield -= 1
1376 errors[i] += delta_primary
1379def interleave_randomly(*iterables):
1380 """Repeatedly select one of the input *iterables* at random and yield the next
1381 item from it.
1383 >>> iterables = [1, 2, 3], 'abc', (True, False, None)
1384 >>> list(interleave_randomly(*iterables)) # doctest: +SKIP
1385 ['a', 'b', 1, 'c', True, False, None, 2, 3]
1387 The relative order of the items in each input iterable will preserved. Note the
1388 sequences of items with this property are not equally likely to be generated.
1390 """
1391 iterators = [iter(e) for e in iterables]
1392 while iterators:
1393 idx = randrange(len(iterators))
1394 try:
1395 yield next(iterators[idx])
1396 except StopIteration:
1397 # equivalent to `list.pop` but slightly faster
1398 iterators[idx] = iterators[-1]
1399 del iterators[-1]
1402def collapse(iterable, base_type=None, levels=None):
1403 """Flatten an iterable with multiple levels of nesting (e.g., a list of
1404 lists of tuples) into non-iterable types.
1406 >>> iterable = [(1, 2), ([3, 4], [[5], [6]])]
1407 >>> list(collapse(iterable))
1408 [1, 2, 3, 4, 5, 6]
1410 Binary and text strings are not considered iterable and
1411 will not be collapsed.
1413 To avoid collapsing other types, specify *base_type*:
1415 >>> iterable = ['ab', ('cd', 'ef'), ['gh', 'ij']]
1416 >>> list(collapse(iterable, base_type=tuple))
1417 ['ab', ('cd', 'ef'), 'gh', 'ij']
1419 Specify *levels* to stop flattening after a certain level:
1421 >>> iterable = [('a', ['b']), ('c', ['d'])]
1422 >>> list(collapse(iterable)) # Fully flattened
1423 ['a', 'b', 'c', 'd']
1424 >>> list(collapse(iterable, levels=1)) # Only one level flattened
1425 ['a', ['b'], 'c', ['d']]
1427 """
1428 stack = deque()
1429 # Add our first node group, treat the iterable as a single node
1430 stack.appendleft((0, repeat(iterable, 1)))
1432 while stack:
1433 node_group = stack.popleft()
1434 level, nodes = node_group
1436 # Check if beyond max level
1437 if levels is not None and level > levels:
1438 yield from nodes
1439 continue
1441 for node in nodes:
1442 # Check if done iterating
1443 if isinstance(node, (str, bytes)) or (
1444 (base_type is not None) and isinstance(node, base_type)
1445 ):
1446 yield node
1447 # Otherwise try to create child nodes
1448 else:
1449 try:
1450 tree = iter(node)
1451 except TypeError:
1452 yield node
1453 else:
1454 # Save our current location
1455 stack.appendleft(node_group)
1456 # Append the new child node
1457 stack.appendleft((level + 1, tree))
1458 # Break to process child node
1459 break
1462def side_effect(func, iterable, chunk_size=None, before=None, after=None):
1463 """Invoke *func* on each item in *iterable* (or on each *chunk_size* group
1464 of items) before yielding the item.
1466 `func` must be a function that takes a single argument. Its return value
1467 will be discarded.
1469 *before* and *after* are optional functions that take no arguments. They
1470 will be executed before iteration starts and after it ends, respectively.
1472 `side_effect` can be used for logging, updating progress bars, or anything
1473 that is not functionally "pure."
1475 Emitting a status message:
1477 >>> from more_itertools import consume
1478 >>> func = lambda item: print('Received {}'.format(item))
1479 >>> consume(side_effect(func, range(2)))
1480 Received 0
1481 Received 1
1483 Operating on chunks of items:
1485 >>> pair_sums = []
1486 >>> func = lambda chunk: pair_sums.append(sum(chunk))
1487 >>> list(side_effect(func, [0, 1, 2, 3, 4, 5], 2))
1488 [0, 1, 2, 3, 4, 5]
1489 >>> list(pair_sums)
1490 [1, 5, 9]
1492 Writing to a file-like object:
1494 >>> from io import StringIO
1495 >>> from more_itertools import consume
1496 >>> f = StringIO()
1497 >>> func = lambda x: print(x, file=f)
1498 >>> before = lambda: print('HEADER', file=f)
1499 >>> after = f.close
1500 >>> it = ['a', 'b', 'c']
1501 >>> consume(side_effect(func, it, before=before, after=after))
1502 >>> f.closed
1503 True
1505 """
1506 try:
1507 if before is not None:
1508 before()
1510 if chunk_size is None:
1511 for item in iterable:
1512 func(item)
1513 yield item
1514 else:
1515 for chunk in chunked(iterable, chunk_size):
1516 func(chunk)
1517 yield from chunk
1518 finally:
1519 if after is not None:
1520 after()
1523def sliced(seq, n, strict=False):
1524 """Yield slices of length *n* from the sequence *seq*.
1526 >>> list(sliced((1, 2, 3, 4, 5, 6), 3))
1527 [(1, 2, 3), (4, 5, 6)]
1529 By the default, the last yielded slice will have fewer than *n* elements
1530 if the length of *seq* is not divisible by *n*:
1532 >>> list(sliced((1, 2, 3, 4, 5, 6, 7, 8), 3))
1533 [(1, 2, 3), (4, 5, 6), (7, 8)]
1535 If the length of *seq* is not divisible by *n* and *strict* is
1536 ``True``, then ``ValueError`` will be raised before the last
1537 slice is yielded.
1539 This function will only work for iterables that support slicing.
1540 For non-sliceable iterables, see :func:`chunked`.
1542 """
1543 if n < 0:
1544 raise ValueError('n must be at least 0')
1546 iterator = takewhile(len, (seq[i : i + n] for i in count(0, n)))
1547 if strict:
1549 def ret():
1550 for _slice in iterator:
1551 if len(_slice) != n:
1552 raise ValueError("seq is not divisible by n.")
1553 yield _slice
1555 return ret()
1556 else:
1557 return iterator
1560def split_at(iterable, pred, maxsplit=-1, keep_separator=False):
1561 """Yield lists of items from *iterable*, where each list is delimited by
1562 an item where callable *pred* returns ``True``.
1564 >>> list(split_at('abcdcba', lambda x: x == 'b'))
1565 [['a'], ['c', 'd', 'c'], ['a']]
1567 >>> list(split_at(range(10), lambda n: n % 2 == 1))
1568 [[0], [2], [4], [6], [8], []]
1570 At most *maxsplit* splits are done. If *maxsplit* is not specified or -1,
1571 then there is no limit on the number of splits:
1573 >>> list(split_at(range(10), lambda n: n % 2 == 1, maxsplit=2))
1574 [[0], [2], [4, 5, 6, 7, 8, 9]]
1576 By default, the delimiting items are not included in the output.
1577 To include them, set *keep_separator* to ``True``.
1579 >>> list(split_at('abcdcba', lambda x: x == 'b', keep_separator=True))
1580 [['a'], ['b'], ['c', 'd', 'c'], ['b'], ['a']]
1582 """
1583 if maxsplit == 0:
1584 yield list(iterable)
1585 return
1587 buf = []
1588 it = iter(iterable)
1589 for item in it:
1590 if pred(item):
1591 yield buf
1592 if keep_separator:
1593 yield [item]
1594 if maxsplit == 1:
1595 yield list(it)
1596 return
1597 buf = []
1598 maxsplit -= 1
1599 else:
1600 buf.append(item)
1601 yield buf
1604def split_before(iterable, pred, maxsplit=-1):
1605 """Yield lists of items from *iterable*, where each list ends just before
1606 an item for which callable *pred* returns ``True``:
1608 >>> list(split_before('OneTwo', lambda s: s.isupper()))
1609 [['O', 'n', 'e'], ['T', 'w', 'o']]
1611 >>> list(split_before(range(10), lambda n: n % 3 == 0))
1612 [[0, 1, 2], [3, 4, 5], [6, 7, 8], [9]]
1614 At most *maxsplit* splits are done. If *maxsplit* is not specified or -1,
1615 then there is no limit on the number of splits:
1617 >>> list(split_before(range(10), lambda n: n % 3 == 0, maxsplit=2))
1618 [[0, 1, 2], [3, 4, 5], [6, 7, 8, 9]]
1619 """
1620 if maxsplit == 0:
1621 buf = list(iterable)
1622 if buf:
1623 yield buf
1624 return
1626 buf = []
1627 it = iter(iterable)
1628 for item in it:
1629 if pred(item) and buf:
1630 yield buf
1631 if maxsplit == 1:
1632 yield [item, *it]
1633 return
1634 buf = []
1635 maxsplit -= 1
1636 buf.append(item)
1637 if buf:
1638 yield buf
1641def split_after(iterable, pred, maxsplit=-1):
1642 """Yield lists of items from *iterable*, where each list ends with an
1643 item where callable *pred* returns ``True``:
1645 >>> list(split_after('one1two2', lambda s: s.isdigit()))
1646 [['o', 'n', 'e', '1'], ['t', 'w', 'o', '2']]
1648 >>> list(split_after(range(10), lambda n: n % 3 == 0))
1649 [[0], [1, 2, 3], [4, 5, 6], [7, 8, 9]]
1651 At most *maxsplit* splits are done. If *maxsplit* is not specified or -1,
1652 then there is no limit on the number of splits:
1654 >>> list(split_after(range(10), lambda n: n % 3 == 0, maxsplit=2))
1655 [[0], [1, 2, 3], [4, 5, 6, 7, 8, 9]]
1657 """
1658 if maxsplit == 0:
1659 buf = list(iterable)
1660 if buf:
1661 yield buf
1662 return
1664 buf = []
1665 it = iter(iterable)
1666 for item in it:
1667 buf.append(item)
1668 if pred(item) and buf:
1669 yield buf
1670 if maxsplit == 1:
1671 buf = list(it)
1672 if buf:
1673 yield buf
1674 return
1675 buf = []
1676 maxsplit -= 1
1677 if buf:
1678 yield buf
1681def split_when(iterable, pred, maxsplit=-1):
1682 """Split *iterable* into pieces based on the output of *pred*.
1683 *pred* should be a function that takes successive pairs of items and
1684 returns ``True`` if the iterable should be split in between them.
1686 For example, to find runs of increasing numbers, split the iterable when
1687 element ``i`` is larger than element ``i + 1``:
1689 >>> list(split_when([1, 2, 3, 3, 2, 5, 2, 4, 2], lambda x, y: x > y))
1690 [[1, 2, 3, 3], [2, 5], [2, 4], [2]]
1692 At most *maxsplit* splits are done. If *maxsplit* is not specified or -1,
1693 then there is no limit on the number of splits:
1695 >>> list(split_when([1, 2, 3, 3, 2, 5, 2, 4, 2],
1696 ... lambda x, y: x > y, maxsplit=2))
1697 [[1, 2, 3, 3], [2, 5], [2, 4, 2]]
1699 """
1700 if maxsplit == 0:
1701 buf = list(iterable)
1702 if buf:
1703 yield buf
1704 return
1706 it = iter(iterable)
1707 try:
1708 cur_item = next(it)
1709 except StopIteration:
1710 return
1712 buf = [cur_item]
1713 for next_item in it:
1714 if pred(cur_item, next_item):
1715 yield buf
1716 if maxsplit == 1:
1717 yield [next_item, *it]
1718 return
1719 buf = []
1720 maxsplit -= 1
1722 buf.append(next_item)
1723 cur_item = next_item
1725 yield buf
1728def split_into(iterable, sizes):
1729 """Yield a list of sequential items from *iterable* of length 'n' for each
1730 integer 'n' in *sizes*.
1732 >>> list(split_into([1,2,3,4,5,6], [1,2,3]))
1733 [[1], [2, 3], [4, 5, 6]]
1735 If the sum of *sizes* is smaller than the length of *iterable*, then the
1736 remaining items of *iterable* will not be returned.
1738 >>> list(split_into([1,2,3,4,5,6], [2,3]))
1739 [[1, 2], [3, 4, 5]]
1741 If the sum of *sizes* is larger than the length of *iterable*, fewer items
1742 will be returned in the iteration that overruns the *iterable* and further
1743 lists will be empty:
1745 >>> list(split_into([1,2,3,4], [1,2,3,4]))
1746 [[1], [2, 3], [4], []]
1748 When a ``None`` object is encountered in *sizes*, the returned list will
1749 contain items up to the end of *iterable* the same way that
1750 :func:`itertools.slice` does:
1752 >>> list(split_into([1,2,3,4,5,6,7,8,9,0], [2,3,None]))
1753 [[1, 2], [3, 4, 5], [6, 7, 8, 9, 0]]
1755 :func:`split_into` can be useful for grouping a series of items where the
1756 sizes of the groups are not uniform. An example would be where in a row
1757 from a table, multiple columns represent elements of the same feature
1758 (e.g. a point represented by x,y,z) but, the format is not the same for
1759 all columns.
1760 """
1761 # convert the iterable argument into an iterator so its contents can
1762 # be consumed by islice in case it is a generator
1763 it = iter(iterable)
1765 for size in sizes:
1766 if size is None:
1767 yield list(it)
1768 return
1769 else:
1770 yield list(islice(it, size))
1773def padded(iterable, fillvalue=None, n=None, next_multiple=False):
1774 """Yield the elements from *iterable*, followed by *fillvalue*, such that
1775 at least *n* items are emitted.
1777 >>> list(padded([1, 2, 3], '?', 5))
1778 [1, 2, 3, '?', '?']
1780 If *next_multiple* is ``True``, *fillvalue* will be emitted until the
1781 number of items emitted is a multiple of *n*:
1783 >>> list(padded([1, 2, 3, 4], n=3, next_multiple=True))
1784 [1, 2, 3, 4, None, None]
1786 If *n* is ``None``, *fillvalue* will be emitted indefinitely.
1788 To create an *iterable* of exactly size *n*, you can truncate with
1789 :func:`islice`.
1791 >>> list(islice(padded([1, 2, 3], '?'), 5))
1792 [1, 2, 3, '?', '?']
1793 >>> list(islice(padded([1, 2, 3, 4, 5, 6, 7, 8], '?'), 5))
1794 [1, 2, 3, 4, 5]
1796 """
1797 iterator = iter(iterable)
1798 iterator_with_repeat = chain(iterator, repeat(fillvalue))
1800 if n is None:
1801 return iterator_with_repeat
1802 elif n < 1:
1803 raise ValueError('n must be at least 1')
1804 elif next_multiple:
1806 def slice_generator():
1807 for first in iterator:
1808 yield (first,)
1809 yield islice(iterator_with_repeat, n - 1)
1811 # While elements exist produce slices of size n
1812 return chain.from_iterable(slice_generator())
1813 else:
1814 # Ensure the first batch is at least size n then iterate
1815 return chain(islice(iterator_with_repeat, n), iterator)
1818def repeat_each(iterable, n=2):
1819 """Repeat each element in *iterable* *n* times.
1821 >>> list(repeat_each('ABC', 3))
1822 ['A', 'A', 'A', 'B', 'B', 'B', 'C', 'C', 'C']
1823 """
1824 return chain.from_iterable(map(repeat, iterable, repeat(n)))
1827def repeat_last(iterable, default=None):
1828 """After the *iterable* is exhausted, keep yielding its last element.
1830 >>> list(islice(repeat_last(range(3)), 5))
1831 [0, 1, 2, 2, 2]
1833 If the iterable is empty, yield *default* forever::
1835 >>> list(islice(repeat_last(range(0), 42), 5))
1836 [42, 42, 42, 42, 42]
1838 """
1839 item = _marker
1840 for item in iterable:
1841 yield item
1842 final = default if item is _marker else item
1843 yield from repeat(final)
1846def distribute(n, iterable):
1847 """Distribute the items from *iterable* among *n* smaller iterables.
1849 >>> group_1, group_2 = distribute(2, [1, 2, 3, 4, 5, 6])
1850 >>> list(group_1)
1851 [1, 3, 5]
1852 >>> list(group_2)
1853 [2, 4, 6]
1855 If the length of *iterable* is not evenly divisible by *n*, then the
1856 length of the returned iterables will not be identical:
1858 >>> children = distribute(3, [1, 2, 3, 4, 5, 6, 7])
1859 >>> [list(c) for c in children]
1860 [[1, 4, 7], [2, 5], [3, 6]]
1862 If the length of *iterable* is smaller than *n*, then the last returned
1863 iterables will be empty:
1865 >>> children = distribute(5, [1, 2, 3])
1866 >>> [list(c) for c in children]
1867 [[1], [2], [3], [], []]
1869 This function uses :func:`itertools.tee` and may require significant
1870 storage.
1872 If you need the order items in the smaller iterables to match the
1873 original iterable, see :func:`divide`.
1875 """
1876 if n < 1:
1877 raise ValueError('n must be at least 1')
1879 children = tee(iterable, n)
1880 return [islice(it, index, None, n) for index, it in enumerate(children)]
1883def stagger(iterable, offsets=(-1, 0, 1), longest=False, fillvalue=None):
1884 """Yield tuples whose elements are offset from *iterable*.
1885 The amount by which the `i`-th item in each tuple is offset is given by
1886 the `i`-th item in *offsets*.
1888 >>> list(stagger([0, 1, 2, 3]))
1889 [(None, 0, 1), (0, 1, 2), (1, 2, 3)]
1890 >>> list(stagger(range(8), offsets=(0, 2, 4)))
1891 [(0, 2, 4), (1, 3, 5), (2, 4, 6), (3, 5, 7)]
1893 By default, the sequence will end when the final element of a tuple is the
1894 last item in the iterable. To continue until the first element of a tuple
1895 is the last item in the iterable, set *longest* to ``True``::
1897 >>> list(stagger([0, 1, 2, 3], longest=True))
1898 [(None, 0, 1), (0, 1, 2), (1, 2, 3), (2, 3, None), (3, None, None)]
1900 By default, ``None`` will be used to replace offsets beyond the end of the
1901 sequence. Specify *fillvalue* to use some other value.
1903 """
1904 children = tee(iterable, len(offsets))
1906 return zip_offset(
1907 *children, offsets=offsets, longest=longest, fillvalue=fillvalue
1908 )
1911def zip_offset(*iterables, offsets, longest=False, fillvalue=None):
1912 """``zip`` the input *iterables* together, but offset the `i`-th iterable
1913 by the `i`-th item in *offsets*.
1915 >>> list(zip_offset('0123', 'abcdef', offsets=(0, 1)))
1916 [('0', 'b'), ('1', 'c'), ('2', 'd'), ('3', 'e')]
1918 This can be used as a lightweight alternative to SciPy or pandas to analyze
1919 data sets in which some series have a lead or lag relationship.
1921 By default, the sequence will end when the shortest iterable is exhausted.
1922 To continue until the longest iterable is exhausted, set *longest* to
1923 ``True``.
1925 >>> list(zip_offset('0123', 'abcdef', offsets=(0, 1), longest=True))
1926 [('0', 'b'), ('1', 'c'), ('2', 'd'), ('3', 'e'), (None, 'f')]
1928 By default, ``None`` will be used to replace offsets beyond the end of the
1929 sequence. Specify *fillvalue* to use some other value.
1931 """
1932 if len(iterables) != len(offsets):
1933 raise ValueError("Number of iterables and offsets didn't match")
1935 staggered = []
1936 for it, n in zip(iterables, offsets):
1937 if n < 0:
1938 staggered.append(chain(repeat(fillvalue, -n), it))
1939 elif n > 0:
1940 staggered.append(islice(it, n, None))
1941 else:
1942 staggered.append(it)
1944 if longest:
1945 return zip_longest(*staggered, fillvalue=fillvalue)
1947 return zip(*staggered)
1950def sort_together(
1951 iterables, key_list=(0,), key=None, reverse=False, strict=False
1952):
1953 """Return the input iterables sorted together, with *key_list* as the
1954 priority for sorting. All iterables are trimmed to the length of the
1955 shortest one.
1957 This can be used like the sorting function in a spreadsheet. If each
1958 iterable represents a column of data, the key list determines which
1959 columns are used for sorting.
1961 By default, all iterables are sorted using the ``0``-th iterable::
1963 >>> iterables = [(4, 3, 2, 1), ('a', 'b', 'c', 'd')]
1964 >>> sort_together(iterables)
1965 [(1, 2, 3, 4), ('d', 'c', 'b', 'a')]
1967 Set a different key list to sort according to another iterable.
1968 Specifying multiple keys dictates how ties are broken::
1970 >>> iterables = [(3, 1, 2), (0, 1, 0), ('c', 'b', 'a')]
1971 >>> sort_together(iterables, key_list=(1, 2))
1972 [(2, 3, 1), (0, 0, 1), ('a', 'c', 'b')]
1974 To sort by a function of the elements of the iterable, pass a *key*
1975 function. Its arguments are the elements of the iterables corresponding to
1976 the key list::
1978 >>> names = ('a', 'b', 'c')
1979 >>> lengths = (1, 2, 3)
1980 >>> widths = (5, 2, 1)
1981 >>> def area(length, width):
1982 ... return length * width
1983 >>> sort_together([names, lengths, widths], key_list=(1, 2), key=area)
1984 [('c', 'b', 'a'), (3, 2, 1), (1, 2, 5)]
1986 Set *reverse* to ``True`` to sort in descending order.
1988 >>> sort_together([(1, 2, 3), ('c', 'b', 'a')], reverse=True)
1989 [(3, 2, 1), ('a', 'b', 'c')]
1991 If the *strict* keyword argument is ``True``, then
1992 ``ValueError`` will be raised if any of the iterables have
1993 different lengths.
1995 """
1996 if key is None:
1997 # if there is no key function, the key argument to sorted is an
1998 # itemgetter
1999 key_argument = itemgetter(*key_list)
2000 else:
2001 # if there is a key function, call it with the items at the offsets
2002 # specified by the key function as arguments
2003 key_list = list(key_list)
2004 if len(key_list) == 1:
2005 # if key_list contains a single item, pass the item at that offset
2006 # as the only argument to the key function
2007 key_offset = key_list[0]
2008 key_argument = lambda zipped_items: key(zipped_items[key_offset])
2009 else:
2010 # if key_list contains multiple items, use itemgetter to return a
2011 # tuple of items, which we pass as *args to the key function
2012 get_key_items = itemgetter(*key_list)
2013 key_argument = lambda zipped_items: key(
2014 *get_key_items(zipped_items)
2015 )
2017 transposed = zip(*iterables, strict=strict)
2018 reordered = sorted(transposed, key=key_argument, reverse=reverse)
2019 untransposed = zip(*reordered, strict=strict)
2020 return list(untransposed)
2023def unzip(iterable):
2024 """The inverse of :func:`zip`, this function disaggregates the elements
2025 of the zipped *iterable*.
2027 The ``i``-th iterable contains the ``i``-th element from each element
2028 of the zipped iterable. The first element is used to determine the
2029 length of the remaining elements.
2031 >>> iterable = [('a', 1), ('b', 2), ('c', 3), ('d', 4)]
2032 >>> letters, numbers = unzip(iterable)
2033 >>> list(letters)
2034 ['a', 'b', 'c', 'd']
2035 >>> list(numbers)
2036 [1, 2, 3, 4]
2038 This is similar to using ``zip(*iterable)``, but it avoids reading
2039 *iterable* into memory. Note, however, that this function uses
2040 :func:`itertools.tee` and thus may require significant storage.
2042 """
2043 head, iterable = spy(iterable)
2044 if not head:
2045 # empty iterable, e.g. zip([], [], [])
2046 return ()
2047 # spy returns a one-length iterable as head
2048 head = head[0]
2049 iterables = tee(iterable, len(head))
2051 # If we have an iterable like iter([(1, 2, 3), (4, 5), (6,)]),
2052 # the second unzipped iterable fails at the third tuple since
2053 # it tries to access (6,)[1].
2054 # Same with the third unzipped iterable and the second tuple.
2055 # To support these "improperly zipped" iterables, we suppress
2056 # the IndexError, which just stops the unzipped iterables at
2057 # first length mismatch.
2058 return tuple(
2059 iter_suppress(map(itemgetter(i), it), IndexError)
2060 for i, it in enumerate(iterables)
2061 )
2064def divide(n, iterable):
2065 """Divide the elements from *iterable* into *n* parts, maintaining
2066 order.
2068 >>> group_1, group_2 = divide(2, [1, 2, 3, 4, 5, 6])
2069 >>> list(group_1)
2070 [1, 2, 3]
2071 >>> list(group_2)
2072 [4, 5, 6]
2074 If the length of *iterable* is not evenly divisible by *n*, then the
2075 length of the returned iterables will not be identical:
2077 >>> children = divide(3, [1, 2, 3, 4, 5, 6, 7])
2078 >>> [list(c) for c in children]
2079 [[1, 2, 3], [4, 5], [6, 7]]
2081 If the length of the iterable is smaller than n, then the last returned
2082 iterables will be empty:
2084 >>> children = divide(5, [1, 2, 3])
2085 >>> [list(c) for c in children]
2086 [[1], [2], [3], [], []]
2088 This function will exhaust the iterable before returning.
2089 If order is not important, see :func:`distribute`, which does not first
2090 pull the iterable into memory.
2092 """
2093 if n < 1:
2094 raise ValueError('n must be at least 1')
2096 try:
2097 iterable[:0]
2098 except TypeError:
2099 seq = tuple(iterable)
2100 else:
2101 seq = iterable
2103 q, r = divmod(len(seq), n)
2105 ret = []
2106 stop = 0
2107 for i in range(1, n + 1):
2108 start = stop
2109 stop += q + 1 if i <= r else q
2110 ret.append(iter(seq[start:stop]))
2112 return ret
2115def always_iterable(obj, base_type=(str, bytes)):
2116 """If *obj* is iterable, return an iterator over its items::
2118 >>> obj = (1, 2, 3)
2119 >>> list(always_iterable(obj))
2120 [1, 2, 3]
2122 If *obj* is not iterable, return a one-item iterable containing *obj*::
2124 >>> obj = 1
2125 >>> list(always_iterable(obj))
2126 [1]
2128 If *obj* is ``None``, return an empty iterable:
2130 >>> obj = None
2131 >>> list(always_iterable(None))
2132 []
2134 By default, binary and text strings are not considered iterable::
2136 >>> obj = 'foo'
2137 >>> list(always_iterable(obj))
2138 ['foo']
2140 If *base_type* is set, objects for which ``isinstance(obj, base_type)``
2141 returns ``True`` won't be considered iterable.
2143 >>> obj = {'a': 1}
2144 >>> list(always_iterable(obj)) # Iterate over the dict's keys
2145 ['a']
2146 >>> list(always_iterable(obj, base_type=dict)) # Treat dicts as a unit
2147 [{'a': 1}]
2149 Set *base_type* to ``None`` to avoid any special handling and treat objects
2150 Python considers iterable as iterable:
2152 >>> obj = 'foo'
2153 >>> list(always_iterable(obj, base_type=None))
2154 ['f', 'o', 'o']
2155 """
2156 if obj is None:
2157 return iter(())
2159 if (base_type is not None) and isinstance(obj, base_type):
2160 return iter((obj,))
2162 try:
2163 return iter(obj)
2164 except TypeError:
2165 return iter((obj,))
2168def adjacent(predicate, iterable, distance=1):
2169 """Return an iterable over `(bool, item)` tuples where the `item` is
2170 drawn from *iterable* and the `bool` indicates whether
2171 that item satisfies the *predicate* or is adjacent to an item that does.
2173 For example, to find whether items are adjacent to a ``3``::
2175 >>> list(adjacent(lambda x: x == 3, range(6)))
2176 [(False, 0), (False, 1), (True, 2), (True, 3), (True, 4), (False, 5)]
2178 Set *distance* to change what counts as adjacent. For example, to find
2179 whether items are two places away from a ``3``:
2181 >>> list(adjacent(lambda x: x == 3, range(6), distance=2))
2182 [(False, 0), (True, 1), (True, 2), (True, 3), (True, 4), (True, 5)]
2184 This is useful for contextualizing the results of a search function.
2185 For example, a code comparison tool might want to identify lines that
2186 have changed, but also surrounding lines to give the viewer of the diff
2187 context.
2189 The predicate function will only be called once for each item in the
2190 iterable.
2192 See also :func:`groupby_transform`, which can be used with this function
2193 to group ranges of items with the same `bool` value.
2195 """
2196 # Allow distance=0 mainly for testing that it reproduces results with map()
2197 if distance < 0:
2198 raise ValueError('distance must be at least 0')
2200 i1, i2 = tee(iterable)
2201 padding = [False] * distance
2202 selected = chain(padding, map(predicate, i1), padding)
2203 adjacent_to_selected = map(any, windowed(selected, 2 * distance + 1))
2204 return zip(adjacent_to_selected, i2)
2207def groupby_transform(iterable, keyfunc=None, valuefunc=None, reducefunc=None):
2208 """An extension of :func:`itertools.groupby` that can apply transformations
2209 to the grouped data.
2211 * *keyfunc* is a function computing a key value for each item in *iterable*
2212 * *valuefunc* is a function that transforms the individual items from
2213 *iterable* after grouping
2214 * *reducefunc* is a function that transforms each group of items
2216 >>> iterable = 'aAAbBBcCC'
2217 >>> keyfunc = lambda k: k.upper()
2218 >>> valuefunc = lambda v: v.lower()
2219 >>> reducefunc = lambda g: ''.join(g)
2220 >>> list(groupby_transform(iterable, keyfunc, valuefunc, reducefunc))
2221 [('A', 'aaa'), ('B', 'bbb'), ('C', 'ccc')]
2223 Each optional argument defaults to an identity function if not specified.
2225 :func:`groupby_transform` is useful when grouping elements of an iterable
2226 using a separate iterable as the key. To do this, :func:`zip` the iterables
2227 and pass a *keyfunc* that extracts the first element and a *valuefunc*
2228 that extracts the second element::
2230 >>> from operator import itemgetter
2231 >>> keys = [0, 0, 1, 1, 1, 2, 2, 2, 3]
2232 >>> values = 'abcdefghi'
2233 >>> iterable = zip(keys, values)
2234 >>> grouper = groupby_transform(iterable, itemgetter(0), itemgetter(1))
2235 >>> [(k, ''.join(g)) for k, g in grouper]
2236 [(0, 'ab'), (1, 'cde'), (2, 'fgh'), (3, 'i')]
2238 Note that the order of items in the iterable is significant.
2239 Only adjacent items are grouped together, so if you don't want any
2240 duplicate groups, you should sort the iterable by the key function
2241 or consider :func:`bucket` or :func:`map_reduce`. :func:`map_reduce`
2242 consumes the iterable immediately and returns a dictionary, while
2243 :func:`bucket` does not.
2245 .. seealso:: :func:`bucket`, :func:`map_reduce`
2247 """
2248 ret = groupby(iterable, keyfunc)
2249 if valuefunc:
2250 ret = ((k, map(valuefunc, g)) for k, g in ret)
2251 if reducefunc:
2252 ret = ((k, reducefunc(g)) for k, g in ret)
2254 return ret
2257class numeric_range(Sequence):
2258 """An extension of the built-in ``range()`` function whose arguments can
2259 be any orderable numeric type.
2261 With only *stop* specified, *start* defaults to ``0`` and *step*
2262 defaults to ``1``. The output items will match the type of *stop*:
2264 >>> list(numeric_range(3.5))
2265 [0.0, 1.0, 2.0, 3.0]
2267 With only *start* and *stop* specified, *step* defaults to ``1``. The
2268 output items will match the type of *start*:
2270 >>> from decimal import Decimal
2271 >>> start = Decimal('2.1')
2272 >>> stop = Decimal('5.1')
2273 >>> list(numeric_range(start, stop))
2274 [Decimal('2.1'), Decimal('3.1'), Decimal('4.1')]
2276 With *start*, *stop*, and *step* specified the output items will match
2277 the type of ``start + step``:
2279 >>> from fractions import Fraction
2280 >>> start = Fraction(1, 2) # Start at 1/2
2281 >>> stop = Fraction(5, 2) # End at 5/2
2282 >>> step = Fraction(1, 2) # Count by 1/2
2283 >>> list(numeric_range(start, stop, step))
2284 [Fraction(1, 2), Fraction(1, 1), Fraction(3, 2), Fraction(2, 1)]
2286 If *step* is zero, ``ValueError`` is raised. Negative steps are supported:
2288 >>> list(numeric_range(3, -1, -1.0))
2289 [3.0, 2.0, 1.0, 0.0]
2291 Be aware of the limitations of floating-point numbers; the representation
2292 of the yielded numbers may be surprising.
2294 ``datetime.datetime`` objects can be used for *start* and *stop*, if *step*
2295 is a ``datetime.timedelta`` object:
2297 >>> import datetime
2298 >>> start = datetime.datetime(2019, 1, 1)
2299 >>> stop = datetime.datetime(2019, 1, 3)
2300 >>> step = datetime.timedelta(days=1)
2301 >>> items = iter(numeric_range(start, stop, step))
2302 >>> next(items)
2303 datetime.datetime(2019, 1, 1, 0, 0)
2304 >>> next(items)
2305 datetime.datetime(2019, 1, 2, 0, 0)
2307 """
2309 _EMPTY_HASH = hash(range(0, 0))
2311 def __init__(self, *args):
2312 argc = len(args)
2313 if argc == 1:
2314 (self._stop,) = args
2315 self._start = type(self._stop)(0)
2316 self._step = type(self._stop - self._start)(1)
2317 elif argc == 2:
2318 self._start, self._stop = args
2319 self._step = type(self._stop - self._start)(1)
2320 elif argc == 3:
2321 self._start, self._stop, self._step = args
2322 elif argc == 0:
2323 raise TypeError(
2324 f'numeric_range expected at least 1 argument, got {argc}'
2325 )
2326 else:
2327 raise TypeError(
2328 f'numeric_range expected at most 3 arguments, got {argc}'
2329 )
2331 self._zero = type(self._step)(0)
2332 if self._step == self._zero:
2333 raise ValueError('numeric_range() arg 3 must not be zero')
2334 self._growing = self._step > self._zero
2336 def __bool__(self):
2337 if self._growing:
2338 return self._start < self._stop
2339 else:
2340 return self._start > self._stop
2342 def __contains__(self, elem):
2343 try:
2344 self.index(elem)
2345 except ValueError:
2346 return False
2347 return True
2349 def __eq__(self, other):
2350 # numeric_range object equality is intended to mirror the built-in range
2351 # object's equality.
2352 # https://github.com/python/cpython/blob/f5c4880151b609e0a0a0b05c292d36b18038c061/Objects/rangeobject.c#L499
2353 if not isinstance(other, numeric_range):
2354 return False
2356 if self is other:
2357 return True
2359 len_self = len(self)
2360 if len_self != len(other):
2361 return False
2363 if not len_self:
2364 return True
2366 if self._start != other._start:
2367 return False
2369 if len_self == 1:
2370 return True
2372 return self._step == other._step
2374 def __getitem__(self, key):
2375 if isinstance(key, int):
2376 return self._get_by_index(key)
2377 elif isinstance(key, slice):
2378 start_idx, stop_idx, step_idx = key.indices(self._len)
2379 return numeric_range(
2380 self._start + start_idx * self._step,
2381 self._start + stop_idx * self._step,
2382 self._step * step_idx,
2383 )
2384 else:
2385 raise TypeError(
2386 'numeric range indices must be '
2387 f'integers or slices, not {type(key).__name__}'
2388 )
2390 def __hash__(self):
2391 # numeric_range hashing is intended to mirror the built-in range object's
2392 # hashing.
2393 # https://github.com/python/cpython/blob/f5c4880151b609e0a0a0b05c292d36b18038c061/Objects/rangeobject.c#L570
2394 len_self = len(self)
2395 if not len_self:
2396 return hash((len_self, None, None))
2397 if len_self == 1:
2398 return hash((len_self, self._start, None))
2399 return hash((len_self, self._start, self._step))
2401 def __iter__(self):
2402 values = (self._start + (n * self._step) for n in count())
2403 if self._growing:
2404 return takewhile(partial(gt, self._stop), values)
2405 else:
2406 return takewhile(partial(lt, self._stop), values)
2408 def __len__(self):
2409 return self._len
2411 @cached_property
2412 def _len(self):
2413 if self._growing:
2414 start = self._start
2415 stop = self._stop
2416 step = self._step
2417 else:
2418 start = self._stop
2419 stop = self._start
2420 step = -self._step
2421 distance = stop - start
2422 if distance <= self._zero:
2423 return 0
2424 else: # distance > 0 and step > 0: regular euclidean division
2425 q, r = divmod(distance, step)
2426 n = int(q) + int(r != self._zero)
2427 # The division above measures the distance to cover, but the items
2428 # are produced by repeated multiplication (see `__iter__`). With
2429 # inexact arithmetic the two disagree at the boundary, so settle
2430 # the count against the items themselves.
2431 while n and not self._before_stop(n - 1):
2432 n -= 1
2433 while self._before_stop(n):
2434 n += 1
2435 return n
2437 def _before_stop(self, i):
2438 value = self._start + i * self._step
2439 return value < self._stop if self._growing else value > self._stop
2441 def __reduce__(self):
2442 return numeric_range, (self._start, self._stop, self._step)
2444 def __repr__(self):
2445 if self._step == 1:
2446 return f"numeric_range({self._start!r}, {self._stop!r})"
2447 return (
2448 f"numeric_range({self._start!r}, {self._stop!r}, {self._step!r})"
2449 )
2451 def __reversed__(self):
2452 # Empty iterator
2453 try:
2454 start = self._get_by_index(-1)
2455 except IndexError:
2456 return iter([])
2458 return iter(
2459 numeric_range(start, self._start - self._step, -self._step)
2460 )
2462 def count(self, value):
2463 return int(value in self)
2465 def index(self, value):
2466 if self._growing:
2467 if self._start <= value < self._stop:
2468 q, _ = divmod(value - self._start, self._step)
2469 return self._index_near(int(q), value)
2470 else:
2471 if self._start >= value > self._stop:
2472 q, _ = divmod(self._start - value, -self._step)
2473 return self._index_near(int(q), value)
2475 raise ValueError(f"{value} is not in numeric range")
2477 def _index_near(self, i, value):
2478 # `i` is the quotient of the division of `value` by the step, which
2479 # locates the value on the grid of items this range produces. For
2480 # exact types that quotient is the index. For inexact ones (floats)
2481 # it can land one short, and the remainder of the division is not a
2482 # reliable membership test either: `numeric_range(0.0, 1.0, 0.1)`
2483 # yields 0.30000000000000004, which leaves a non-zero remainder.
2484 # So compare against the item the range actually produces there.
2485 for candidate in (i, i + 1):
2486 if 0 <= candidate < self._len:
2487 if self._start + candidate * self._step == value:
2488 return candidate
2490 raise ValueError(f"{value} is not in numeric range")
2492 def _get_by_index(self, i):
2493 if i < 0:
2494 i += self._len
2495 if i < 0 or i >= self._len:
2496 raise IndexError("numeric range object index out of range")
2497 return self._start + i * self._step
2500def count_cycle(iterable, n=None):
2501 """Cycle through the items from *iterable* up to *n* times, yielding
2502 the number of completed cycles along with each item. If *n* is omitted the
2503 process repeats indefinitely.
2505 >>> list(count_cycle('AB', 3))
2506 [(0, 'A'), (0, 'B'), (1, 'A'), (1, 'B'), (2, 'A'), (2, 'B')]
2508 """
2509 if n is not None:
2510 return product(range(n), iterable)
2511 seq = tuple(iterable)
2512 if not seq:
2513 return iter(())
2514 return zip(repeat_each(count(), len(seq)), cycle(seq))
2517def mark_ends(iterable):
2518 """Yield 3-tuples of the form ``(is_first, is_last, item)``.
2520 >>> list(mark_ends('ABC'))
2521 [(True, False, 'A'), (False, False, 'B'), (False, True, 'C')]
2523 Use this when looping over an iterable to take special action on its first
2524 and/or last items:
2526 >>> iterable = ['Header', 100, 200, 'Footer']
2527 >>> total = 0
2528 >>> for is_first, is_last, item in mark_ends(iterable):
2529 ... if is_first:
2530 ... continue # Skip the header
2531 ... if is_last:
2532 ... continue # Skip the footer
2533 ... total += item
2534 >>> print(total)
2535 300
2536 """
2537 it = iter(iterable)
2538 for a in it:
2539 first = True
2540 for b in it:
2541 yield first, False, a
2542 a = b
2543 first = False
2544 yield first, True, a
2547def locate(iterable, pred=bool, window_size=None):
2548 """Yield the index of each item in *iterable* for which *pred* returns
2549 ``True``.
2551 *pred* defaults to :func:`bool`, which will select truthy items:
2553 >>> list(locate([0, 1, 1, 0, 1, 0, 0]))
2554 [1, 2, 4]
2556 Set *pred* to a custom function to, e.g., find the indexes for a particular
2557 item.
2559 >>> list(locate(['a', 'b', 'c', 'b'], lambda x: x == 'b'))
2560 [1, 3]
2562 If *window_size* is given, then the *pred* function will be called with
2563 the values in each window. This enables searching for sub-sequences.
2564 Note that *pred* may receive fewer than *window_size* arguments at the end of
2565 the iterable.
2567 >>> iterable = [0, 1, 2, 3, 0, 1, 2, 3, 0, 1, 2, 3]
2568 >>> pred = lambda *args: args == (1, 2, 3)
2569 >>> list(locate(iterable, pred=pred, window_size=3))
2570 [1, 5, 9]
2572 Use with :func:`seekable` to find indexes and then retrieve the associated
2573 items:
2575 >>> from itertools import count
2576 >>> from more_itertools import seekable
2577 >>> source = (3 * n + 1 if (n % 2) else n // 2 for n in count())
2578 >>> it = seekable(source)
2579 >>> pred = lambda x: x > 100
2580 >>> indexes = locate(it, pred=pred)
2581 >>> i = next(indexes)
2582 >>> it.seek(i)
2583 >>> next(it)
2584 106
2586 """
2587 if window_size is None:
2588 return compress(count(), map(pred, iterable))
2590 if window_size < 1:
2591 raise ValueError('window size must be at least 1')
2593 it = windowed(iterable, window_size, fillvalue=_marker)
2594 return compress(
2595 count(),
2596 (pred(*(x for x in w if x is not _marker)) for w in it),
2597 )
2600def longest_common_prefix(iterables):
2601 """Yield elements of the longest common prefix among given *iterables*.
2603 >>> ''.join(longest_common_prefix(['abcd', 'abc', 'abf']))
2604 'ab'
2606 """
2607 return (c[0] for c in takewhile(all_equal, zip(*iterables)))
2610def lstrip(iterable, pred):
2611 """Yield the items from *iterable*, but strip any from the beginning
2612 for which *pred* returns ``True``.
2614 For example, to remove a set of items from the start of an iterable:
2616 >>> iterable = (None, False, None, 1, 2, None, 3, False, None)
2617 >>> pred = lambda x: x in {None, False, ''}
2618 >>> list(lstrip(iterable, pred))
2619 [1, 2, None, 3, False, None]
2621 This function is analogous to :func:`str.lstrip`, and is essentially
2622 a wrapper for :func:`itertools.dropwhile`.
2624 """
2625 return dropwhile(pred, iterable)
2628def rstrip(iterable, pred):
2629 """Yield the items from *iterable*, but strip any from the end
2630 for which *pred* returns ``True``.
2632 For example, to remove a set of items from the end of an iterable:
2634 >>> iterable = (None, False, None, 1, 2, None, 3, False, None)
2635 >>> pred = lambda x: x in {None, False, ''}
2636 >>> list(rstrip(iterable, pred))
2637 [None, False, None, 1, 2, None, 3]
2639 This function is analogous to :func:`str.rstrip`.
2641 """
2642 cache = []
2643 cache_append = cache.append
2644 cache_clear = cache.clear
2645 for x in iterable:
2646 if pred(x):
2647 cache_append(x)
2648 else:
2649 yield from cache
2650 cache_clear()
2651 yield x
2654def strip(iterable, pred):
2655 """Yield the items from *iterable*, but strip any from the
2656 beginning and end for which *pred* returns ``True``.
2658 For example, to remove a set of items from both ends of an iterable:
2660 >>> iterable = (None, False, None, 1, 2, None, 3, False, None)
2661 >>> pred = lambda x: x in {None, False, ''}
2662 >>> list(strip(iterable, pred))
2663 [1, 2, None, 3]
2665 This function is analogous to :func:`str.strip`.
2667 """
2668 return rstrip(lstrip(iterable, pred), pred)
2671class islice_extended:
2672 """An extension of :func:`itertools.islice` that supports negative values
2673 for *stop*, *start*, and *step*.
2675 >>> iterator = iter('abcdefgh')
2676 >>> list(islice_extended(iterator, -4, -1))
2677 ['e', 'f', 'g']
2679 Slices with negative values require some caching of *iterable*, but this
2680 function takes care to minimize the amount of memory required.
2682 For example, you can use a negative step with an infinite iterator:
2684 >>> from itertools import count
2685 >>> list(islice_extended(count(), 110, 99, -2))
2686 [110, 108, 106, 104, 102, 100]
2688 You can also use slice notation directly:
2690 >>> iterator = map(str, count())
2691 >>> it = islice_extended(iterator)[10:20:2]
2692 >>> list(it)
2693 ['10', '12', '14', '16', '18']
2695 """
2697 def __init__(self, iterable, *args):
2698 it = iter(iterable)
2699 if args:
2700 self._iterator = _islice_helper(it, slice(*args))
2701 else:
2702 self._iterator = it
2704 def __iter__(self):
2705 return self
2707 def __next__(self):
2708 return next(self._iterator)
2710 def __getitem__(self, key):
2711 if isinstance(key, slice):
2712 return islice_extended(_islice_helper(self._iterator, key))
2714 raise TypeError('islice_extended.__getitem__ argument must be a slice')
2717def _islice_helper(it, s):
2718 start = s.start
2719 stop = s.stop
2720 if s.step == 0:
2721 raise ValueError('step argument must be a non-zero integer or None.')
2722 step = s.step or 1
2724 if step > 0:
2725 start = 0 if (start is None) else start
2727 if start < 0:
2728 # Consume all but the last -start items
2729 counter = count(1)
2730 wrapper = compress(it, counter)
2731 cache = deque(wrapper, maxlen=-start)
2732 len_iter = next(counter) - 1
2734 # Adjust start to be positive
2735 i = max(len_iter + start, 0)
2737 # Adjust stop to be positive
2738 if stop is None:
2739 j = len_iter
2740 elif stop >= 0:
2741 j = min(stop, len_iter)
2742 else:
2743 j = max(len_iter + stop, 0)
2745 # Slice the cache
2746 n = j - i
2747 if n <= 0:
2748 return
2750 for index in range(n):
2751 if index % step == 0:
2752 # pop and yield the item.
2753 # We don't want to use an intermediate variable
2754 # it would extend the lifetime of the current item
2755 yield cache.popleft()
2756 else:
2757 # just pop and discard the item
2758 cache.popleft()
2759 elif (stop is not None) and (stop < 0):
2760 # Advance to the start position
2761 next(islice(it, start, start), None)
2763 # When stop is negative, we have to carry -stop items while
2764 # iterating
2765 cache = deque(islice(it, -stop), maxlen=-stop)
2767 for index, item in enumerate(it):
2768 if index % step == 0:
2769 # pop and yield the item.
2770 # We don't want to use an intermediate variable
2771 # it would extend the lifetime of the current item
2772 yield cache.popleft()
2773 else:
2774 # just pop and discard the item
2775 cache.popleft()
2776 cache.append(item)
2777 else:
2778 # When both start and stop are positive we have the normal case
2779 yield from islice(it, start, stop, step)
2780 else:
2781 start = -1 if (start is None) else start
2783 if (stop is not None) and (stop < 0):
2784 # Consume all but the last items
2785 n = -stop - 1
2786 counter = count(1)
2787 wrapper = compress(it, counter)
2788 cache = deque(wrapper, maxlen=n)
2789 len_iter = next(counter) - 1
2791 # If start and stop are both negative they are comparable and
2792 # we can just slice. Otherwise we can adjust start to be negative
2793 # and then slice.
2794 if start < 0:
2795 i, j = start, stop
2796 else:
2797 i, j = min(start - len_iter, -1), None
2799 yield from list(cache)[i:j:step]
2800 else:
2801 # Advance to the stop position
2802 if stop is not None:
2803 m = stop + 1
2804 next(islice(it, m, m), None)
2806 # stop is positive, so if start is negative they are not comparable
2807 # and we need the rest of the items.
2808 if start < 0:
2809 i = start
2810 n = None
2811 # stop is None and start is positive, so we just need items up to
2812 # the start index.
2813 elif stop is None:
2814 i = None
2815 n = start + 1
2816 # Both stop and start are positive, so they are comparable.
2817 else:
2818 i = None
2819 n = start - stop
2820 if n <= 0:
2821 return
2823 cache = list(islice(it, n))
2825 yield from cache[i::step]
2828def always_reversible(iterable):
2829 """An extension of :func:`reversed` that supports all iterables, not
2830 just those which implement the ``Reversible`` or ``Sequence`` protocols.
2832 >>> print(*always_reversible(x for x in range(3)))
2833 2 1 0
2835 If the iterable is already reversible, this function returns the
2836 result of :func:`reversed()`. If the iterable is not reversible,
2837 this function will cache the remaining items in the iterable and
2838 yield them in reverse order, which may require significant storage.
2839 """
2840 try:
2841 return reversed(iterable)
2842 except TypeError:
2843 return reversed(list(iterable))
2846def consecutive_groups(iterable, ordering=None):
2847 """Yield groups of consecutive items using :func:`itertools.groupby`.
2848 The *ordering* function determines whether two items are adjacent by
2849 returning their position.
2851 By default, the ordering function is the identity function. This is
2852 suitable for finding runs of numbers:
2854 >>> iterable = [1, 10, 11, 12, 20, 30, 31, 32, 33, 40]
2855 >>> for group in consecutive_groups(iterable):
2856 ... print(list(group))
2857 [1]
2858 [10, 11, 12]
2859 [20]
2860 [30, 31, 32, 33]
2861 [40]
2863 To find runs of adjacent letters, apply :func:`ord` function
2864 to convert letters to ordinals.
2866 >>> iterable = 'abcdfgilmnop'
2867 >>> ordering = ord
2868 >>> for group in consecutive_groups(iterable, ordering):
2869 ... print(list(group))
2870 ['a', 'b', 'c', 'd']
2871 ['f', 'g']
2872 ['i']
2873 ['l', 'm', 'n', 'o', 'p']
2875 Each group of consecutive items is an iterator that shares its source with
2876 *iterable*. When an output group is advanced, the previous group is
2877 no longer available unless its elements are copied (e.g., into a ``list``).
2879 >>> iterable = [1, 2, 11, 12, 21, 22]
2880 >>> saved_groups = []
2881 >>> for group in consecutive_groups(iterable):
2882 ... saved_groups.append(list(group)) # Copy group elements
2883 >>> saved_groups
2884 [[1, 2], [11, 12], [21, 22]]
2886 """
2887 if ordering is None:
2888 key = lambda x: x[0] - x[1]
2889 else:
2890 key = lambda x: x[0] - ordering(x[1])
2892 for k, g in groupby(enumerate(iterable), key=key):
2893 yield map(itemgetter(1), g)
2896def difference(iterable, func=sub, *, initial=None):
2897 """This function is the inverse of :func:`itertools.accumulate`. By default
2898 it will compute the first difference of *iterable* using
2899 :func:`operator.sub`:
2901 >>> from itertools import accumulate
2902 >>> iterable = accumulate([0, 1, 2, 3, 4]) # produces 0, 1, 3, 6, 10
2903 >>> list(difference(iterable))
2904 [0, 1, 2, 3, 4]
2906 *func* defaults to :func:`operator.sub`, but other functions can be
2907 specified. They will be applied as follows::
2909 A, B, C, D, ... --> A, func(B, A), func(C, B), func(D, C), ...
2911 For example, to do progressive division:
2913 >>> iterable = [1, 2, 6, 24, 120]
2914 >>> func = lambda x, y: x // y
2915 >>> list(difference(iterable, func))
2916 [1, 2, 3, 4, 5]
2918 If the *initial* keyword is set, the first element will be skipped when
2919 computing successive differences.
2921 >>> it = [10, 11, 13, 16] # from accumulate([1, 2, 3], initial=10)
2922 >>> list(difference(it, initial=10))
2923 [1, 2, 3]
2925 """
2926 a, b = tee(iterable)
2927 try:
2928 first = [next(b)]
2929 except StopIteration:
2930 return iter([])
2932 if initial is not None:
2933 return map(func, b, a)
2935 return chain(first, map(func, b, a))
2938class SequenceView(Sequence):
2939 """Return a read-only view of the sequence object *target*.
2941 :class:`SequenceView` objects are analogous to Python's built-in
2942 "dictionary view" types. They provide a dynamic view of a sequence's items,
2943 meaning that when the sequence updates, so does the view.
2945 >>> seq = ['0', '1', '2']
2946 >>> view = SequenceView(seq)
2947 >>> view
2948 SequenceView(['0', '1', '2'])
2949 >>> seq.append('3')
2950 >>> view
2951 SequenceView(['0', '1', '2', '3'])
2953 Sequence views support indexing, slicing, and length queries. They act
2954 like the underlying sequence, except they don't allow assignment:
2956 >>> view[1]
2957 '1'
2958 >>> view[1:-1]
2959 ['1', '2']
2960 >>> len(view)
2961 4
2963 Sequence views are useful as an alternative to copying, as they don't
2964 require (much) extra storage.
2966 """
2968 def __init__(self, target):
2969 if not isinstance(target, Sequence):
2970 raise TypeError
2971 self._target = target
2973 def __getitem__(self, index):
2974 return self._target[index]
2976 def __len__(self):
2977 return len(self._target)
2979 def __repr__(self):
2980 return f'{self.__class__.__name__}({self._target!r})'
2983class seekable:
2984 """Wrap an iterator to allow for seeking backward and forward. This
2985 progressively caches the items in the source iterable so they can be
2986 re-visited.
2988 Call :meth:`seek` with an index to seek to that position in the source
2989 iterable.
2991 To "reset" an iterator, seek to ``0``:
2993 >>> from itertools import count
2994 >>> it = seekable((str(n) for n in count()))
2995 >>> next(it), next(it), next(it)
2996 ('0', '1', '2')
2997 >>> it.seek(0)
2998 >>> next(it), next(it), next(it)
2999 ('0', '1', '2')
3001 You can also seek forward:
3003 >>> it = seekable((str(n) for n in range(20)))
3004 >>> it.seek(10)
3005 >>> next(it)
3006 '10'
3007 >>> it.seek(20) # Seeking past the end of the source isn't a problem
3008 >>> list(it)
3009 []
3010 >>> it.seek(0) # Resetting works even after hitting the end
3011 >>> next(it)
3012 '0'
3014 Call :meth:`relative_seek` to seek relative to the source iterator's
3015 current position.
3017 >>> it = seekable((str(n) for n in range(20)))
3018 >>> next(it), next(it), next(it)
3019 ('0', '1', '2')
3020 >>> it.relative_seek(2)
3021 >>> next(it)
3022 '5'
3023 >>> it.relative_seek(-3) # Source is at '6', we move back to '3'
3024 >>> next(it)
3025 '3'
3026 >>> it.relative_seek(-3) # Source is at '4', we move back to '1'
3027 >>> next(it)
3028 '1'
3031 Call :meth:`peek` to look ahead one item without advancing the iterator:
3033 >>> it = seekable('1234')
3034 >>> it.peek()
3035 '1'
3036 >>> list(it)
3037 ['1', '2', '3', '4']
3038 >>> it.peek(default='empty')
3039 'empty'
3041 Before the iterator is at its end, calling :func:`bool` on it will return
3042 ``True``. After it will return ``False``:
3044 >>> it = seekable('5678')
3045 >>> bool(it)
3046 True
3047 >>> list(it)
3048 ['5', '6', '7', '8']
3049 >>> bool(it)
3050 False
3052 You may view the contents of the cache with the :meth:`elements` method.
3053 That returns a :class:`SequenceView`, a view that updates automatically:
3055 >>> it = seekable((str(n) for n in range(10)))
3056 >>> next(it), next(it), next(it)
3057 ('0', '1', '2')
3058 >>> elements = it.elements()
3059 >>> elements
3060 SequenceView(['0', '1', '2'])
3061 >>> next(it)
3062 '3'
3063 >>> elements
3064 SequenceView(['0', '1', '2', '3'])
3066 Indexing the :class:`seekable` directly returns items from the cache:
3068 >>> it = seekable((str(n) for n in range(10)))
3069 >>> next(it), next(it), next(it)
3070 ('0', '1', '2')
3071 >>> it[-1]
3072 '2'
3073 >>> it[0]
3074 '0'
3076 By default, the cache grows as the source iterable progresses, so beware of
3077 wrapping very large or infinite iterables. Supply *maxlen* to limit the
3078 size of the cache (this of course limits how far back you can seek).
3080 >>> from itertools import count
3081 >>> it = seekable((str(n) for n in count()), maxlen=2)
3082 >>> next(it), next(it), next(it), next(it)
3083 ('0', '1', '2', '3')
3084 >>> list(it.elements())
3085 ['2', '3']
3086 >>> it.seek(0)
3087 >>> next(it), next(it), next(it), next(it)
3088 ('2', '3', '4', '5')
3089 >>> next(it)
3090 '6'
3092 """
3094 def __init__(self, iterable, maxlen=None):
3095 self._source = iter(iterable)
3096 if maxlen is None:
3097 self._cache = []
3098 else:
3099 self._cache = deque([], maxlen)
3100 self._index = None
3102 def __iter__(self):
3103 return self
3105 def __next__(self):
3106 if self._index is not None:
3107 try:
3108 item = self._cache[self._index]
3109 except IndexError:
3110 self._index = None
3111 else:
3112 self._index += 1
3113 return item
3115 item = next(self._source)
3116 self._cache.append(item)
3117 return item
3119 def __bool__(self):
3120 try:
3121 self.peek()
3122 except StopIteration:
3123 return False
3124 return True
3126 def peek(self, default=_marker):
3127 try:
3128 peeked = next(self)
3129 except StopIteration:
3130 if default is _marker:
3131 raise
3132 return default
3133 if self._index is None:
3134 self._index = len(self._cache)
3135 self._index -= 1
3136 return peeked
3138 def elements(self):
3139 return SequenceView(self._cache)
3141 def seek(self, index):
3142 self._index = index
3143 remainder = index - len(self._cache)
3144 if remainder > 0:
3145 consume(self, remainder)
3147 def relative_seek(self, count):
3148 if self._index is None:
3149 self._index = len(self._cache)
3151 self.seek(max(self._index + count, 0))
3153 def __getitem__(self, index):
3154 return self._cache[index]
3157class run_length:
3158 """
3159 :func:`run_length.encode` compresses an iterable with run-length encoding.
3160 It yields groups of repeated items with the count of how many times they
3161 were repeated:
3163 >>> uncompressed = 'abbcccdddd'
3164 >>> list(run_length.encode(uncompressed))
3165 [('a', 1), ('b', 2), ('c', 3), ('d', 4)]
3167 :func:`run_length.decode` decompresses an iterable that was previously
3168 compressed with run-length encoding. It yields the items of the
3169 decompressed iterable:
3171 >>> compressed = [('a', 1), ('b', 2), ('c', 3), ('d', 4)]
3172 >>> list(run_length.decode(compressed))
3173 ['a', 'b', 'b', 'c', 'c', 'c', 'd', 'd', 'd', 'd']
3175 """
3177 @staticmethod
3178 def encode(iterable):
3179 return ((k, ilen(g)) for k, g in groupby(iterable))
3181 @staticmethod
3182 def decode(iterable):
3183 return chain.from_iterable(starmap(repeat, iterable))
3186def exactly_n(iterable, n, predicate=bool):
3187 """Return ``True`` if exactly ``n`` items in the iterable are ``True``
3188 according to the *predicate* function.
3190 >>> exactly_n([True, True, False], 2)
3191 True
3192 >>> exactly_n([True, True, False], 1)
3193 False
3194 >>> exactly_n([0, 1, 2, 3, 4, 5], 3, lambda x: x < 3)
3195 True
3197 The iterable will be advanced until ``n + 1`` truthy items are encountered,
3198 so avoid calling it on infinite iterables.
3200 """
3201 iterator = filter(predicate, iterable)
3202 if n <= 0:
3203 if n < 0:
3204 return False
3205 for _ in iterator:
3206 return False
3207 return True
3209 iterator = islice(iterator, n - 1, None)
3210 for _ in iterator:
3211 for _ in iterator:
3212 return False
3213 return True
3214 return False
3217def circular_shifts(iterable, steps=1):
3218 """Yield the circular shifts of *iterable*.
3220 >>> list(circular_shifts(range(4)))
3221 [(0, 1, 2, 3), (1, 2, 3, 0), (2, 3, 0, 1), (3, 0, 1, 2)]
3223 Set *steps* to the number of places to rotate to the left
3224 (or to the right if negative). Defaults to 1.
3226 >>> list(circular_shifts(range(4), 2))
3227 [(0, 1, 2, 3), (2, 3, 0, 1)]
3229 >>> list(circular_shifts(range(4), -1))
3230 [(0, 1, 2, 3), (3, 0, 1, 2), (2, 3, 0, 1), (1, 2, 3, 0)]
3232 """
3233 buffer = deque(iterable)
3234 if steps == 0:
3235 raise ValueError('Steps should be a non-zero integer')
3237 buffer.rotate(steps)
3238 steps = -steps
3239 n = len(buffer)
3240 n //= math.gcd(n, steps)
3242 for _ in repeat(None, n):
3243 buffer.rotate(steps)
3244 yield tuple(buffer)
3247def make_decorator(wrapping_func, result_index=0):
3248 """Return a decorator version of *wrapping_func*, which is a function that
3249 modifies an iterable. *result_index* is the position in that function's
3250 signature where the iterable goes.
3252 This lets you use itertools on the "production end," i.e. at function
3253 definition. This can augment what the function returns without changing the
3254 function's code.
3256 For example, to produce a decorator version of :func:`chunked`:
3258 >>> from more_itertools import chunked
3259 >>> chunker = make_decorator(chunked, result_index=0)
3260 >>> @chunker(3)
3261 ... def iter_range(n):
3262 ... return iter(range(n))
3263 ...
3264 >>> list(iter_range(9))
3265 [[0, 1, 2], [3, 4, 5], [6, 7, 8]]
3267 To only allow truthy items to be returned:
3269 >>> truth_serum = make_decorator(filter, result_index=1)
3270 >>> @truth_serum(bool)
3271 ... def boolean_test():
3272 ... return [0, 1, '', ' ', False, True]
3273 ...
3274 >>> list(boolean_test())
3275 [1, ' ', True]
3277 The :func:`peekable` and :func:`seekable` wrappers make for practical
3278 decorators:
3280 >>> from more_itertools import peekable
3281 >>> peekable_function = make_decorator(peekable)
3282 >>> @peekable_function()
3283 ... def str_range(*args):
3284 ... return (str(x) for x in range(*args))
3285 ...
3286 >>> it = str_range(1, 20, 2)
3287 >>> next(it), next(it), next(it)
3288 ('1', '3', '5')
3289 >>> it.peek()
3290 '7'
3291 >>> next(it)
3292 '7'
3294 """
3296 # See https://sites.google.com/site/bbayles/index/decorator_factory for
3297 # notes on how this works.
3298 def decorator(*wrapping_args, **wrapping_kwargs):
3299 def outer_wrapper(f):
3300 def inner_wrapper(*args, **kwargs):
3301 result = f(*args, **kwargs)
3302 wrapping_args_ = list(wrapping_args)
3303 wrapping_args_.insert(result_index, result)
3304 return wrapping_func(*wrapping_args_, **wrapping_kwargs)
3306 return inner_wrapper
3308 return outer_wrapper
3310 return decorator
3313def map_reduce(iterable, keyfunc, valuefunc=None, reducefunc=None):
3314 """Return a dictionary that maps the items in *iterable* to categories
3315 defined by *keyfunc*, transforms them with *valuefunc*, and
3316 then summarizes them by category with *reducefunc*.
3318 *valuefunc* defaults to the identity function if it is unspecified.
3319 If *reducefunc* is unspecified, no summarization takes place:
3321 >>> keyfunc = lambda x: x.upper()
3322 >>> result = map_reduce('abbccc', keyfunc)
3323 >>> sorted(result.items())
3324 [('A', ['a']), ('B', ['b', 'b']), ('C', ['c', 'c', 'c'])]
3326 Specifying *valuefunc* transforms the categorized items:
3328 >>> keyfunc = lambda x: x.upper()
3329 >>> valuefunc = lambda x: 1
3330 >>> result = map_reduce('abbccc', keyfunc, valuefunc)
3331 >>> sorted(result.items())
3332 [('A', [1]), ('B', [1, 1]), ('C', [1, 1, 1])]
3334 Specifying *reducefunc* summarizes the categorized items:
3336 >>> keyfunc = lambda x: x.upper()
3337 >>> valuefunc = lambda x: 1
3338 >>> reducefunc = sum
3339 >>> result = map_reduce('abbccc', keyfunc, valuefunc, reducefunc)
3340 >>> sorted(result.items())
3341 [('A', 1), ('B', 2), ('C', 3)]
3343 You may want to filter the input iterable before applying the map/reduce
3344 procedure:
3346 >>> all_items = range(30)
3347 >>> items = [x for x in all_items if 10 <= x <= 20] # Filter
3348 >>> keyfunc = lambda x: x % 2 # Evens map to 0; odds to 1
3349 >>> categories = map_reduce(items, keyfunc=keyfunc)
3350 >>> sorted(categories.items())
3351 [(0, [10, 12, 14, 16, 18, 20]), (1, [11, 13, 15, 17, 19])]
3352 >>> summaries = map_reduce(items, keyfunc=keyfunc, reducefunc=sum)
3353 >>> sorted(summaries.items())
3354 [(0, 90), (1, 75)]
3356 Note that all items in the iterable are gathered into a list before the
3357 summarization step, which may require significant storage.
3359 The returned object is a :obj:`collections.defaultdict` with the
3360 ``default_factory`` set to ``None``, such that it behaves like a normal
3361 dictionary.
3363 .. seealso:: :func:`bucket`, :func:`groupby_transform`
3365 If storage is a concern, :func:`bucket` can be used without consuming the
3366 entire iterable right away. If the elements with the same key are already
3367 adjacent, :func:`groupby_transform` or :func:`itertools.groupby` can be
3368 used without any caching overhead.
3370 """
3372 ret = defaultdict(list)
3374 if valuefunc is None:
3375 for item in iterable:
3376 key = keyfunc(item)
3377 ret[key].append(item)
3379 else:
3380 for item in iterable:
3381 key = keyfunc(item)
3382 value = valuefunc(item)
3383 ret[key].append(value)
3385 if reducefunc is not None:
3386 for key, value_list in ret.items():
3387 ret[key] = reducefunc(value_list)
3389 ret.default_factory = None
3390 return ret
3393def rlocate(iterable, pred=bool, window_size=None):
3394 """Yield the index of each item in *iterable* for which *pred* returns
3395 ``True``, starting from the right and moving left.
3397 *pred* defaults to :func:`bool`, which will select truthy items:
3399 >>> list(rlocate([0, 1, 1, 0, 1, 0, 0])) # Truthy at 1, 2, and 4
3400 [4, 2, 1]
3402 Set *pred* to a custom function to, e.g., find the indexes for a particular
3403 item:
3405 >>> iterator = iter('abcb')
3406 >>> pred = lambda x: x == 'b'
3407 >>> list(rlocate(iterator, pred))
3408 [3, 1]
3410 If *window_size* is given, then the *pred* function will be called with
3411 that many items. This enables searching for sub-sequences:
3413 >>> iterable = [0, 1, 2, 3, 0, 1, 2, 3, 0, 1, 2, 3]
3414 >>> pred = lambda *args: args == (1, 2, 3)
3415 >>> list(rlocate(iterable, pred=pred, window_size=3))
3416 [9, 5, 1]
3418 Beware, this function won't return anything for infinite iterables.
3419 If *iterable* is reversible, ``rlocate`` will reverse it and search from
3420 the right. Otherwise, it will search from the left and return the results
3421 in reverse order.
3423 See :func:`locate` to for other example applications.
3425 """
3426 if window_size is None:
3427 try:
3428 len_iter = len(iterable)
3429 return (len_iter - i - 1 for i in locate(reversed(iterable), pred))
3430 except TypeError:
3431 pass
3433 return reversed(list(locate(iterable, pred, window_size)))
3436def replace(iterable, pred, substitutes, count=None, window_size=1):
3437 """Yield the items from *iterable*, replacing the items for which *pred*
3438 returns ``True`` with the items from the iterable *substitutes*.
3440 >>> iterable = [1, 1, 0, 1, 1, 0, 1, 1]
3441 >>> pred = lambda x: x == 0
3442 >>> substitutes = (2, 3)
3443 >>> list(replace(iterable, pred, substitutes))
3444 [1, 1, 2, 3, 1, 1, 2, 3, 1, 1]
3446 If *count* is given, the number of replacements will be limited:
3448 >>> iterable = [1, 1, 0, 1, 1, 0, 1, 1, 0]
3449 >>> pred = lambda x: x == 0
3450 >>> substitutes = [None]
3451 >>> list(replace(iterable, pred, substitutes, count=2))
3452 [1, 1, None, 1, 1, None, 1, 1, 0]
3454 Use *window_size* to control the number of items passed as arguments to
3455 *pred*. This allows for locating and replacing subsequences.
3457 >>> iterable = [0, 1, 2, 5, 0, 1, 2, 5]
3458 >>> window_size = 3
3459 >>> pred = lambda *args: args == (0, 1, 2) # 3 items passed to pred
3460 >>> substitutes = [3, 4] # Splice in these items
3461 >>> list(replace(iterable, pred, substitutes, window_size=window_size))
3462 [3, 4, 5, 3, 4, 5]
3464 *pred* may receive fewer than *window_size* arguments at the end of
3465 the iterable and should be able to handle this.
3467 """
3468 if window_size < 1:
3469 raise ValueError('window_size must be at least 1')
3471 # Save the substitutes iterable, since it's used more than once
3472 substitutes = tuple(substitutes)
3474 # Add padding such that the number of windows matches the length of the
3475 # iterable
3476 it = chain(iterable, repeat(_marker, window_size - 1))
3477 windows = windowed(it, window_size)
3479 n = 0
3480 for w in windows:
3481 # Strip any _marker padding so pred never sees internal sentinels.
3482 # Near the end of the iterable, pred will receive fewer arguments.
3483 args = tuple(x for x in w if x is not _marker)
3485 # If the current window matches our predicate (and we haven't hit
3486 # our maximum number of replacements), splice in the substitutes
3487 # and then consume the following windows that overlap with this one.
3488 # For example, if the iterable is (0, 1, 2, 3, 4...)
3489 # and the window size is 2, we have (0, 1), (1, 2), (2, 3)...
3490 # If the predicate matches on (0, 1), we need to zap (0, 1) and (1, 2)
3491 if args and pred(*args):
3492 if (count is None) or (n < count):
3493 n += 1
3494 yield from substitutes
3495 consume(windows, window_size - 1)
3496 continue
3498 # If there was no match (or we've reached the replacement limit),
3499 # yield the first item from the window.
3500 if args:
3501 yield args[0]
3504def partitions(iterable):
3505 """Yield all possible order-preserving partitions of *iterable*.
3507 >>> iterable = 'abc'
3508 >>> for part in partitions(iterable):
3509 ... print([''.join(p) for p in part])
3510 ['abc']
3511 ['a', 'bc']
3512 ['ab', 'c']
3513 ['a', 'b', 'c']
3515 This is unrelated to :func:`partition`.
3517 """
3518 sequence = list(iterable)
3519 n = len(sequence)
3520 for i in powerset(range(1, n)):
3521 yield [sequence[i:j] for i, j in zip((0,) + i, i + (n,))]
3524def set_partitions(iterable, k=None, min_size=None, max_size=None):
3525 """
3526 Yield the set partitions of *iterable* into *k* parts. Set partitions are
3527 not order-preserving.
3529 >>> iterable = 'abc'
3530 >>> for part in set_partitions(iterable, 2):
3531 ... print([''.join(p) for p in part])
3532 ['a', 'bc']
3533 ['ab', 'c']
3534 ['b', 'ac']
3537 If *k* is not given, every set partition is generated.
3539 >>> iterable = 'abc'
3540 >>> for part in set_partitions(iterable):
3541 ... print([''.join(p) for p in part])
3542 ['abc']
3543 ['a', 'bc']
3544 ['ab', 'c']
3545 ['b', 'ac']
3546 ['a', 'b', 'c']
3548 if *min_size* and/or *max_size* are given, the minimum and/or maximum size
3549 per block in partition is set.
3551 >>> iterable = 'abc'
3552 >>> for part in set_partitions(iterable, min_size=2):
3553 ... print([''.join(p) for p in part])
3554 ['abc']
3555 >>> for part in set_partitions(iterable, max_size=2):
3556 ... print([''.join(p) for p in part])
3557 ['a', 'bc']
3558 ['ab', 'c']
3559 ['b', 'ac']
3560 ['a', 'b', 'c']
3562 """
3563 L = list(iterable)
3564 n = len(L)
3565 if k is not None:
3566 if k < 1:
3567 raise ValueError(
3568 "Can't partition in a negative or zero number of groups"
3569 )
3570 elif k > n:
3571 return
3573 min_size = min_size if min_size is not None else 0
3574 max_size = max_size if max_size is not None else n
3575 if min_size > max_size:
3576 return
3578 def set_partitions_helper(L, k):
3579 n = len(L)
3580 if k == 1:
3581 yield [L]
3582 elif n == k:
3583 yield [[s] for s in L]
3584 else:
3585 e, *M = L
3586 for p in set_partitions_helper(M, k - 1):
3587 yield [[e], *p]
3588 for p in set_partitions_helper(M, k):
3589 for i in range(len(p)):
3590 yield p[:i] + [[e] + p[i]] + p[i + 1 :]
3592 if k is None:
3593 for k in range(1, n + 1):
3594 yield from filter(
3595 lambda z: all(min_size <= len(bk) <= max_size for bk in z),
3596 set_partitions_helper(L, k),
3597 )
3598 else:
3599 yield from filter(
3600 lambda z: all(min_size <= len(bk) <= max_size for bk in z),
3601 set_partitions_helper(L, k),
3602 )
3605class time_limited:
3606 """
3607 Yield items from *iterable* until *limit_seconds* have passed.
3608 If the time limit expires before all items have been yielded, the
3609 ``timed_out`` parameter will be set to ``True``.
3611 >>> from time import sleep
3612 >>> def generator():
3613 ... yield 1
3614 ... yield 2
3615 ... sleep(0.2)
3616 ... yield 3
3617 >>> iterable = time_limited(0.1, generator())
3618 >>> list(iterable)
3619 [1, 2]
3620 >>> iterable.timed_out
3621 True
3623 Note that the time is checked before each item is yielded, and iteration
3624 stops if the time elapsed is greater than *limit_seconds*. If your time
3625 limit is 1 second, but it takes 2 seconds to generate the first item from
3626 the iterable, the function will run for 2 seconds and not yield anything.
3627 As a special case, when *limit_seconds* is zero, the iterator never
3628 returns anything.
3630 """
3632 def __init__(self, limit_seconds, iterable):
3633 if limit_seconds < 0:
3634 raise ValueError('limit_seconds must be positive')
3635 self.limit_seconds = limit_seconds
3636 self._iterator = iter(iterable)
3637 self._start_time = monotonic()
3638 self.timed_out = False
3640 def __iter__(self):
3641 return self
3643 def __next__(self):
3644 if self.limit_seconds == 0:
3645 self.timed_out = True
3646 raise StopIteration
3647 item = next(self._iterator)
3648 if monotonic() - self._start_time > self.limit_seconds:
3649 self.timed_out = True
3650 raise StopIteration
3652 return item
3655def only(iterable, default=None, too_long=None):
3656 """If *iterable* has only one item, return it.
3657 If it has zero items, return *default*.
3658 If it has more than one item, raise the exception given by *too_long*,
3659 which is ``ValueError`` by default.
3661 >>> only([], default='missing')
3662 'missing'
3663 >>> only([1])
3664 1
3665 >>> only([1, 2]) # doctest: +IGNORE_EXCEPTION_DETAIL
3666 Traceback (most recent call last):
3667 ...
3668 ValueError: Expected exactly one item in iterable, but got 1, 2,
3669 and perhaps more.'
3670 >>> only([1, 2], too_long=TypeError) # doctest: +IGNORE_EXCEPTION_DETAIL
3671 Traceback (most recent call last):
3672 ...
3673 TypeError
3675 Note that :func:`only` attempts to advance *iterable* twice to ensure there
3676 is only one item. See :func:`spy` or :func:`peekable` to check
3677 iterable contents less destructively.
3679 """
3680 iterator = iter(iterable)
3681 for first in iterator:
3682 for second in iterator:
3683 msg = (
3684 f'Expected exactly one item in iterable, but got {first!r}, '
3685 f'{second!r}, and perhaps more.'
3686 )
3687 raise too_long or ValueError(msg)
3688 return first
3689 return default
3692def ichunked(iterable, n):
3693 """Break *iterable* into sub-iterables with *n* elements each.
3694 :func:`ichunked` is like :func:`chunked`, but it yields iterables
3695 instead of lists.
3697 If the sub-iterables are read in order, the elements of *iterable*
3698 won't be stored in memory.
3699 If they are read out of order, :func:`itertools.tee` is used to cache
3700 elements as necessary.
3702 >>> from itertools import count
3703 >>> all_chunks = ichunked(count(), 4)
3704 >>> c_1, c_2, c_3 = next(all_chunks), next(all_chunks), next(all_chunks)
3705 >>> list(c_2) # c_1's elements have been cached; c_3's haven't been
3706 [4, 5, 6, 7]
3707 >>> list(c_1)
3708 [0, 1, 2, 3]
3709 >>> list(c_3)
3710 [8, 9, 10, 11]
3712 """
3713 iterator = iter(iterable)
3714 for first in iterator:
3715 rest = islice(iterator, n - 1)
3716 cache, cacher = tee(rest)
3717 yield chain([first], rest, cache)
3718 consume(cacher)
3721def iequals(*iterables):
3722 """Return ``True`` if all given *iterables* are equal to each other,
3723 which means that they contain the same elements in the same order.
3725 The function is useful for comparing iterables of different data types
3726 or iterables that do not support equality checks.
3728 >>> iequals("abc", ['a', 'b', 'c'], ('a', 'b', 'c'), iter("abc"))
3729 True
3731 >>> iequals("abc", "acb")
3732 False
3734 Not to be confused with :func:`all_equal`, which checks whether all
3735 elements of iterable are equal to each other.
3737 """
3738 try:
3739 return all(map(all_equal, zip(*iterables, strict=True)))
3740 except ValueError:
3741 return False
3744def distinct_combinations(iterable, r):
3745 """Yield the distinct combinations of *r* items taken from *iterable*.
3747 >>> list(distinct_combinations([0, 0, 1], 2))
3748 [(0, 0), (0, 1)]
3750 Equivalent to ``set(combinations(iterable))``, except duplicates are not
3751 generated and thrown away. For larger input sequences this is much more
3752 efficient.
3754 """
3755 if r < 0:
3756 raise ValueError('r must be non-negative')
3757 elif r == 0:
3758 yield ()
3759 return
3760 pool = tuple(iterable)
3761 generators = [unique_everseen(enumerate(pool), key=itemgetter(1))]
3762 current_combo = [None] * r
3763 level = 0
3764 while generators:
3765 try:
3766 cur_idx, p = next(generators[-1])
3767 except StopIteration:
3768 generators.pop()
3769 level -= 1
3770 continue
3771 current_combo[level] = p
3772 if level + 1 == r:
3773 yield tuple(current_combo)
3774 else:
3775 generators.append(
3776 unique_everseen(
3777 enumerate(pool[cur_idx + 1 :], cur_idx + 1),
3778 key=itemgetter(1),
3779 )
3780 )
3781 level += 1
3784def filter_except(validator, iterable, *exceptions):
3785 """Yield the items from *iterable* for which the *validator* function does
3786 not raise one of the specified *exceptions*.
3788 *validator* is called for each item in *iterable*.
3789 It should be a function that accepts one argument and raises an exception
3790 if that item is not valid.
3792 >>> iterable = ['1', '2', 'three', '4', None]
3793 >>> list(filter_except(int, iterable, ValueError, TypeError))
3794 ['1', '2', '4']
3796 If an exception other than one given by *exceptions* is raised by
3797 *validator*, it is raised like normal.
3798 """
3799 for item in iterable:
3800 try:
3801 validator(item)
3802 except exceptions:
3803 pass
3804 else:
3805 yield item
3808def map_except(function, iterable, *exceptions):
3809 """Transform each item from *iterable* with *function* and yield the
3810 result, unless *function* raises one of the specified *exceptions*.
3812 *function* is called to transform each item in *iterable*.
3813 It should accept one argument.
3815 >>> iterable = ['1', '2', 'three', '4', None]
3816 >>> list(map_except(int, iterable, ValueError, TypeError))
3817 [1, 2, 4]
3819 If an exception other than one given by *exceptions* is raised by
3820 *function*, it is raised like normal.
3821 """
3822 for item in iterable:
3823 try:
3824 yield function(item)
3825 except exceptions:
3826 pass
3829def map_if(iterable, pred, func, func_else=None):
3830 """Evaluate each item from *iterable* using *pred*. If the result is
3831 equivalent to ``True``, transform the item with *func* and yield it.
3832 Otherwise, transform the item with *func_else* and yield it.
3834 *pred*, *func*, and *func_else* should each be functions that accept
3835 one argument. By default, *func_else* is the identity function.
3837 >>> from math import sqrt
3838 >>> iterable = list(range(-5, 5))
3839 >>> iterable
3840 [-5, -4, -3, -2, -1, 0, 1, 2, 3, 4]
3841 >>> list(map_if(iterable, lambda x: x > 3, lambda x: 'toobig'))
3842 [-5, -4, -3, -2, -1, 0, 1, 2, 3, 'toobig']
3843 >>> list(map_if(iterable, lambda x: x >= 0,
3844 ... lambda x: f'{sqrt(x):.2f}', lambda x: None))
3845 [None, None, None, None, None, '0.00', '1.00', '1.41', '1.73', '2.00']
3846 """
3848 if func_else is None:
3849 for item in iterable:
3850 yield func(item) if pred(item) else item
3852 else:
3853 for item in iterable:
3854 yield func(item) if pred(item) else func_else(item)
3857def _sample_unweighted(iterator, k, strict):
3858 # Algorithm L in the 1994 paper by Kim-Hung Li:
3859 # "Reservoir-Sampling Algorithms of Time Complexity O(n(1+log(N/n)))".
3861 reservoir = list(islice(iterator, k))
3862 if strict and len(reservoir) < k:
3863 raise ValueError('Sample larger than population')
3864 W = 1.0
3866 with suppress(StopIteration):
3867 while True:
3868 W *= random() ** (1 / k)
3869 skip = floor(log(random()) / log1p(-W))
3870 element = next(islice(iterator, skip, None))
3871 reservoir[randrange(k)] = element
3873 shuffle(reservoir)
3874 return reservoir
3877def _sample_weighted(iterator, k, weights, strict):
3878 # Implementation of "A-ExpJ" from the 2006 paper by Efraimidis et al. :
3879 # "Weighted random sampling with a reservoir".
3881 # Log-transform for numerical stability for weights that are small/large
3882 weight_keys = (log(random()) / weight for weight in weights)
3884 # Fill up the reservoir (collection of samples) with the first `k`
3885 # weight-keys and elements, then heapify the list.
3886 reservoir = take(k, zip(weight_keys, iterator))
3887 if strict and len(reservoir) < k:
3888 raise ValueError('Sample larger than population')
3890 heapify(reservoir)
3892 # The number of jumps before changing the reservoir is a random variable
3893 # with an exponential distribution. Sample it using random() and logs.
3894 smallest_weight_key, _ = reservoir[0]
3895 weights_to_skip = log(random()) / smallest_weight_key
3897 for weight, element in zip(weights, iterator):
3898 if weight >= weights_to_skip:
3899 # The notation here is consistent with the paper, but we store
3900 # the weight-keys in log-space for better numerical stability.
3901 smallest_weight_key, _ = reservoir[0]
3902 t_w = exp(weight * smallest_weight_key)
3903 r_2 = uniform(t_w, 1) # generate U(t_w, 1)
3904 weight_key = log(r_2) / weight
3905 heapreplace(reservoir, (weight_key, element))
3906 smallest_weight_key, _ = reservoir[0]
3907 weights_to_skip = log(random()) / smallest_weight_key
3908 else:
3909 weights_to_skip -= weight
3911 ret = [element for weight_key, element in reservoir]
3912 shuffle(ret)
3913 return ret
3916def _sample_counted(population, k, counts, strict):
3917 element = None
3918 remaining = 0
3920 def feed(i):
3921 # Advance *i* steps ahead and consume an element
3922 nonlocal element, remaining
3924 while i + 1 > remaining:
3925 i = i - remaining
3926 element = next(population)
3927 remaining = next(counts)
3928 remaining -= i + 1
3929 return element
3931 with suppress(StopIteration):
3932 reservoir = []
3933 for _ in range(k):
3934 reservoir.append(feed(0))
3936 if strict and len(reservoir) < k:
3937 raise ValueError('Sample larger than population')
3939 with suppress(StopIteration):
3940 W = 1.0
3941 while True:
3942 W *= random() ** (1 / k)
3943 skip = floor(log(random()) / log1p(-W))
3944 element = feed(skip)
3945 reservoir[randrange(k)] = element
3947 shuffle(reservoir)
3948 return reservoir
3951def sample(iterable, k, weights=None, *, counts=None, strict=False):
3952 """Return a *k*-length list of elements chosen (without replacement)
3953 from the *iterable*.
3955 Similar to :func:`random.sample`, but works on inputs that aren't
3956 indexable (such as sets and dictionaries) and on inputs where the
3957 size isn't known in advance (such as generators).
3959 >>> iterable = range(100)
3960 >>> sample(iterable, 5) # doctest: +SKIP
3961 [81, 60, 96, 16, 4]
3963 For iterables with repeated elements, you may supply *counts* to
3964 indicate the repeats.
3966 >>> iterable = ['a', 'b']
3967 >>> counts = [3, 4] # Equivalent to 'a', 'a', 'a', 'b', 'b', 'b', 'b'
3968 >>> sample(iterable, k=3, counts=counts) # doctest: +SKIP
3969 ['a', 'a', 'b']
3971 An iterable with *weights* may be given:
3973 >>> iterable = range(100)
3974 >>> weights = (i * i + 1 for i in range(100))
3975 >>> sampled = sample(iterable, 5, weights=weights) # doctest: +SKIP
3976 [79, 67, 74, 66, 78]
3978 Weighted selections are made without replacement.
3979 After an element is selected, it is removed from the pool and the
3980 relative weights of the other elements increase (this
3981 does not match the behavior of :func:`random.sample`'s *counts*
3982 parameter). Note that *weights* may not be used with *counts*.
3984 If the length of *iterable* is less than *k*,
3985 ``ValueError`` is raised if *strict* is ``True`` and
3986 all elements are returned (in shuffled order) if *strict* is ``False``.
3988 By default, the `Algorithm L <https://w.wiki/ANrM>`__ reservoir sampling
3989 technique is used. When *weights* are provided,
3990 `Algorithm A-ExpJ <https://w.wiki/ANrS>`__ is used instead.
3992 Notes on reproducibility:
3994 * The algorithms rely on inexact floating-point functions provided
3995 by the underlying math library (e.g. ``log``, ``log1p``, and ``pow``).
3996 Those functions can `produce slightly different results
3997 <https://members.loria.fr/PZimmermann/papers/accuracy.pdf>`_ on
3998 different builds. Accordingly, selections can vary across builds
3999 even for the same seed.
4001 * The algorithms loop over the input and make selections based on
4002 ordinal position, so selections from unordered collections (such as
4003 sets) won't reproduce across sessions on the same platform using the
4004 same seed. For example, this won't reproduce::
4006 >> seed(8675309)
4007 >> sample(set('abcdefghijklmnopqrstuvwxyz'), 10)
4008 ['c', 'p', 'e', 'w', 's', 'a', 'j', 'd', 'n', 't']
4010 """
4011 iterator = iter(iterable)
4013 if k < 0:
4014 raise ValueError('k must be non-negative')
4016 if k == 0:
4017 return []
4019 if weights is not None and counts is not None:
4020 raise TypeError('weights and counts are mutually exclusive')
4022 elif weights is not None:
4023 weights = iter(weights)
4024 return _sample_weighted(iterator, k, weights, strict)
4026 elif counts is not None:
4027 counts = iter(counts)
4028 return _sample_counted(iterator, k, counts, strict)
4030 else:
4031 return _sample_unweighted(iterator, k, strict)
4034def is_sorted(iterable, key=None, reverse=False, strict=False):
4035 """Returns ``True`` if the items of iterable are in sorted order, and
4036 ``False`` otherwise. *key* and *reverse* have the same meaning that they do
4037 in the built-in :func:`sorted` function.
4039 >>> is_sorted(['1', '2', '3', '4', '5'], key=int)
4040 True
4041 >>> is_sorted([5, 4, 3, 1, 2], reverse=True)
4042 False
4044 If *strict*, tests for strict sorting, that is, returns ``False`` if equal
4045 elements are found:
4047 >>> is_sorted([1, 2, 2])
4048 True
4049 >>> is_sorted([1, 2, 2], strict=True)
4050 False
4052 The function returns ``False`` after encountering the first out-of-order
4053 item, which means it may produce results that differ from the built-in
4054 :func:`sorted` function for objects with unusual comparison dynamics
4055 (like ``math.nan``). If there are no out-of-order items, the iterable is
4056 exhausted.
4057 """
4058 it = iterable if (key is None) else map(key, iterable)
4059 a, b = tee(it)
4060 next(b, None)
4061 if reverse:
4062 b, a = a, b
4063 return all(map(lt, a, b)) if strict else not any(map(lt, b, a))
4066class AbortThread(BaseException):
4067 pass
4070class callback_iter:
4071 """Convert a function that uses callbacks to an iterator.
4073 .. deprecated:: 11.0.0
4074 Will be removed in a future major release.
4076 Let *func* be a function that takes a `callback` keyword argument.
4077 For example:
4079 >>> def func(callback=None):
4080 ... for i, c in [(1, 'a'), (2, 'b'), (3, 'c')]:
4081 ... if callback:
4082 ... callback(i, c)
4083 ... return 4
4086 Use ``with callback_iter(func)`` to get an iterator over the parameters
4087 that are delivered to the callback.
4089 >>> with callback_iter(func) as it:
4090 ... for args, kwargs in it:
4091 ... print(args)
4092 (1, 'a')
4093 (2, 'b')
4094 (3, 'c')
4096 The function will be called in a background thread. The ``done`` property
4097 indicates whether it has completed execution.
4099 >>> it.done
4100 True
4102 If it completes successfully, its return value will be available
4103 in the ``result`` property.
4105 >>> it.result
4106 4
4108 Notes:
4110 * If the function uses some keyword argument besides ``callback``, supply
4111 *callback_kwd*.
4112 * If it finished executing, but raised an exception, accessing the
4113 ``result`` property will raise the same exception.
4114 * If it hasn't finished executing, accessing the ``result``
4115 property from within the ``with`` block will raise ``RuntimeError``.
4116 * If it hasn't finished executing, accessing the ``result`` property from
4117 outside the ``with`` block will raise a
4118 ``more_itertools.AbortThread`` exception.
4119 * Provide *wait_seconds* to adjust how frequently the it is polled for
4120 output.
4122 """
4124 def __init__(self, func, callback_kwd='callback', wait_seconds=0.1):
4125 self._func = func
4126 self._callback_kwd = callback_kwd
4127 self._aborted = False
4128 self._future = None
4129 self._wait_seconds = wait_seconds
4131 # Lazily import concurrent.future
4132 self._module = __import__('concurrent.futures').futures
4133 self._executor = self._module.ThreadPoolExecutor(max_workers=1)
4134 self._iterator = self._reader()
4136 def __enter__(self):
4137 return self
4139 def __exit__(self, exc_type, exc_value, traceback):
4140 self._aborted = True
4141 self._executor.shutdown()
4143 def __iter__(self):
4144 return self
4146 def __next__(self):
4147 return next(self._iterator)
4149 @property
4150 def done(self):
4151 if self._future is None:
4152 return False
4153 return self._future.done()
4155 @property
4156 def result(self):
4157 if self._future:
4158 try:
4159 return self._future.result(timeout=0)
4160 except self._module.TimeoutError:
4161 pass
4163 raise RuntimeError('Function has not yet completed')
4165 def _reader(self):
4166 q = Queue()
4168 def callback(*args, **kwargs):
4169 if self._aborted:
4170 raise AbortThread('canceled by user')
4172 q.put((args, kwargs))
4174 self._future = self._executor.submit(
4175 self._func, **{self._callback_kwd: callback}
4176 )
4178 while True:
4179 try:
4180 item = q.get(timeout=self._wait_seconds)
4181 except Empty:
4182 pass
4183 else:
4184 q.task_done()
4185 yield item
4187 if self._future.done():
4188 break
4190 remaining = []
4191 while True:
4192 try:
4193 item = q.get_nowait()
4194 except Empty:
4195 break
4196 else:
4197 q.task_done()
4198 remaining.append(item)
4199 q.join()
4200 yield from remaining
4203def windowed_complete(iterable, n):
4204 """
4205 Yield ``(beginning, middle, end)`` tuples, where:
4207 * Each ``middle`` has *n* items from *iterable*
4208 * Each ``beginning`` has the items before the ones in ``middle``
4209 * Each ``end`` has the items after the ones in ``middle``
4211 >>> iterable = range(7)
4212 >>> n = 3
4213 >>> for beginning, middle, end in windowed_complete(iterable, n):
4214 ... print(beginning, middle, end)
4215 () (0, 1, 2) (3, 4, 5, 6)
4216 (0,) (1, 2, 3) (4, 5, 6)
4217 (0, 1) (2, 3, 4) (5, 6)
4218 (0, 1, 2) (3, 4, 5) (6,)
4219 (0, 1, 2, 3) (4, 5, 6) ()
4221 Note that *n* must be at least 0 and most equal to the length of
4222 *iterable*.
4224 This function will exhaust the iterable and may require significant
4225 storage.
4226 """
4227 if n < 0:
4228 raise ValueError('n must be >= 0')
4230 seq = tuple(iterable)
4231 size = len(seq)
4233 if n > size:
4234 raise ValueError('n must be <= len(seq)')
4236 for i in range(size - n + 1):
4237 beginning = seq[:i]
4238 middle = seq[i : i + n]
4239 end = seq[i + n :]
4240 yield beginning, middle, end
4243def all_unique(iterable, key=None):
4244 """
4245 Returns ``True`` if all the elements of *iterable* are unique (no two
4246 elements are equal).
4248 >>> all_unique('ABCB')
4249 False
4251 If a *key* function is specified, it will be used to make comparisons.
4253 >>> all_unique('ABCb')
4254 True
4255 >>> all_unique('ABCb', str.lower)
4256 False
4258 The function returns as soon as the first non-unique element is
4259 encountered. Iterables with a mix of hashable and unhashable items can
4260 be used, but the function will be slower for unhashable items.
4261 """
4262 seenset = set()
4263 seenset_add = seenset.add
4264 seenlist = []
4265 seenlist_add = seenlist.append
4266 for element in map(key, iterable) if key else iterable:
4267 try:
4268 if element in seenset:
4269 return False
4270 seenset_add(element)
4271 except TypeError:
4272 if element in seenlist:
4273 return False
4274 seenlist_add(element)
4275 return True
4278def nth_product(index, *iterables, repeat=1):
4279 """Equivalent to ``list(product(*iterables, repeat=repeat))[index]``.
4281 The products of *iterables* can be ordered lexicographically.
4282 :func:`nth_product` computes the product at sort position *index* without
4283 computing the previous products.
4285 >>> nth_product(8, range(2), range(2), range(2), range(2))
4286 (1, 0, 0, 0)
4288 The *repeat* keyword argument specifies the number of repetitions
4289 of the iterables. The above example is equivalent to::
4291 >>> nth_product(8, range(2), repeat=4)
4292 (1, 0, 0, 0)
4294 ``IndexError`` will be raised if the given *index* is invalid.
4295 """
4296 pools = tuple(map(tuple, reversed(iterables))) * repeat
4297 ns = tuple(map(len, pools))
4299 c = prod(ns)
4301 if index < 0:
4302 index += c
4303 if not 0 <= index < c:
4304 raise IndexError
4306 result = []
4307 for pool, n in zip(pools, ns):
4308 result.append(pool[index % n])
4309 index //= n
4311 return tuple(reversed(result))
4314def nth_permutation(iterable, r, index):
4315 """Equivalent to ``list(permutations(iterable, r))[index]```
4317 The subsequences of *iterable* that are of length *r* where order is
4318 important can be ordered lexicographically. :func:`nth_permutation`
4319 computes the subsequence at sort position *index* directly, without
4320 computing the previous subsequences.
4322 >>> nth_permutation('ghijk', 2, 5)
4323 ('h', 'i')
4325 ``ValueError`` will be raised If *r* is negative.
4326 ``IndexError`` will be raised if the given *index* is invalid.
4327 """
4328 pool = list(iterable)
4329 n = len(pool)
4330 if r is None:
4331 r = n
4332 c = perm(n, r)
4334 if index < 0:
4335 index += c
4336 if not 0 <= index < c:
4337 raise IndexError
4339 result = [0] * r
4340 q = index
4341 for d in range(n - r + 1, n + 1):
4342 q, i = divmod(q, d)
4343 result[n - d] = i
4344 if q == 0:
4345 break
4347 return tuple(map(pool.pop, result))
4350def nth_combination_with_replacement(iterable, r, index):
4351 """Equivalent to
4352 ``list(combinations_with_replacement(iterable, r))[index]``.
4355 The subsequences with repetition of *iterable* that are of length *r* can
4356 be ordered lexicographically. :func:`nth_combination_with_replacement`
4357 computes the subsequence at sort position *index* directly, without
4358 computing the previous subsequences with replacement.
4360 >>> nth_combination_with_replacement(range(5), 3, 5)
4361 (0, 1, 1)
4363 ``ValueError`` will be raised If *r* is negative.
4364 ``IndexError`` will be raised if the given *index* is invalid.
4365 """
4366 pool = tuple(iterable)
4367 n = len(pool)
4368 if r < 0:
4369 raise ValueError
4370 c = comb(n + r - 1, r) if n else 0 if r else 1
4372 if index < 0:
4373 index += c
4374 if not 0 <= index < c:
4375 raise IndexError
4377 result = []
4378 i = 0
4379 while r:
4380 r -= 1
4381 while n >= 0:
4382 num_combs = comb(n + r - 1, r)
4383 if index < num_combs:
4384 break
4385 n -= 1
4386 i += 1
4387 index -= num_combs
4388 result.append(pool[i])
4390 return tuple(result)
4393def value_chain(*args):
4394 """Yield all arguments passed to the function in the same order in which
4395 they were passed. If an argument itself is iterable then iterate over its
4396 values.
4398 >>> list(value_chain(1, 2, 3, [4, 5, 6]))
4399 [1, 2, 3, 4, 5, 6]
4401 Binary and text strings are not considered iterable and are emitted
4402 as-is:
4404 >>> list(value_chain('12', '34', ['56', '78']))
4405 ['12', '34', '56', '78']
4407 Pre- or postpend a single element to an iterable:
4409 >>> list(value_chain(1, [2, 3, 4, 5, 6]))
4410 [1, 2, 3, 4, 5, 6]
4411 >>> list(value_chain([1, 2, 3, 4, 5], 6))
4412 [1, 2, 3, 4, 5, 6]
4414 Multiple levels of nesting are not flattened.
4416 """
4417 scalar_types = (str, bytes)
4418 for value in args:
4419 if isinstance(value, scalar_types):
4420 yield value
4421 continue
4422 try:
4423 it = iter(value)
4424 except TypeError:
4425 yield value
4426 else:
4427 yield from it
4430def product_index(element, *iterables, repeat=1):
4431 """Equivalent to ``list(product(*iterables, repeat=repeat)).index(tuple(element))``
4433 The products of *iterables* can be ordered lexicographically.
4434 :func:`product_index` computes the first index of *element* without
4435 computing the previous products.
4437 >>> product_index([8, 2], range(10), range(5))
4438 42
4440 The *repeat* keyword argument specifies the number of repetitions
4441 of the iterables::
4443 >>> product_index([8, 0, 7], range(10), repeat=3)
4444 807
4446 ``ValueError`` will be raised if the given *element* isn't in the product
4447 of *args*.
4448 """
4449 elements = tuple(element)
4450 pools = tuple(map(tuple, iterables)) * repeat
4451 if len(elements) != len(pools):
4452 raise ValueError('element is not a product of args')
4454 index = 0
4455 for elem, pool in zip(elements, pools):
4456 index = index * len(pool) + pool.index(elem)
4457 return index
4460def combination_index(element, iterable):
4461 """Equivalent to ``list(combinations(iterable, r)).index(element)``
4463 The subsequences of *iterable* that are of length *r* can be ordered
4464 lexicographically. :func:`combination_index` computes the index of the
4465 first *element*, without computing the previous combinations.
4467 >>> combination_index('adf', 'abcdefg')
4468 10
4470 ``ValueError`` will be raised if the given *element* isn't one of the
4471 combinations of *iterable*.
4472 """
4473 element = enumerate(element)
4474 k, y = next(element, (None, None))
4475 if k is None:
4476 return 0
4478 indexes = []
4479 pool = enumerate(iterable)
4480 for n, x in pool:
4481 if x == y:
4482 indexes.append(n)
4483 tmp, y = next(element, (None, None))
4484 if tmp is None:
4485 break
4486 else:
4487 k = tmp
4488 else:
4489 raise ValueError('element is not a combination of iterable')
4491 n, _ = last(pool, default=(n, None))
4493 index = 1
4494 for i, j in enumerate(reversed(indexes), start=1):
4495 j = n - j
4496 if i <= j:
4497 index += comb(j, i)
4499 return comb(n + 1, k + 1) - index
4502def combination_with_replacement_index(element, iterable):
4503 """Equivalent to
4504 ``list(combinations_with_replacement(iterable, r)).index(element)``
4506 The subsequences with repetition of *iterable* that are of length *r* can
4507 be ordered lexicographically. :func:`combination_with_replacement_index`
4508 computes the index of the first *element*, without computing the previous
4509 combinations with replacement.
4511 >>> combination_with_replacement_index('adf', 'abcdefg')
4512 20
4514 ``ValueError`` will be raised if the given *element* isn't one of the
4515 combinations with replacement of *iterable*.
4516 """
4517 element = tuple(element)
4518 l = len(element)
4519 element = enumerate(element)
4521 k, y = next(element, (None, None))
4522 if k is None:
4523 return 0
4525 indexes = []
4526 pool = tuple(iterable)
4527 for n, x in enumerate(pool):
4528 while x == y:
4529 indexes.append(n)
4530 tmp, y = next(element, (None, None))
4531 if tmp is None:
4532 break
4533 else:
4534 k = tmp
4535 if y is None:
4536 break
4537 else:
4538 raise ValueError(
4539 'element is not a combination with replacement of iterable'
4540 )
4542 n = len(pool)
4543 occupations = [0] * n
4544 for p in indexes:
4545 occupations[p] += 1
4547 index = 0
4548 cumulative_sum = 0
4549 for k in range(1, n):
4550 cumulative_sum += occupations[k - 1]
4551 j = l + n - 1 - k - cumulative_sum
4552 i = n - k
4553 if i <= j:
4554 index += comb(j, i)
4556 return index
4559def permutation_index(element, iterable):
4560 """Equivalent to ``list(permutations(iterable, r)).index(element)```
4562 The subsequences of *iterable* that are of length *r* where order is
4563 important can be ordered lexicographically. :func:`permutation_index`
4564 computes the index of the first *element* directly, without computing
4565 the previous permutations.
4567 >>> permutation_index([1, 3, 2], range(5))
4568 19
4570 ``ValueError`` will be raised if the given *element* isn't one of the
4571 permutations of *iterable*.
4572 """
4573 index = 0
4574 pool = list(iterable)
4575 for i, x in zip(range(len(pool), -1, -1), element):
4576 r = pool.index(x)
4577 index = index * i + r
4578 del pool[r]
4580 return index
4583class countable:
4584 """Wrap *iterable* and keep a count of how many items have been consumed.
4586 The ``items_seen`` attribute starts at ``0`` and increments as the iterable
4587 is consumed:
4589 >>> iterable = map(str, range(10))
4590 >>> it = countable(iterable)
4591 >>> it.items_seen
4592 0
4593 >>> next(it), next(it)
4594 ('0', '1')
4595 >>> list(it)
4596 ['2', '3', '4', '5', '6', '7', '8', '9']
4597 >>> it.items_seen
4598 10
4599 """
4601 def __init__(self, iterable):
4602 self._iterator = iter(iterable)
4603 self.items_seen = 0
4605 def __iter__(self):
4606 return self
4608 def __next__(self):
4609 item = next(self._iterator)
4610 self.items_seen += 1
4612 return item
4615def chunked_even(iterable, n):
4616 """Break *iterable* into lists of approximately length *n*.
4617 Items are distributed such the lengths of the lists differ by at most
4618 1 item.
4620 >>> iterable = [1, 2, 3, 4, 5, 6, 7]
4621 >>> n = 3
4622 >>> list(chunked_even(iterable, n)) # List lengths: 3, 2, 2
4623 [[1, 2, 3], [4, 5], [6, 7]]
4624 >>> list(chunked(iterable, n)) # List lengths: 3, 3, 1
4625 [[1, 2, 3], [4, 5, 6], [7]]
4627 """
4628 iterator = iter(iterable)
4630 # Initialize a buffer to process the chunks while keeping
4631 # some back to fill any underfilled chunks
4632 min_buffer = (n - 1) * (n - 2)
4633 buffer = list(islice(iterator, min_buffer))
4635 # Append items until we have a completed chunk
4636 for _ in islice(map(buffer.append, iterator), n, None, n):
4637 yield buffer[:n]
4638 del buffer[:n]
4640 # Check if any chunks need addition processing
4641 if not buffer:
4642 return
4643 length = len(buffer)
4645 # Chunks are either size `full_size <= n` or `partial_size = full_size - 1`
4646 q, r = divmod(length, n)
4647 num_lists = q + (1 if r > 0 else 0)
4648 q, r = divmod(length, num_lists)
4649 full_size = q + (1 if r > 0 else 0)
4650 partial_size = full_size - 1
4651 num_full = length - partial_size * num_lists
4653 # Yield chunks of full size
4654 partial_start_idx = num_full * full_size
4655 if full_size > 0:
4656 for i in range(0, partial_start_idx, full_size):
4657 yield buffer[i : i + full_size]
4659 # Yield chunks of partial size
4660 if partial_size > 0:
4661 for i in range(partial_start_idx, length, partial_size):
4662 yield buffer[i : i + partial_size]
4665def zip_broadcast(*objects, scalar_types=(str, bytes), strict=False):
4666 """A version of :func:`zip` that "broadcasts" any scalar
4667 (i.e., non-iterable) items into output tuples.
4669 >>> iterable_1 = [1, 2, 3]
4670 >>> iterable_2 = ['a', 'b', 'c']
4671 >>> scalar = '_'
4672 >>> list(zip_broadcast(iterable_1, iterable_2, scalar))
4673 [(1, 'a', '_'), (2, 'b', '_'), (3, 'c', '_')]
4675 The *scalar_types* keyword argument determines what types are considered
4676 scalar. It is set to ``(str, bytes)`` by default. Set it to ``None`` to
4677 treat strings and byte strings as iterable:
4679 >>> list(zip_broadcast('abc', 0, 'xyz', scalar_types=None))
4680 [('a', 0, 'x'), ('b', 0, 'y'), ('c', 0, 'z')]
4682 If the *strict* keyword argument is ``True``, then
4683 ``ValueError`` will be raised if any of the iterables have
4684 different lengths.
4685 """
4687 def is_scalar(obj):
4688 if scalar_types and isinstance(obj, scalar_types):
4689 return True
4690 try:
4691 iter(obj)
4692 except TypeError:
4693 return True
4694 else:
4695 return False
4697 size = len(objects)
4698 if not size:
4699 return
4701 new_item = [None] * size
4702 iterables, iterable_positions = [], []
4703 for i, obj in enumerate(objects):
4704 if is_scalar(obj):
4705 new_item[i] = obj
4706 else:
4707 iterables.append(iter(obj))
4708 iterable_positions.append(i)
4710 if not iterables:
4711 yield tuple(objects)
4712 return
4714 for item in zip(*iterables, strict=strict):
4715 for i, new_item[i] in zip(iterable_positions, item):
4716 pass
4717 yield tuple(new_item)
4720def unique_in_window(iterable, n, key=None):
4721 """Yield the items from *iterable* that haven't been seen recently.
4722 *n* is the size of the sliding window.
4724 >>> iterable = [0, 1, 0, 2, 3, 0]
4725 >>> n = 3
4726 >>> list(unique_in_window(iterable, n))
4727 [0, 1, 2, 3, 0]
4729 The *key* function, if provided, will be used to determine uniqueness:
4731 >>> list(unique_in_window('abAcda', 3, key=lambda x: x.lower()))
4732 ['a', 'b', 'c', 'd', 'a']
4734 Updates a sliding window no larger than n and yields a value
4735 if the item only occurs once in the updated window.
4737 When `n == 1`, *unique_in_window* is memoryless:
4739 >>> list(unique_in_window('aab', n=1))
4740 ['a', 'a', 'b']
4742 The items in *iterable* must be hashable.
4744 """
4745 if n <= 0:
4746 raise ValueError('n must be greater than 0')
4748 window = deque(maxlen=n)
4749 counts = Counter()
4750 use_key = key is not None
4752 for item in iterable:
4753 if len(window) == n:
4754 to_discard = window[0]
4755 if counts[to_discard] == 1:
4756 del counts[to_discard]
4757 else:
4758 counts[to_discard] -= 1
4760 k = key(item) if use_key else item
4761 if k not in counts:
4762 yield item
4763 counts[k] += 1
4764 window.append(k)
4767def duplicates_everseen(iterable, key=None):
4768 """Yield duplicate elements after their first appearance.
4770 >>> list(duplicates_everseen('mississippi'))
4771 ['s', 'i', 's', 's', 'i', 'p', 'i']
4772 >>> list(duplicates_everseen('AaaBbbCccAaa', str.lower))
4773 ['a', 'a', 'b', 'b', 'c', 'c', 'A', 'a', 'a']
4775 This function is analogous to :func:`unique_everseen` and is subject to
4776 the same performance considerations.
4778 If you would like each duplicate to only appear once, much like ``uniq -d``
4779 in the Unix shell or ``Itertools::duplicates`` from the Rust ``itertools``
4780 crate, pass the return value of this function into :func:`unique_everseen`
4781 with the same ``key``.
4783 """
4784 seen_set = set()
4785 seen_list = []
4786 use_key = key is not None
4788 for element in iterable:
4789 k = key(element) if use_key else element
4790 try:
4791 if k not in seen_set:
4792 seen_set.add(k)
4793 else:
4794 yield element
4795 except TypeError:
4796 if k not in seen_list:
4797 seen_list.append(k)
4798 else:
4799 yield element
4802def duplicates_justseen(iterable, key=None):
4803 """Yields serially-duplicate elements after their first appearance.
4805 >>> list(duplicates_justseen('mississippi'))
4806 ['s', 's', 'p']
4807 >>> list(duplicates_justseen('AaaBbbCccAaa', str.lower))
4808 ['a', 'a', 'b', 'b', 'c', 'c', 'a', 'a']
4810 This function is analogous to :func:`unique_justseen`.
4812 """
4813 return flatten(g for _, g in groupby(iterable, key) for _ in g)
4816def classify_unique(iterable, key=None):
4817 """Classify each element in terms of its uniqueness.
4819 For each element in the input iterable, return a 3-tuple consisting of:
4821 1. The element itself
4822 2. ``False`` if the element is equal to the one preceding it in the input,
4823 ``True`` otherwise (i.e. the equivalent of :func:`unique_justseen`)
4824 3. ``False`` if this element has been seen anywhere in the input before,
4825 ``True`` otherwise (i.e. the equivalent of :func:`unique_everseen`)
4827 >>> list(classify_unique('otto')) # doctest: +NORMALIZE_WHITESPACE
4828 [('o', True, True),
4829 ('t', True, True),
4830 ('t', False, False),
4831 ('o', True, False)]
4833 This function is analogous to :func:`unique_everseen` and is subject to
4834 the same performance considerations.
4836 """
4837 seen_set = set()
4838 seen_list = []
4839 use_key = key is not None
4840 previous = None
4842 for i, element in enumerate(iterable):
4843 k = key(element) if use_key else element
4844 is_unique_justseen = not i or previous != k
4845 previous = k
4846 is_unique_everseen = False
4847 try:
4848 if k not in seen_set:
4849 seen_set.add(k)
4850 is_unique_everseen = True
4851 except TypeError:
4852 if k not in seen_list:
4853 seen_list.append(k)
4854 is_unique_everseen = True
4855 yield element, is_unique_justseen, is_unique_everseen
4858def minmax(iterable_or_value, *others, key=None, default=_marker):
4859 """Returns both the smallest and largest items from an iterable
4860 or from two or more arguments.
4862 >>> minmax([3, 1, 5])
4863 (1, 5)
4865 >>> minmax(4, 2, 6)
4866 (2, 6)
4868 If a *key* function is provided, it will be used to transform the input
4869 items for comparison.
4871 >>> minmax([5, 30], key=str) # '30' sorts before '5'
4872 (30, 5)
4874 If a *default* value is provided, it will be returned if there are no
4875 input items.
4877 >>> minmax([], default=(0, 0))
4878 (0, 0)
4880 Otherwise ``ValueError`` is raised.
4882 This function makes a single pass over the input elements and takes care to
4883 minimize the number of comparisons made during processing.
4885 Note that unlike the builtin ``max`` function, which always returns the first
4886 item with the maximum value, this function may return another item when there are
4887 ties.
4889 This function is based on the
4890 `recipe <https://code.activestate.com/recipes/577916-fast-minmax-function>`__ by
4891 Raymond Hettinger.
4892 """
4893 iterable = (iterable_or_value, *others) if others else iterable_or_value
4895 it = iter(iterable)
4897 try:
4898 lo = hi = next(it)
4899 except StopIteration as exc:
4900 if default is _marker:
4901 raise ValueError(
4902 '`minmax()` argument is an empty iterable. '
4903 'Provide a `default` value to suppress this error.'
4904 ) from exc
4905 return default
4907 # Different branches depending on the presence of key. This saves a lot
4908 # of unimportant copies which would slow the "key=None" branch
4909 # significantly down.
4910 if key is None:
4911 for x, y in zip_longest(it, it, fillvalue=lo):
4912 if y < x:
4913 if y < lo:
4914 lo = y
4915 if hi < x:
4916 hi = x
4917 else:
4918 if x < lo:
4919 lo = x
4920 if hi < y:
4921 hi = y
4923 else:
4924 lo_key = hi_key = key(lo)
4926 for x, y in zip_longest(it, it, fillvalue=lo):
4927 x_key, y_key = key(x), key(y)
4929 if y_key < x_key:
4930 if y_key < lo_key:
4931 lo, lo_key = y, y_key
4932 if hi_key < x_key:
4933 hi, hi_key = x, x_key
4934 else:
4935 if x_key < lo_key:
4936 lo, lo_key = x, x_key
4937 if hi_key < y_key:
4938 hi, hi_key = y, y_key
4940 return lo, hi
4943def constrained_batches(
4944 iterable, max_size, max_count=None, get_len=len, strict=True
4945):
4946 """Yield batches of items from *iterable* with a combined size limited by
4947 *max_size*.
4949 >>> iterable = [b'12345', b'123', b'12345678', b'1', b'1', b'12', b'1']
4950 >>> list(constrained_batches(iterable, 10))
4951 [(b'12345', b'123'), (b'12345678', b'1', b'1'), (b'12', b'1')]
4953 If a *max_count* is supplied, the number of items per batch is also
4954 limited:
4956 >>> iterable = [b'12345', b'123', b'12345678', b'1', b'1', b'12', b'1']
4957 >>> list(constrained_batches(iterable, 10, max_count = 2))
4958 [(b'12345', b'123'), (b'12345678', b'1'), (b'1', b'12'), (b'1',)]
4960 If a *get_len* function is supplied, use that instead of :func:`len` to
4961 determine item size.
4963 If *strict* is ``True``, raise ``ValueError`` if any single item is bigger
4964 than *max_size*. Otherwise, allow single items to exceed *max_size*.
4965 """
4966 if max_size <= 0:
4967 raise ValueError('maximum size must be greater than zero')
4969 batch = []
4970 batch_size = 0
4971 batch_count = 0
4972 for item in iterable:
4973 item_len = get_len(item)
4974 if strict and item_len > max_size:
4975 raise ValueError('item size exceeds maximum size')
4977 reached_count = batch_count == max_count
4978 reached_size = item_len + batch_size > max_size
4979 if batch_count and (reached_size or reached_count):
4980 yield tuple(batch)
4981 batch.clear()
4982 batch_size = 0
4983 batch_count = 0
4985 batch.append(item)
4986 batch_size += item_len
4987 batch_count += 1
4989 if batch:
4990 yield tuple(batch)
4993def gray_product(*iterables, repeat=1):
4994 """Like :func:`itertools.product`, but return tuples in an order such
4995 that only one element in the generated tuple changes from one iteration
4996 to the next.
4998 >>> list(gray_product('AB','CD'))
4999 [('A', 'C'), ('B', 'C'), ('B', 'D'), ('A', 'D')]
5001 The *repeat* keyword argument specifies the number of repetitions
5002 of the iterables. For example, ``gray_product('AB', repeat=3)`` is
5003 equivalent to ``gray_product('AB', 'AB', 'AB')``.
5005 This function consumes all of the input iterables before producing output.
5006 If any of the input iterables have fewer than two items, ``ValueError``
5007 is raised.
5009 For information on the algorithm, see
5010 `this section <https://www-cs-faculty.stanford.edu/~knuth/fasc2a.ps.gz>`__
5011 of Donald Knuth's *The Art of Computer Programming*.
5012 """
5013 all_iterables = tuple(map(tuple, iterables)) * repeat
5014 iterable_count = len(all_iterables)
5015 for iterable in all_iterables:
5016 if len(iterable) < 2:
5017 raise ValueError("each iterable must have two or more items")
5019 # This is based on "Algorithm H" from section 7.2.1.1, page 20.
5020 # a holds the indexes of the source iterables for the n-tuple to be yielded
5021 # f is the array of "focus pointers"
5022 # o is the array of "directions"
5023 a = [0] * iterable_count
5024 f = list(range(iterable_count + 1))
5025 o = [1] * iterable_count
5026 while True:
5027 yield tuple(all_iterables[i][a[i]] for i in range(iterable_count))
5028 j = f[0]
5029 f[0] = 0
5030 if j == iterable_count:
5031 break
5032 a[j] = a[j] + o[j]
5033 if a[j] == 0 or a[j] == len(all_iterables[j]) - 1:
5034 o[j] = -o[j]
5035 f[j] = f[j + 1]
5036 f[j + 1] = j + 1
5039def partial_product(*iterables, repeat=1):
5040 """Yields tuples containing one item from each iterator, with subsequent
5041 tuples changing a single item at a time by advancing each iterator until it
5042 is exhausted. This sequence guarantees every value in each iterable is
5043 output at least once without generating all possible combinations.
5045 This may be useful, for example, when testing an expensive function.
5047 >>> list(partial_product('AB', 'C', 'DEF'))
5048 [('A', 'C', 'D'), ('B', 'C', 'D'), ('B', 'C', 'E'), ('B', 'C', 'F')]
5050 The *repeat* keyword argument specifies the number of repetitions
5051 of the iterables. For example, ``partial_product('AB', repeat=3)`` is
5052 equivalent to ``partial_product('AB', 'AB', 'AB')``.
5053 """
5055 all_iterables = tuple(map(tuple, iterables)) * repeat
5056 iterators = tuple(map(iter, all_iterables))
5058 try:
5059 prod = [next(it) for it in iterators]
5060 except StopIteration:
5061 return
5062 yield tuple(prod)
5064 for i, it in enumerate(iterators):
5065 for prod[i] in it:
5066 yield tuple(prod)
5069def takewhile_inclusive(predicate, iterable):
5070 """A variant of :func:`takewhile` that yields one additional element.
5072 >>> list(takewhile_inclusive(lambda x: x < 5, [1, 4, 6, 4, 1]))
5073 [1, 4, 6]
5075 :func:`takewhile` would return ``[1, 4]``.
5076 """
5077 for x in iterable:
5078 yield x
5079 if not predicate(x):
5080 break
5083def outer_product(func, xs, ys, *args, **kwargs):
5084 """A generalized outer product that applies a binary function to all
5085 pairs of items. Returns a 2D matrix with ``len(xs)`` rows and ``len(ys)``
5086 columns.
5087 Also accepts ``*args`` and ``**kwargs`` that are passed to ``func``.
5089 Multiplication table:
5091 >>> from operator import mul
5092 >>> list(outer_product(mul, range(1, 4), range(1, 6)))
5093 [(1, 2, 3, 4, 5), (2, 4, 6, 8, 10), (3, 6, 9, 12, 15)]
5095 Cross tabulation:
5097 >>> xs = ['A', 'B', 'A', 'A', 'B', 'B', 'A', 'A', 'B', 'B']
5098 >>> ys = ['X', 'X', 'X', 'Y', 'Z', 'Z', 'Y', 'Y', 'Z', 'Z']
5099 >>> pair_counts = Counter(zip(xs, ys))
5100 >>> count_rows = lambda x, y: pair_counts[x, y]
5101 >>> list(outer_product(count_rows, sorted(set(xs)), sorted(set(ys))))
5102 [(2, 3, 0), (1, 0, 4)]
5104 Usage with ``*args`` and ``**kwargs``:
5106 >>> animals = ['cat', 'wolf', 'mouse']
5107 >>> list(outer_product(min, animals, animals, key=len))
5108 [('cat', 'cat', 'cat'), ('cat', 'wolf', 'wolf'), ('cat', 'wolf', 'mouse')]
5109 """
5110 ys = tuple(ys)
5111 return batched(
5112 starmap(lambda x, y: func(x, y, *args, **kwargs), product(xs, ys)),
5113 n=len(ys),
5114 )
5117def iter_suppress(iterable, *exceptions):
5118 """Yield each of the items from *iterable*. If the iteration raises one of
5119 the specified *exceptions*, that exception will be suppressed and iteration
5120 will stop.
5122 >>> from itertools import chain
5123 >>> def breaks_at_five(x):
5124 ... while True:
5125 ... if x >= 5:
5126 ... raise RuntimeError
5127 ... yield x
5128 ... x += 1
5129 >>> it_1 = iter_suppress(breaks_at_five(1), RuntimeError)
5130 >>> it_2 = iter_suppress(breaks_at_five(2), RuntimeError)
5131 >>> list(chain(it_1, it_2))
5132 [1, 2, 3, 4, 2, 3, 4]
5133 """
5134 try:
5135 yield from iterable
5136 except exceptions:
5137 return
5140def filter_map(func, iterable):
5141 """Apply *func* to every element of *iterable*, yielding only those which
5142 are not ``None``.
5144 >>> elems = ['1', 'a', '2', 'b', '3']
5145 >>> list(filter_map(lambda s: int(s) if s.isnumeric() else None, elems))
5146 [1, 2, 3]
5147 """
5148 for x in iterable:
5149 y = func(x)
5150 if y is not None:
5151 yield y
5154def powerset_of_sets(iterable, *, baseset=set):
5155 """Yields all possible subsets of the iterable.
5157 >>> list(powerset_of_sets([1, 2, 3])) # doctest: +SKIP
5158 [set(), {1}, {2}, {3}, {1, 2}, {1, 3}, {2, 3}, {1, 2, 3}]
5159 >>> list(powerset_of_sets([1, 1, 0])) # doctest: +SKIP
5160 [set(), {1}, {0}, {0, 1}]
5162 :func:`powerset_of_sets` takes care to minimize the number
5163 of hash operations performed.
5165 The *baseset* parameter determines what kind of sets are
5166 constructed, either *set* or *frozenset*.
5167 """
5168 sets = tuple(dict.fromkeys(map(frozenset, zip(iterable))))
5169 union = baseset().union
5170 return chain.from_iterable(
5171 starmap(union, combinations(sets, r)) for r in range(len(sets) + 1)
5172 )
5175def join_mappings(**field_to_map):
5176 """
5177 Joins multiple mappings together using their common keys.
5179 >>> user_scores = {'elliot': 50, 'claris': 60}
5180 >>> user_times = {'elliot': 30, 'claris': 40}
5181 >>> join_mappings(score=user_scores, time=user_times)
5182 {'elliot': {'score': 50, 'time': 30}, 'claris': {'score': 60, 'time': 40}}
5183 """
5184 ret = defaultdict(dict)
5186 for field_name, mapping in field_to_map.items():
5187 for key, value in mapping.items():
5188 ret[key][field_name] = value
5190 return dict(ret)
5193def _complex_sumprod(v1, v2):
5194 """High precision sumprod() for complex numbers.
5195 Used by :func:`dft` and :func:`idft`.
5196 """
5198 real = attrgetter('real')
5199 imag = attrgetter('imag')
5200 r1 = chain(map(real, v1), map(neg, map(imag, v1)))
5201 r2 = chain(map(real, v2), map(imag, v2))
5202 i1 = chain(map(real, v1), map(imag, v1))
5203 i2 = chain(map(imag, v2), map(real, v2))
5204 return complex(_fsumprod(r1, r2), _fsumprod(i1, i2))
5207def dft(xarr):
5208 """Discrete Fourier Transform. *xarr* is a sequence of complex numbers.
5209 Yields the components of the corresponding transformed output vector.
5211 >>> import cmath
5212 >>> xarr = [1, 2-1j, -1j, -1+2j] # time domain
5213 >>> Xarr = [2, -2-2j, -2j, 4+4j] # frequency domain
5214 >>> magnitudes, phases = zip(*map(cmath.polar, Xarr))
5215 >>> all(map(cmath.isclose, dft(xarr), Xarr))
5216 True
5218 Inputs are restricted to numeric types that can add and multiply
5219 with a complex number. This includes int, float, complex, and
5220 Fraction, but excludes Decimal.
5222 See :func:`idft` for the inverse Discrete Fourier Transform.
5223 """
5224 N = len(xarr)
5225 roots_of_unity = [e ** (n / N * tau * -1j) for n in range(N)]
5226 for k in range(N):
5227 coeffs = [roots_of_unity[k * n % N] for n in range(N)]
5228 yield _complex_sumprod(xarr, coeffs)
5231def idft(Xarr):
5232 """Inverse Discrete Fourier Transform. *Xarr* is a sequence of
5233 complex numbers. Yields the components of the corresponding
5234 inverse-transformed output vector.
5236 >>> import cmath
5237 >>> xarr = [1, 2-1j, -1j, -1+2j] # time domain
5238 >>> Xarr = [2, -2-2j, -2j, 4+4j] # frequency domain
5239 >>> all(map(cmath.isclose, idft(Xarr), xarr))
5240 True
5242 Inputs are restricted to numeric types that can add and multiply
5243 with a complex number. This includes int, float, complex, and
5244 Fraction, but excludes Decimal.
5246 See :func:`dft` for the Discrete Fourier Transform.
5247 """
5248 N = len(Xarr)
5249 roots_of_unity = [e ** (n / N * tau * 1j) for n in range(N)]
5250 for k in range(N):
5251 coeffs = [roots_of_unity[k * n % N] for n in range(N)]
5252 yield _complex_sumprod(Xarr, coeffs) / N
5255def doublestarmap(func, iterable):
5256 """Apply *func* to every item of *iterable* by dictionary unpacking
5257 the item into *func*.
5259 The difference between :func:`itertools.starmap` and :func:`doublestarmap`
5260 parallels the distinction between ``func(*a)`` and ``func(**a)``.
5262 >>> iterable = [{'a': 1, 'b': 2}, {'a': 40, 'b': 60}]
5263 >>> list(doublestarmap(lambda a, b: a + b, iterable))
5264 [3, 100]
5266 ``TypeError`` will be raised if *func*'s signature doesn't match the
5267 mapping contained in *iterable* or if *iterable* does not contain mappings.
5268 """
5269 for item in iterable:
5270 yield func(**item)
5273def _nth_prime_bounds(n):
5274 """Bounds for the nth prime (counting from 1): lb < p_n < ub."""
5275 # At and above 688,383, the lb/ub spread is under 0.003 * p_n.
5277 if n < 1:
5278 raise ValueError
5280 if n < 6:
5281 return (n, 2.25 * n)
5283 # https://en.wikipedia.org/wiki/Prime-counting_function#Inequalities
5284 upper_bound = n * log(n * log(n))
5285 lower_bound = upper_bound - n
5286 if n >= 688_383:
5287 upper_bound -= n * (1.0 - (log(log(n)) - 2.0) / log(n))
5289 return lower_bound, upper_bound
5292def nth_prime(n, *, approximate=False):
5293 """Return the nth prime (counting from 0).
5295 >>> nth_prime(0)
5296 2
5297 >>> nth_prime(100)
5298 547
5300 If *approximate* is set to True, will return a prime close
5301 to the nth prime. The estimation is much faster than computing
5302 an exact result.
5304 >>> nth_prime(200_000_000, approximate=True) # Exact result is 4222234763
5305 4217820427
5307 """
5308 lb, ub = _nth_prime_bounds(n + 1)
5310 if not approximate or n <= 1_000_000:
5311 return nth(sieve(ceil(ub)), n)
5313 # Search from the midpoint and return the first odd prime
5314 odd = floor((lb + ub) / 2) | 1
5315 return first_true(count(odd, step=2), pred=is_prime)
5318def argmin(iterable, *, key=None):
5319 """
5320 Index of the first occurrence of a minimum value in an iterable.
5322 >>> argmin('efghabcdijkl')
5323 4
5324 >>> argmin([3, 2, 1, 0, 4, 2, 1, 0])
5325 3
5327 For example, look up a label corresponding to the position
5328 of a value that minimizes a cost function::
5330 >>> def cost(x):
5331 ... "Days for a wound to heal given a subject's age."
5332 ... return x**2 - 20*x + 150
5333 ...
5334 >>> labels = ['homer', 'marge', 'bart', 'lisa', 'maggie']
5335 >>> ages = [ 35, 30, 10, 9, 1 ]
5337 # Fastest healing family member
5338 >>> labels[argmin(ages, key=cost)]
5339 'bart'
5341 # Age with fastest healing
5342 >>> min(ages, key=cost)
5343 10
5345 """
5346 if key is not None:
5347 iterable = map(key, iterable)
5348 return min(enumerate(iterable), key=itemgetter(1))[0]
5351def argmax(iterable, *, key=None):
5352 """
5353 Index of the first occurrence of a maximum value in an iterable.
5355 >>> argmax('abcdefghabcd')
5356 7
5357 >>> argmax([0, 1, 2, 3, 3, 2, 1, 0])
5358 3
5360 For example, identify the best machine learning model::
5362 >>> models = ['svm', 'random forest', 'knn', 'naïve bayes']
5363 >>> accuracy = [ 68, 61, 84, 72 ]
5365 # Most accurate model
5366 >>> models[argmax(accuracy)]
5367 'knn'
5369 # Best accuracy
5370 >>> max(accuracy)
5371 84
5373 """
5374 if key is not None:
5375 iterable = map(key, iterable)
5376 return max(enumerate(iterable), key=itemgetter(1))[0]
5379def _extract_monotonic(iterator, indices):
5380 'Non-decreasing indices, lazily consumed'
5381 num_read = 0
5382 for index in indices:
5383 advance = index - num_read
5384 try:
5385 value = next(islice(iterator, advance, None))
5386 except ValueError:
5387 if advance != -1 or index < 0:
5388 raise ValueError(f'Invalid index: {index}') from None
5389 except StopIteration:
5390 raise IndexError(index) from None
5391 else:
5392 num_read += advance + 1
5393 yield value
5396def _extract_buffered(iterator, index_and_position):
5397 'Arbitrary index order, greedily consumed'
5398 buffer = {}
5399 iterator_position = -1
5400 next_to_emit = 0
5402 for index, order in index_and_position:
5403 advance = index - iterator_position
5404 if advance:
5405 try:
5406 value = next(islice(iterator, advance - 1, None))
5407 except StopIteration:
5408 raise IndexError(index) from None
5409 iterator_position = index
5411 buffer[order] = value
5413 while next_to_emit in buffer:
5414 yield buffer.pop(next_to_emit)
5415 next_to_emit += 1
5418def extract(iterable, indices, *, monotonic=False):
5419 """Yield values at the specified indices.
5421 Example:
5423 >>> data = 'abcdefghijklmnopqrstuvwxyz'
5424 >>> list(extract(data, [7, 4, 11, 11, 14]))
5425 ['h', 'e', 'l', 'l', 'o']
5427 The *iterable* is consumed lazily and can be infinite.
5429 When *monotonic* is false, the *indices* are consumed immediately
5430 and must be finite. When *monotonic* is true, *indices* are consumed
5431 lazily and can be infinite but must be non-decreasing.
5433 Raises ``IndexError`` if an index lies beyond the iterable.
5434 Raises ``ValueError`` for a negative index or for a decreasing
5435 index when *monotonic* is true.
5436 """
5438 iterator = iter(iterable)
5439 indices = iter(indices)
5441 if monotonic:
5442 return _extract_monotonic(iterator, indices)
5444 index_and_position = sorted(zip(indices, count()))
5445 if index_and_position and index_and_position[0][0] < 0:
5446 raise ValueError('Indices must be non-negative')
5447 return _extract_buffered(iterator, index_and_position)
5450class serialize:
5451 """Wrap a non-concurrent iterator with a lock to enforce sequential access.
5453 Applies a non-reentrant lock around calls to ``__next__``, allowing
5454 iterator and generator instances to be shared by multiple consumer
5455 threads.
5456 """
5458 __slots__ = ('_iterator', '_lock')
5460 def __init__(self, iterable):
5461 self._iterator = iter(iterable)
5462 self._lock = Lock()
5464 def __iter__(self):
5465 return self
5467 def __next__(self):
5468 with self._lock:
5469 return next(self._iterator)
5471 def send(self, value, /):
5472 """Send a value to a generator.
5474 Raises AttributeError if not a generator.
5475 """
5476 with self._lock:
5477 return self._iterator.send(value)
5479 def throw(self, *args):
5480 """Call throw() on a generator.
5482 Raises AttributeError if not a generator.
5483 """
5484 with self._lock:
5485 return self._iterator.throw(*args)
5487 def close(self):
5488 """Call close() on a generator.
5490 Raises AttributeError if not a generator.
5491 """
5492 with self._lock:
5493 return self._iterator.close()
5496def synchronized(func):
5497 """Wrap an iterator-returning callable to make its iterators thread-safe.
5499 Existing itertools and more-itertools can be wrapped so that their
5500 iterator instances are serialized.
5502 For example, ``itertools.count`` does not make thread-safe instances,
5503 but that is easily fixed with::
5505 atomic_counter = synchronized(itertools.count)
5507 Can also be used as a decorator for generator functions definitions
5508 so that the generator instances are serialized::
5510 @synchronized
5511 def enumerate_and_timestamp(iterable):
5512 for count, value in enumerate(iterable):
5513 yield count, time_ns(), value
5515 """
5517 @wraps(func)
5518 def inner(*args, **kwargs):
5519 iterator = func(*args, **kwargs)
5520 return serialize(iterator)
5522 return inner
5525def concurrent_tee(iterable, n=2):
5526 """Variant of itertools.tee() but with guaranteed threading semantics.
5528 Takes a non-threadsafe iterator as an input and creates concurrent
5529 tee objects for other threads to have reliable independent copies of
5530 the data stream.
5532 The new iterators are only thread-safe if consumed within a single thread.
5533 To share just one of the new iterators across multiple threads, wrap it
5534 with :func:`serialize`.
5535 """
5537 if n < 0:
5538 raise ValueError
5539 if n == 0:
5540 return ()
5541 iterator = _concurrent_tee(iterable)
5542 result = [iterator]
5543 for _ in range(n - 1):
5544 result.append(_concurrent_tee(iterator))
5545 return tuple(result)
5548class _concurrent_tee:
5549 __slots__ = ('iterator', 'link', 'lock')
5551 def __init__(self, iterable):
5552 if isinstance(iterable, _concurrent_tee):
5553 self.iterator = iterable.iterator
5554 self.link = iterable.link
5555 self.lock = iterable.lock
5556 else:
5557 self.iterator = iter(iterable)
5558 self.link = [None, None]
5559 self.lock = Lock()
5561 def __iter__(self):
5562 return self
5564 def __next__(self):
5565 link = self.link
5566 if link[1] is None:
5567 with self.lock:
5568 if link[1] is None:
5569 link[0] = next(self.iterator)
5570 link[1] = [None, None]
5571 value, self.link = link
5572 return value
5575def subfactorial(n):
5576 """Number of permutations of *n* elements with no fixed points.
5578 The :func:`subfactorial` function computes the length of
5579 :func:`derangements`. For example, there are 1,854 ways to
5580 rearrange the letters in word "epsilon" without leaving any
5581 letter in its original position:
5583 >>> from more_itertools import derangements, ilen
5584 >>> ilen(derangements('epsilon'))
5585 1854
5586 >>> subfactorial(len('epsilon'))
5587 1854
5589 Reference: https://oeis.org/A000166
5591 """
5592 if n < 0:
5593 raise ValueError
5594 sf = adj = 1
5595 for i in range(n + 1):
5596 sf = sf * i + adj
5597 adj = -adj
5598 return sf