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 Duplicate permutations arise when there are duplicated elements in the
759 input iterable. The number of items returned is
760 `n! / (x_1! * x_2! * ... * x_n!)`, where `n` is the total number of
761 items input, and each `x_i` is the count of a distinct item in the input
762 sequence. The function :func:`multinomial` computes this directly.
764 If *r* is given, only the *r*-length permutations are yielded.
766 >>> sorted(distinct_permutations([1, 0, 1], r=2))
767 [(0, 1), (1, 0), (1, 1)]
768 >>> sorted(distinct_permutations(range(3), r=2))
769 [(0, 1), (0, 2), (1, 0), (1, 2), (2, 0), (2, 1)]
771 *iterable* need not be sortable, but note that using equal (``x == y``)
772 but non-identical (``id(x) != id(y)``) elements may produce surprising
773 behavior. For example, ``1`` and ``True`` are equal but non-identical:
775 >>> list(distinct_permutations([1, True, '3'])) # doctest: +SKIP
776 [
777 (1, True, '3'),
778 (1, '3', True),
779 ('3', 1, True)
780 ]
781 >>> list(distinct_permutations([1, 2, '3'])) # doctest: +SKIP
782 [
783 (1, 2, '3'),
784 (1, '3', 2),
785 (2, 1, '3'),
786 (2, '3', 1),
787 ('3', 1, 2),
788 ('3', 2, 1)
789 ]
790 """
792 # Algorithm: https://w.wiki/Qai
793 def _full(A):
794 while True:
795 # Yield the permutation we have
796 yield tuple(A)
798 # Find the largest index i such that A[i] < A[i + 1]
799 for i in range(size - 2, -1, -1):
800 if A[i] < A[i + 1]:
801 break
802 # If no such index exists, this permutation is the last one
803 else:
804 return
806 # Find the largest index j greater than j such that A[i] < A[j]
807 for j in range(size - 1, i, -1):
808 if A[i] < A[j]:
809 break
811 # Swap the value of A[i] with that of A[j], then reverse the
812 # sequence from A[i + 1] to form the new permutation
813 A[i], A[j] = A[j], A[i]
814 A[i + 1 :] = A[: i - size : -1] # A[i + 1:][::-1]
816 # Algorithm: modified from the above
817 def _partial(A, r):
818 # Split A into the first r items and the last r items
819 head, tail = A[:r], A[r:]
820 right_head_indexes = range(r - 1, -1, -1)
821 left_tail_indexes = range(len(tail))
823 while True:
824 # Yield the permutation we have
825 yield tuple(head)
827 # Starting from the right, find the first index of the head with
828 # value smaller than the maximum value of the tail - call it i.
829 pivot = tail[-1]
830 for i in right_head_indexes:
831 if head[i] < pivot:
832 break
833 pivot = head[i]
834 else:
835 return
837 # Starting from the left, find the first value of the tail
838 # with a value greater than head[i] and swap.
839 for j in left_tail_indexes:
840 if tail[j] > head[i]:
841 head[i], tail[j] = tail[j], head[i]
842 break
843 # If we didn't find one, start from the right and find the first
844 # index of the head with a value greater than head[i] and swap.
845 else:
846 for j in right_head_indexes:
847 if head[j] > head[i]:
848 head[i], head[j] = head[j], head[i]
849 break
851 # Reverse head[i + 1:] and swap it with tail[:r - (i + 1)]
852 tail += head[: i - r : -1] # head[i + 1:][::-1]
853 i += 1
854 head[i:], tail[:] = tail[: r - i], tail[r - i :]
856 items = list(iterable)
858 try:
859 items.sort()
860 sortable = True
861 except TypeError:
862 sortable = False
864 indices_dict = defaultdict(list)
866 for item in items:
867 indices_dict[items.index(item)].append(item)
869 indices = [items.index(item) for item in items]
870 indices.sort()
872 equivalent_items = {k: cycle(v) for k, v in indices_dict.items()}
874 def permuted_items(permuted_indices):
875 return tuple(
876 next(equivalent_items[index]) for index in permuted_indices
877 )
879 size = len(items)
880 if r is None:
881 r = size
883 # functools.partial(_partial, ... )
884 algorithm = _full if (r == size) else partial(_partial, r=r)
886 if 0 < r <= size:
887 if sortable:
888 return algorithm(items)
889 else:
890 return (
891 permuted_items(permuted_indices)
892 for permuted_indices in algorithm(indices)
893 )
895 return iter(() if r else ((),))
898def derangements(iterable, r=None):
899 """Yield successive derangements of the elements in *iterable*.
901 A derangement is a permutation in which no element appears at its original
902 index. In other words, a derangement is a permutation that has no fixed points.
904 Suppose Alice, Bob, Carol, and Dave are playing Secret Santa.
905 The code below outputs all of the different ways to assign gift recipients
906 such that nobody is assigned to himself or herself:
908 >>> for d in derangements(['Alice', 'Bob', 'Carol', 'Dave']):
909 ... print(', '.join(d))
910 Bob, Alice, Dave, Carol
911 Bob, Carol, Dave, Alice
912 Bob, Dave, Alice, Carol
913 Carol, Alice, Dave, Bob
914 Carol, Dave, Alice, Bob
915 Carol, Dave, Bob, Alice
916 Dave, Alice, Bob, Carol
917 Dave, Carol, Alice, Bob
918 Dave, Carol, Bob, Alice
920 If *r* is given, only the *r*-length derangements are yielded.
922 >>> sorted(derangements(range(3), 2))
923 [(1, 0), (1, 2), (2, 0)]
924 >>> sorted(derangements([0, 2, 3], 2))
925 [(2, 0), (2, 3), (3, 0)]
927 Elements are treated as unique based on their position, not on their value.
929 Consider the Secret Santa example with two *different* people who have
930 the *same* name. Then there are two valid gift assignments even though
931 it might appear that a person is assigned to themselves:
933 >>> names = ['Alice', 'Bob', 'Bob']
934 >>> list(derangements(names))
935 [('Bob', 'Bob', 'Alice'), ('Bob', 'Alice', 'Bob')]
937 To avoid confusion, make the inputs distinct:
939 >>> deduped = [f'{name}{index}' for index, name in enumerate(names)]
940 >>> list(derangements(deduped))
941 [('Bob1', 'Bob2', 'Alice0'), ('Bob2', 'Alice0', 'Bob1')]
943 The number of derangements of a set of size *n* is known as the
944 "subfactorial of n". For n > 0, the subfactorial is:
945 ``round(math.factorial(n) / math.e)``. The more-itertools function
946 :func:`subfactorial` computes this directly.
948 References:
950 * Article: https://www.numberanalytics.com/blog/ultimate-guide-to-derangements-in-combinatorics
951 * Sizes: https://oeis.org/A000166
952 """
953 xs = tuple(iterable)
954 ys = tuple(range(len(xs)))
955 return compress(
956 permutations(xs, r=r),
957 map(all, map(map, repeat(is_not), repeat(ys), permutations(ys, r=r))),
958 )
961def intersperse(e, iterable, n=1):
962 """Intersperse filler element *e* among the items in *iterable*, leaving
963 *n* items between each filler element.
965 >>> list(intersperse('!', [1, 2, 3, 4, 5]))
966 [1, '!', 2, '!', 3, '!', 4, '!', 5]
968 >>> list(intersperse(None, [1, 2, 3, 4, 5], n=2))
969 [1, 2, None, 3, 4, None, 5]
971 """
972 if n == 0:
973 raise ValueError('n must be > 0')
974 elif n == 1:
975 # interleave(repeat(e), iterable) -> e, x_0, e, x_1, e, x_2...
976 # islice(..., 1, None) -> x_0, e, x_1, e, x_2...
977 return islice(interleave(repeat(e), iterable), 1, None)
978 else:
979 # interleave(filler, chunks) -> [e], [x_0, x_1], [e], [x_2, x_3]...
980 # islice(..., 1, None) -> [x_0, x_1], [e], [x_2, x_3]...
981 # flatten(...) -> x_0, x_1, e, x_2, x_3...
982 filler = repeat([e])
983 chunks = chunked(iterable, n)
984 return flatten(islice(interleave(filler, chunks), 1, None))
987def unique_to_each(*iterables):
988 """Return the elements from each of the input iterables that aren't in the
989 other input iterables.
991 For example, suppose you have a set of packages, each with a set of
992 dependencies::
994 {'pkg_1': {'A', 'B'}, 'pkg_2': {'B', 'C'}, 'pkg_3': {'B', 'D'}}
996 If you remove one package, which dependencies can also be removed?
998 If ``pkg_1`` is removed, then ``A`` is no longer necessary - it is not
999 associated with ``pkg_2`` or ``pkg_3``. Similarly, ``C`` is only needed for
1000 ``pkg_2``, and ``D`` is only needed for ``pkg_3``::
1002 >>> unique_to_each({'A', 'B'}, {'B', 'C'}, {'B', 'D'})
1003 [['A'], ['C'], ['D']]
1005 If there are duplicates in one input iterable that aren't in the others
1006 they will be duplicated in the output. Input order is preserved::
1008 >>> unique_to_each("mississippi", "missouri")
1009 [['p', 'p'], ['o', 'u', 'r']]
1011 It is assumed that the elements of each iterable are hashable.
1013 """
1014 pool = [list(it) for it in iterables]
1015 counts = Counter(chain.from_iterable(map(set, pool)))
1016 uniques = {element for element in counts if counts[element] == 1}
1017 return [list(filter(uniques.__contains__, it)) for it in pool]
1020def windowed(seq, n, fillvalue=None, step=1):
1021 """Return a sliding window of width *n* over the given iterable.
1023 >>> all_windows = windowed([1, 2, 3, 4, 5], 3)
1024 >>> list(all_windows)
1025 [(1, 2, 3), (2, 3, 4), (3, 4, 5)]
1027 When the window is larger than the iterable, *fillvalue* is used in place
1028 of missing values:
1030 >>> list(windowed([1, 2, 3], 4))
1031 [(1, 2, 3, None)]
1033 Each window will advance in increments of *step*:
1035 >>> list(windowed([1, 2, 3, 4, 5, 6], 3, fillvalue='!', step=2))
1036 [(1, 2, 3), (3, 4, 5), (5, 6, '!')]
1038 To slide into the iterable's items, use :func:`chain` to add filler items
1039 to the left:
1041 >>> iterable = [1, 2, 3, 4]
1042 >>> n = 3
1043 >>> padding = [None] * (n - 1)
1044 >>> list(windowed(chain(padding, iterable), 3))
1045 [(None, None, 1), (None, 1, 2), (1, 2, 3), (2, 3, 4)]
1046 """
1047 if n <= 0:
1048 raise ValueError('n must be > 0')
1049 if step < 1:
1050 raise ValueError('step must be >= 1')
1052 iterator = iter(seq)
1054 # Generate first window
1055 window = deque(islice(iterator, n), maxlen=n)
1057 # Deal with the first window not being full
1058 if not window:
1059 return
1060 if len(window) < n:
1061 yield tuple(window) + ((fillvalue,) * (n - len(window)))
1062 return
1063 yield tuple(window)
1065 # Create the filler for the next windows. The padding ensures
1066 # we have just enough elements to fill the last window.
1067 padding = (fillvalue,) * (n - 1 if step >= n else step - 1)
1068 filler = map(window.append, chain(iterator, padding))
1070 # Generate the rest of the windows
1071 for _ in islice(filler, step - 1, None, step):
1072 yield tuple(window)
1075def substrings(iterable):
1076 """Yield all of the substrings of *iterable*.
1078 >>> [''.join(s) for s in substrings('more')]
1079 ['m', 'o', 'r', 'e', 'mo', 'or', 're', 'mor', 'ore', 'more']
1081 Note that non-string iterables can also be subdivided.
1083 >>> list(substrings([0, 1, 2]))
1084 [(0,), (1,), (2,), (0, 1), (1, 2), (0, 1, 2)]
1086 Like subslices() but returns tuples instead of lists
1087 and returns the shortest substrings first.
1089 """
1090 seq = tuple(iterable)
1091 item_count = len(seq)
1092 for n in range(1, item_count + 1):
1093 slices = map(slice, range(item_count), range(n, item_count + 1))
1094 yield from map(getitem, repeat(seq), slices)
1097def substrings_indexes(seq, reverse=False):
1098 """Yield all substrings and their positions in *seq*
1100 The items yielded will be a tuple of the form ``(substr, i, j)``, where
1101 ``substr == seq[i:j]``.
1103 This function only works for iterables that support slicing, such as
1104 ``str`` objects.
1106 >>> for item in substrings_indexes('more'):
1107 ... print(item)
1108 ('m', 0, 1)
1109 ('o', 1, 2)
1110 ('r', 2, 3)
1111 ('e', 3, 4)
1112 ('mo', 0, 2)
1113 ('or', 1, 3)
1114 ('re', 2, 4)
1115 ('mor', 0, 3)
1116 ('ore', 1, 4)
1117 ('more', 0, 4)
1119 Set *reverse* to ``True`` to yield the same items in the opposite order.
1122 """
1123 r = range(1, len(seq) + 1)
1124 if reverse:
1125 r = reversed(r)
1126 return (
1127 (seq[i : i + L], i, i + L) for L in r for i in range(len(seq) - L + 1)
1128 )
1131class bucket:
1132 """Wrap *iterable* and return an object that buckets the iterable into
1133 child iterables based on a *key* function.
1135 >>> iterable = ['a1', 'b1', 'c1', 'a2', 'b2', 'c2', 'b3']
1136 >>> s = bucket(iterable, key=lambda x: x[0]) # Bucket by 1st character
1137 >>> sorted(list(s)) # Get the keys
1138 ['a', 'b', 'c']
1139 >>> a_iterable = s['a']
1140 >>> next(a_iterable)
1141 'a1'
1142 >>> next(a_iterable)
1143 'a2'
1144 >>> list(s['b'])
1145 ['b1', 'b2', 'b3']
1147 The original iterable will be advanced and its items will be cached until
1148 they are used by the child iterables. This may require significant storage.
1150 By default, attempting to select a bucket to which no items belong will
1151 exhaust the iterable and cache all values.
1152 If you specify a *validator* function, selected buckets will instead be
1153 checked against it.
1155 >>> from itertools import count
1156 >>> it = count(1, 2) # Infinite sequence of odd numbers
1157 >>> key = lambda x: x % 10 # Bucket by last digit
1158 >>> validator = lambda x: x in {1, 3, 5, 7, 9} # Odd digits only
1159 >>> s = bucket(it, key=key, validator=validator)
1160 >>> 2 in s
1161 False
1162 >>> list(s[2])
1163 []
1165 .. seealso:: :func:`map_reduce`, :func:`groupby_transform`
1167 If storage is not a concern, :func:`map_reduce` returns a Python
1168 dictionary, which is generally easier to work with. If the elements
1169 with the same key are already adjacent, :func:`groupby_transform`
1170 or :func:`itertools.groupby` can be used without any caching overhead.
1172 """
1174 def __init__(self, iterable, key, validator=None):
1175 self._it = iter(iterable)
1176 self._key = key
1177 self._cache = defaultdict(deque)
1178 self._validator = validator or (lambda x: True)
1180 def __contains__(self, value):
1181 if not self._validator(value):
1182 return False
1184 try:
1185 item = next(self[value])
1186 except StopIteration:
1187 return False
1188 else:
1189 self._cache[value].appendleft(item)
1191 return True
1193 def _get_values(self, value):
1194 """
1195 Helper to yield items from the parent iterator that match *value*.
1196 Items that don't match are stored in the local cache as they
1197 are encountered.
1198 """
1199 while True:
1200 # If we've cached some items that match the target value, emit
1201 # the first one and evict it from the cache.
1202 if self._cache[value]:
1203 yield self._cache[value].popleft()
1204 # Otherwise we need to advance the parent iterator to search for
1205 # a matching item, caching the rest.
1206 else:
1207 while True:
1208 try:
1209 item = next(self._it)
1210 except StopIteration:
1211 return
1212 item_value = self._key(item)
1213 if item_value == value:
1214 yield item
1215 break
1216 elif self._validator(item_value):
1217 self._cache[item_value].append(item)
1219 def __iter__(self):
1220 for item in self._it:
1221 item_value = self._key(item)
1222 if self._validator(item_value):
1223 self._cache[item_value].append(item)
1225 return iter(self._cache)
1227 def __getitem__(self, value):
1228 if not self._validator(value):
1229 return iter(())
1231 return self._get_values(value)
1234def spy(iterable, n=1):
1235 """Return a 2-tuple with a list containing the first *n* elements of
1236 *iterable*, and an iterator with the same items as *iterable*.
1237 This allows you to "look ahead" at the items in the iterable without
1238 advancing it.
1240 There is one item in the list by default:
1242 >>> iterable = 'abcdefg'
1243 >>> head, iterable = spy(iterable)
1244 >>> head
1245 ['a']
1246 >>> list(iterable)
1247 ['a', 'b', 'c', 'd', 'e', 'f', 'g']
1249 You may use unpacking to retrieve items instead of lists:
1251 >>> (head,), iterable = spy('abcdefg')
1252 >>> head
1253 'a'
1254 >>> (first, second), iterable = spy('abcdefg', 2)
1255 >>> first
1256 'a'
1257 >>> second
1258 'b'
1260 The number of items requested can be larger than the number of items in
1261 the iterable:
1263 >>> iterable = [1, 2, 3, 4, 5]
1264 >>> head, iterable = spy(iterable, 10)
1265 >>> head
1266 [1, 2, 3, 4, 5]
1267 >>> list(iterable)
1268 [1, 2, 3, 4, 5]
1270 """
1271 p, q = tee(iterable)
1272 return take(n, q), p
1275def interleave(*iterables):
1276 """Return a new iterable yielding from each iterable in turn,
1277 until the shortest is exhausted.
1279 >>> list(interleave([1, 2, 3], [4, 5], [6, 7, 8]))
1280 [1, 4, 6, 2, 5, 7]
1282 For a version that doesn't terminate after the shortest iterable is
1283 exhausted, see :func:`interleave_longest`.
1285 """
1286 return chain.from_iterable(zip(*iterables))
1289def interleave_longest(*iterables):
1290 """Return a new iterable yielding from each iterable in turn,
1291 skipping any that are exhausted.
1293 >>> list(interleave_longest([1, 2, 3], [4, 5], [6, 7, 8]))
1294 [1, 4, 6, 2, 5, 7, 3, 8]
1296 This function produces the same output as :func:`roundrobin`, but may
1297 perform better for some inputs (in particular when the number of iterables
1298 is large).
1300 """
1301 for xs in zip_longest(*iterables, fillvalue=_marker):
1302 for x in xs:
1303 if x is not _marker:
1304 yield x
1307def interleave_evenly(iterables, lengths=None):
1308 """
1309 Interleave multiple iterables so that their elements are evenly distributed
1310 throughout the output sequence.
1312 >>> iterables = [1, 2, 3, 4, 5], ['a', 'b']
1313 >>> list(interleave_evenly(iterables))
1314 [1, 2, 'a', 3, 4, 'b', 5]
1316 >>> iterables = [[1, 2, 3], [4, 5], [6, 7, 8]]
1317 >>> list(interleave_evenly(iterables))
1318 [1, 6, 4, 2, 7, 3, 8, 5]
1320 This function requires iterables of known length. Iterables without
1321 ``__len__()`` can be used by manually specifying lengths with *lengths*:
1323 >>> from itertools import combinations, repeat
1324 >>> iterables = [combinations(range(4), 2), ['a', 'b', 'c']]
1325 >>> lengths = [4 * (4 - 1) // 2, 3]
1326 >>> list(interleave_evenly(iterables, lengths=lengths))
1327 [(0, 1), (0, 2), 'a', (0, 3), (1, 2), 'b', (1, 3), (2, 3), 'c']
1329 Based on Bresenham's algorithm.
1330 """
1331 if lengths is None:
1332 try:
1333 lengths = [len(it) for it in iterables]
1334 except TypeError:
1335 raise ValueError(
1336 'Iterable lengths could not be determined automatically. '
1337 'Specify them with the lengths keyword.'
1338 )
1339 elif len(iterables) != len(lengths):
1340 raise ValueError('Mismatching number of iterables and lengths.')
1342 dims = len(lengths)
1344 if not dims:
1345 return
1347 # sort iterables by length, descending
1348 lengths_permute = sorted(
1349 range(dims), key=lambda i: lengths[i], reverse=True
1350 )
1351 lengths_desc = [lengths[i] for i in lengths_permute]
1352 iters_desc = [iter(iterables[i]) for i in lengths_permute]
1354 # the longest iterable is the primary one (Bresenham: the longest
1355 # distance along an axis)
1356 delta_primary, deltas_secondary = lengths_desc[0], lengths_desc[1:]
1357 iter_primary, iters_secondary = iters_desc[0], iters_desc[1:]
1358 errors = [delta_primary // dims] * len(deltas_secondary)
1360 to_yield = sum(lengths)
1361 while to_yield:
1362 yield next(iter_primary)
1363 to_yield -= 1
1364 # update errors for each secondary iterable
1365 errors = [e - delta for e, delta in zip(errors, deltas_secondary)]
1367 # those iterables for which the error is negative are yielded
1368 # ("diagonal step" in Bresenham)
1369 for i, e_ in enumerate(errors):
1370 if e_ < 0:
1371 yield next(iters_secondary[i])
1372 to_yield -= 1
1373 errors[i] += delta_primary
1376def interleave_randomly(*iterables):
1377 """Repeatedly select one of the input *iterables* at random and yield the next
1378 item from it.
1380 >>> iterables = [1, 2, 3], 'abc', (True, False, None)
1381 >>> list(interleave_randomly(*iterables)) # doctest: +SKIP
1382 ['a', 'b', 1, 'c', True, False, None, 2, 3]
1384 The relative order of the items in each input iterable will preserved. Note the
1385 sequences of items with this property are not equally likely to be generated.
1387 """
1388 iterators = [iter(e) for e in iterables]
1389 while iterators:
1390 idx = randrange(len(iterators))
1391 try:
1392 yield next(iterators[idx])
1393 except StopIteration:
1394 # equivalent to `list.pop` but slightly faster
1395 iterators[idx] = iterators[-1]
1396 del iterators[-1]
1399def collapse(iterable, base_type=None, levels=None):
1400 """Flatten an iterable with multiple levels of nesting (e.g., a list of
1401 lists of tuples) into non-iterable types.
1403 >>> iterable = [(1, 2), ([3, 4], [[5], [6]])]
1404 >>> list(collapse(iterable))
1405 [1, 2, 3, 4, 5, 6]
1407 Binary and text strings are not considered iterable and
1408 will not be collapsed.
1410 To avoid collapsing other types, specify *base_type*:
1412 >>> iterable = ['ab', ('cd', 'ef'), ['gh', 'ij']]
1413 >>> list(collapse(iterable, base_type=tuple))
1414 ['ab', ('cd', 'ef'), 'gh', 'ij']
1416 Specify *levels* to stop flattening after a certain level:
1418 >>> iterable = [('a', ['b']), ('c', ['d'])]
1419 >>> list(collapse(iterable)) # Fully flattened
1420 ['a', 'b', 'c', 'd']
1421 >>> list(collapse(iterable, levels=1)) # Only one level flattened
1422 ['a', ['b'], 'c', ['d']]
1424 """
1425 stack = deque()
1426 # Add our first node group, treat the iterable as a single node
1427 stack.appendleft((0, repeat(iterable, 1)))
1429 while stack:
1430 node_group = stack.popleft()
1431 level, nodes = node_group
1433 # Check if beyond max level
1434 if levels is not None and level > levels:
1435 yield from nodes
1436 continue
1438 for node in nodes:
1439 # Check if done iterating
1440 if isinstance(node, (str, bytes)) or (
1441 (base_type is not None) and isinstance(node, base_type)
1442 ):
1443 yield node
1444 # Otherwise try to create child nodes
1445 else:
1446 try:
1447 tree = iter(node)
1448 except TypeError:
1449 yield node
1450 else:
1451 # Save our current location
1452 stack.appendleft(node_group)
1453 # Append the new child node
1454 stack.appendleft((level + 1, tree))
1455 # Break to process child node
1456 break
1459def side_effect(func, iterable, chunk_size=None, before=None, after=None):
1460 """Invoke *func* on each item in *iterable* (or on each *chunk_size* group
1461 of items) before yielding the item.
1463 `func` must be a function that takes a single argument. Its return value
1464 will be discarded.
1466 *before* and *after* are optional functions that take no arguments. They
1467 will be executed before iteration starts and after it ends, respectively.
1469 `side_effect` can be used for logging, updating progress bars, or anything
1470 that is not functionally "pure."
1472 Emitting a status message:
1474 >>> from more_itertools import consume
1475 >>> func = lambda item: print('Received {}'.format(item))
1476 >>> consume(side_effect(func, range(2)))
1477 Received 0
1478 Received 1
1480 Operating on chunks of items:
1482 >>> pair_sums = []
1483 >>> func = lambda chunk: pair_sums.append(sum(chunk))
1484 >>> list(side_effect(func, [0, 1, 2, 3, 4, 5], 2))
1485 [0, 1, 2, 3, 4, 5]
1486 >>> list(pair_sums)
1487 [1, 5, 9]
1489 Writing to a file-like object:
1491 >>> from io import StringIO
1492 >>> from more_itertools import consume
1493 >>> f = StringIO()
1494 >>> func = lambda x: print(x, file=f)
1495 >>> before = lambda: print('HEADER', file=f)
1496 >>> after = f.close
1497 >>> it = ['a', 'b', 'c']
1498 >>> consume(side_effect(func, it, before=before, after=after))
1499 >>> f.closed
1500 True
1502 """
1503 try:
1504 if before is not None:
1505 before()
1507 if chunk_size is None:
1508 for item in iterable:
1509 func(item)
1510 yield item
1511 else:
1512 for chunk in chunked(iterable, chunk_size):
1513 func(chunk)
1514 yield from chunk
1515 finally:
1516 if after is not None:
1517 after()
1520def sliced(seq, n, strict=False):
1521 """Yield slices of length *n* from the sequence *seq*.
1523 >>> list(sliced((1, 2, 3, 4, 5, 6), 3))
1524 [(1, 2, 3), (4, 5, 6)]
1526 By the default, the last yielded slice will have fewer than *n* elements
1527 if the length of *seq* is not divisible by *n*:
1529 >>> list(sliced((1, 2, 3, 4, 5, 6, 7, 8), 3))
1530 [(1, 2, 3), (4, 5, 6), (7, 8)]
1532 If the length of *seq* is not divisible by *n* and *strict* is
1533 ``True``, then ``ValueError`` will be raised before the last
1534 slice is yielded.
1536 This function will only work for iterables that support slicing.
1537 For non-sliceable iterables, see :func:`chunked`.
1539 """
1540 if n < 0:
1541 raise ValueError('n must be at least 0')
1543 iterator = takewhile(len, (seq[i : i + n] for i in count(0, n)))
1544 if strict:
1546 def ret():
1547 for _slice in iterator:
1548 if len(_slice) != n:
1549 raise ValueError("seq is not divisible by n.")
1550 yield _slice
1552 return ret()
1553 else:
1554 return iterator
1557def split_at(iterable, pred, maxsplit=-1, keep_separator=False):
1558 """Yield lists of items from *iterable*, where each list is delimited by
1559 an item where callable *pred* returns ``True``.
1561 >>> list(split_at('abcdcba', lambda x: x == 'b'))
1562 [['a'], ['c', 'd', 'c'], ['a']]
1564 >>> list(split_at(range(10), lambda n: n % 2 == 1))
1565 [[0], [2], [4], [6], [8], []]
1567 At most *maxsplit* splits are done. If *maxsplit* is not specified or -1,
1568 then there is no limit on the number of splits:
1570 >>> list(split_at(range(10), lambda n: n % 2 == 1, maxsplit=2))
1571 [[0], [2], [4, 5, 6, 7, 8, 9]]
1573 By default, the delimiting items are not included in the output.
1574 To include them, set *keep_separator* to ``True``.
1576 >>> list(split_at('abcdcba', lambda x: x == 'b', keep_separator=True))
1577 [['a'], ['b'], ['c', 'd', 'c'], ['b'], ['a']]
1579 """
1580 if maxsplit == 0:
1581 yield list(iterable)
1582 return
1584 buf = []
1585 it = iter(iterable)
1586 for item in it:
1587 if pred(item):
1588 yield buf
1589 if keep_separator:
1590 yield [item]
1591 if maxsplit == 1:
1592 yield list(it)
1593 return
1594 buf = []
1595 maxsplit -= 1
1596 else:
1597 buf.append(item)
1598 yield buf
1601def split_before(iterable, pred, maxsplit=-1):
1602 """Yield lists of items from *iterable*, where each list ends just before
1603 an item for which callable *pred* returns ``True``:
1605 >>> list(split_before('OneTwo', lambda s: s.isupper()))
1606 [['O', 'n', 'e'], ['T', 'w', 'o']]
1608 >>> list(split_before(range(10), lambda n: n % 3 == 0))
1609 [[0, 1, 2], [3, 4, 5], [6, 7, 8], [9]]
1611 At most *maxsplit* splits are done. If *maxsplit* is not specified or -1,
1612 then there is no limit on the number of splits:
1614 >>> list(split_before(range(10), lambda n: n % 3 == 0, maxsplit=2))
1615 [[0, 1, 2], [3, 4, 5], [6, 7, 8, 9]]
1616 """
1617 if maxsplit == 0:
1618 yield list(iterable)
1619 return
1621 buf = []
1622 it = iter(iterable)
1623 for item in it:
1624 if pred(item) and buf:
1625 yield buf
1626 if maxsplit == 1:
1627 yield [item, *it]
1628 return
1629 buf = []
1630 maxsplit -= 1
1631 buf.append(item)
1632 if buf:
1633 yield buf
1636def split_after(iterable, pred, maxsplit=-1):
1637 """Yield lists of items from *iterable*, where each list ends with an
1638 item where callable *pred* returns ``True``:
1640 >>> list(split_after('one1two2', lambda s: s.isdigit()))
1641 [['o', 'n', 'e', '1'], ['t', 'w', 'o', '2']]
1643 >>> list(split_after(range(10), lambda n: n % 3 == 0))
1644 [[0], [1, 2, 3], [4, 5, 6], [7, 8, 9]]
1646 At most *maxsplit* splits are done. If *maxsplit* is not specified or -1,
1647 then there is no limit on the number of splits:
1649 >>> list(split_after(range(10), lambda n: n % 3 == 0, maxsplit=2))
1650 [[0], [1, 2, 3], [4, 5, 6, 7, 8, 9]]
1652 """
1653 if maxsplit == 0:
1654 yield list(iterable)
1655 return
1657 buf = []
1658 it = iter(iterable)
1659 for item in it:
1660 buf.append(item)
1661 if pred(item) and buf:
1662 yield buf
1663 if maxsplit == 1:
1664 buf = list(it)
1665 if buf:
1666 yield buf
1667 return
1668 buf = []
1669 maxsplit -= 1
1670 if buf:
1671 yield buf
1674def split_when(iterable, pred, maxsplit=-1):
1675 """Split *iterable* into pieces based on the output of *pred*.
1676 *pred* should be a function that takes successive pairs of items and
1677 returns ``True`` if the iterable should be split in between them.
1679 For example, to find runs of increasing numbers, split the iterable when
1680 element ``i`` is larger than element ``i + 1``:
1682 >>> list(split_when([1, 2, 3, 3, 2, 5, 2, 4, 2], lambda x, y: x > y))
1683 [[1, 2, 3, 3], [2, 5], [2, 4], [2]]
1685 At most *maxsplit* splits are done. If *maxsplit* is not specified or -1,
1686 then there is no limit on the number of splits:
1688 >>> list(split_when([1, 2, 3, 3, 2, 5, 2, 4, 2],
1689 ... lambda x, y: x > y, maxsplit=2))
1690 [[1, 2, 3, 3], [2, 5], [2, 4, 2]]
1692 """
1693 if maxsplit == 0:
1694 yield list(iterable)
1695 return
1697 it = iter(iterable)
1698 try:
1699 cur_item = next(it)
1700 except StopIteration:
1701 return
1703 buf = [cur_item]
1704 for next_item in it:
1705 if pred(cur_item, next_item):
1706 yield buf
1707 if maxsplit == 1:
1708 yield [next_item, *it]
1709 return
1710 buf = []
1711 maxsplit -= 1
1713 buf.append(next_item)
1714 cur_item = next_item
1716 yield buf
1719def split_into(iterable, sizes):
1720 """Yield a list of sequential items from *iterable* of length 'n' for each
1721 integer 'n' in *sizes*.
1723 >>> list(split_into([1,2,3,4,5,6], [1,2,3]))
1724 [[1], [2, 3], [4, 5, 6]]
1726 If the sum of *sizes* is smaller than the length of *iterable*, then the
1727 remaining items of *iterable* will not be returned.
1729 >>> list(split_into([1,2,3,4,5,6], [2,3]))
1730 [[1, 2], [3, 4, 5]]
1732 If the sum of *sizes* is larger than the length of *iterable*, fewer items
1733 will be returned in the iteration that overruns the *iterable* and further
1734 lists will be empty:
1736 >>> list(split_into([1,2,3,4], [1,2,3,4]))
1737 [[1], [2, 3], [4], []]
1739 When a ``None`` object is encountered in *sizes*, the returned list will
1740 contain items up to the end of *iterable* the same way that
1741 :func:`itertools.slice` does:
1743 >>> list(split_into([1,2,3,4,5,6,7,8,9,0], [2,3,None]))
1744 [[1, 2], [3, 4, 5], [6, 7, 8, 9, 0]]
1746 :func:`split_into` can be useful for grouping a series of items where the
1747 sizes of the groups are not uniform. An example would be where in a row
1748 from a table, multiple columns represent elements of the same feature
1749 (e.g. a point represented by x,y,z) but, the format is not the same for
1750 all columns.
1751 """
1752 # convert the iterable argument into an iterator so its contents can
1753 # be consumed by islice in case it is a generator
1754 it = iter(iterable)
1756 for size in sizes:
1757 if size is None:
1758 yield list(it)
1759 return
1760 else:
1761 yield list(islice(it, size))
1764def padded(iterable, fillvalue=None, n=None, next_multiple=False):
1765 """Yield the elements from *iterable*, followed by *fillvalue*, such that
1766 at least *n* items are emitted.
1768 >>> list(padded([1, 2, 3], '?', 5))
1769 [1, 2, 3, '?', '?']
1771 If *next_multiple* is ``True``, *fillvalue* will be emitted until the
1772 number of items emitted is a multiple of *n*:
1774 >>> list(padded([1, 2, 3, 4], n=3, next_multiple=True))
1775 [1, 2, 3, 4, None, None]
1777 If *n* is ``None``, *fillvalue* will be emitted indefinitely.
1779 To create an *iterable* of exactly size *n*, you can truncate with
1780 :func:`islice`.
1782 >>> list(islice(padded([1, 2, 3], '?'), 5))
1783 [1, 2, 3, '?', '?']
1784 >>> list(islice(padded([1, 2, 3, 4, 5, 6, 7, 8], '?'), 5))
1785 [1, 2, 3, 4, 5]
1787 """
1788 iterator = iter(iterable)
1789 iterator_with_repeat = chain(iterator, repeat(fillvalue))
1791 if n is None:
1792 return iterator_with_repeat
1793 elif n < 1:
1794 raise ValueError('n must be at least 1')
1795 elif next_multiple:
1797 def slice_generator():
1798 for first in iterator:
1799 yield (first,)
1800 yield islice(iterator_with_repeat, n - 1)
1802 # While elements exist produce slices of size n
1803 return chain.from_iterable(slice_generator())
1804 else:
1805 # Ensure the first batch is at least size n then iterate
1806 return chain(islice(iterator_with_repeat, n), iterator)
1809def repeat_each(iterable, n=2):
1810 """Repeat each element in *iterable* *n* times.
1812 >>> list(repeat_each('ABC', 3))
1813 ['A', 'A', 'A', 'B', 'B', 'B', 'C', 'C', 'C']
1814 """
1815 return chain.from_iterable(map(repeat, iterable, repeat(n)))
1818def repeat_last(iterable, default=None):
1819 """After the *iterable* is exhausted, keep yielding its last element.
1821 >>> list(islice(repeat_last(range(3)), 5))
1822 [0, 1, 2, 2, 2]
1824 If the iterable is empty, yield *default* forever::
1826 >>> list(islice(repeat_last(range(0), 42), 5))
1827 [42, 42, 42, 42, 42]
1829 """
1830 item = _marker
1831 for item in iterable:
1832 yield item
1833 final = default if item is _marker else item
1834 yield from repeat(final)
1837def distribute(n, iterable):
1838 """Distribute the items from *iterable* among *n* smaller iterables.
1840 >>> group_1, group_2 = distribute(2, [1, 2, 3, 4, 5, 6])
1841 >>> list(group_1)
1842 [1, 3, 5]
1843 >>> list(group_2)
1844 [2, 4, 6]
1846 If the length of *iterable* is not evenly divisible by *n*, then the
1847 length of the returned iterables will not be identical:
1849 >>> children = distribute(3, [1, 2, 3, 4, 5, 6, 7])
1850 >>> [list(c) for c in children]
1851 [[1, 4, 7], [2, 5], [3, 6]]
1853 If the length of *iterable* is smaller than *n*, then the last returned
1854 iterables will be empty:
1856 >>> children = distribute(5, [1, 2, 3])
1857 >>> [list(c) for c in children]
1858 [[1], [2], [3], [], []]
1860 This function uses :func:`itertools.tee` and may require significant
1861 storage.
1863 If you need the order items in the smaller iterables to match the
1864 original iterable, see :func:`divide`.
1866 """
1867 if n < 1:
1868 raise ValueError('n must be at least 1')
1870 children = tee(iterable, n)
1871 return [islice(it, index, None, n) for index, it in enumerate(children)]
1874def stagger(iterable, offsets=(-1, 0, 1), longest=False, fillvalue=None):
1875 """Yield tuples whose elements are offset from *iterable*.
1876 The amount by which the `i`-th item in each tuple is offset is given by
1877 the `i`-th item in *offsets*.
1879 >>> list(stagger([0, 1, 2, 3]))
1880 [(None, 0, 1), (0, 1, 2), (1, 2, 3)]
1881 >>> list(stagger(range(8), offsets=(0, 2, 4)))
1882 [(0, 2, 4), (1, 3, 5), (2, 4, 6), (3, 5, 7)]
1884 By default, the sequence will end when the final element of a tuple is the
1885 last item in the iterable. To continue until the first element of a tuple
1886 is the last item in the iterable, set *longest* to ``True``::
1888 >>> list(stagger([0, 1, 2, 3], longest=True))
1889 [(None, 0, 1), (0, 1, 2), (1, 2, 3), (2, 3, None), (3, None, None)]
1891 By default, ``None`` will be used to replace offsets beyond the end of the
1892 sequence. Specify *fillvalue* to use some other value.
1894 """
1895 children = tee(iterable, len(offsets))
1897 return zip_offset(
1898 *children, offsets=offsets, longest=longest, fillvalue=fillvalue
1899 )
1902def zip_offset(*iterables, offsets, longest=False, fillvalue=None):
1903 """``zip`` the input *iterables* together, but offset the `i`-th iterable
1904 by the `i`-th item in *offsets*.
1906 >>> list(zip_offset('0123', 'abcdef', offsets=(0, 1)))
1907 [('0', 'b'), ('1', 'c'), ('2', 'd'), ('3', 'e')]
1909 This can be used as a lightweight alternative to SciPy or pandas to analyze
1910 data sets in which some series have a lead or lag relationship.
1912 By default, the sequence will end when the shortest iterable is exhausted.
1913 To continue until the longest iterable is exhausted, set *longest* to
1914 ``True``.
1916 >>> list(zip_offset('0123', 'abcdef', offsets=(0, 1), longest=True))
1917 [('0', 'b'), ('1', 'c'), ('2', 'd'), ('3', 'e'), (None, 'f')]
1919 By default, ``None`` will be used to replace offsets beyond the end of the
1920 sequence. Specify *fillvalue* to use some other value.
1922 """
1923 if len(iterables) != len(offsets):
1924 raise ValueError("Number of iterables and offsets didn't match")
1926 staggered = []
1927 for it, n in zip(iterables, offsets):
1928 if n < 0:
1929 staggered.append(chain(repeat(fillvalue, -n), it))
1930 elif n > 0:
1931 staggered.append(islice(it, n, None))
1932 else:
1933 staggered.append(it)
1935 if longest:
1936 return zip_longest(*staggered, fillvalue=fillvalue)
1938 return zip(*staggered)
1941def sort_together(
1942 iterables, key_list=(0,), key=None, reverse=False, strict=False
1943):
1944 """Return the input iterables sorted together, with *key_list* as the
1945 priority for sorting. All iterables are trimmed to the length of the
1946 shortest one.
1948 This can be used like the sorting function in a spreadsheet. If each
1949 iterable represents a column of data, the key list determines which
1950 columns are used for sorting.
1952 By default, all iterables are sorted using the ``0``-th iterable::
1954 >>> iterables = [(4, 3, 2, 1), ('a', 'b', 'c', 'd')]
1955 >>> sort_together(iterables)
1956 [(1, 2, 3, 4), ('d', 'c', 'b', 'a')]
1958 Set a different key list to sort according to another iterable.
1959 Specifying multiple keys dictates how ties are broken::
1961 >>> iterables = [(3, 1, 2), (0, 1, 0), ('c', 'b', 'a')]
1962 >>> sort_together(iterables, key_list=(1, 2))
1963 [(2, 3, 1), (0, 0, 1), ('a', 'c', 'b')]
1965 To sort by a function of the elements of the iterable, pass a *key*
1966 function. Its arguments are the elements of the iterables corresponding to
1967 the key list::
1969 >>> names = ('a', 'b', 'c')
1970 >>> lengths = (1, 2, 3)
1971 >>> widths = (5, 2, 1)
1972 >>> def area(length, width):
1973 ... return length * width
1974 >>> sort_together([names, lengths, widths], key_list=(1, 2), key=area)
1975 [('c', 'b', 'a'), (3, 2, 1), (1, 2, 5)]
1977 Set *reverse* to ``True`` to sort in descending order.
1979 >>> sort_together([(1, 2, 3), ('c', 'b', 'a')], reverse=True)
1980 [(3, 2, 1), ('a', 'b', 'c')]
1982 If the *strict* keyword argument is ``True``, then
1983 ``ValueError`` will be raised if any of the iterables have
1984 different lengths.
1986 """
1987 if key is None:
1988 # if there is no key function, the key argument to sorted is an
1989 # itemgetter
1990 key_argument = itemgetter(*key_list)
1991 else:
1992 # if there is a key function, call it with the items at the offsets
1993 # specified by the key function as arguments
1994 key_list = list(key_list)
1995 if len(key_list) == 1:
1996 # if key_list contains a single item, pass the item at that offset
1997 # as the only argument to the key function
1998 key_offset = key_list[0]
1999 key_argument = lambda zipped_items: key(zipped_items[key_offset])
2000 else:
2001 # if key_list contains multiple items, use itemgetter to return a
2002 # tuple of items, which we pass as *args to the key function
2003 get_key_items = itemgetter(*key_list)
2004 key_argument = lambda zipped_items: key(
2005 *get_key_items(zipped_items)
2006 )
2008 transposed = zip(*iterables, strict=strict)
2009 reordered = sorted(transposed, key=key_argument, reverse=reverse)
2010 untransposed = zip(*reordered, strict=strict)
2011 return list(untransposed)
2014def unzip(iterable):
2015 """The inverse of :func:`zip`, this function disaggregates the elements
2016 of the zipped *iterable*.
2018 The ``i``-th iterable contains the ``i``-th element from each element
2019 of the zipped iterable. The first element is used to determine the
2020 length of the remaining elements.
2022 >>> iterable = [('a', 1), ('b', 2), ('c', 3), ('d', 4)]
2023 >>> letters, numbers = unzip(iterable)
2024 >>> list(letters)
2025 ['a', 'b', 'c', 'd']
2026 >>> list(numbers)
2027 [1, 2, 3, 4]
2029 This is similar to using ``zip(*iterable)``, but it avoids reading
2030 *iterable* into memory. Note, however, that this function uses
2031 :func:`itertools.tee` and thus may require significant storage.
2033 """
2034 head, iterable = spy(iterable)
2035 if not head:
2036 # empty iterable, e.g. zip([], [], [])
2037 return ()
2038 # spy returns a one-length iterable as head
2039 head = head[0]
2040 iterables = tee(iterable, len(head))
2042 # If we have an iterable like iter([(1, 2, 3), (4, 5), (6,)]),
2043 # the second unzipped iterable fails at the third tuple since
2044 # it tries to access (6,)[1].
2045 # Same with the third unzipped iterable and the second tuple.
2046 # To support these "improperly zipped" iterables, we suppress
2047 # the IndexError, which just stops the unzipped iterables at
2048 # first length mismatch.
2049 return tuple(
2050 iter_suppress(map(itemgetter(i), it), IndexError)
2051 for i, it in enumerate(iterables)
2052 )
2055def divide(n, iterable):
2056 """Divide the elements from *iterable* into *n* parts, maintaining
2057 order.
2059 >>> group_1, group_2 = divide(2, [1, 2, 3, 4, 5, 6])
2060 >>> list(group_1)
2061 [1, 2, 3]
2062 >>> list(group_2)
2063 [4, 5, 6]
2065 If the length of *iterable* is not evenly divisible by *n*, then the
2066 length of the returned iterables will not be identical:
2068 >>> children = divide(3, [1, 2, 3, 4, 5, 6, 7])
2069 >>> [list(c) for c in children]
2070 [[1, 2, 3], [4, 5], [6, 7]]
2072 If the length of the iterable is smaller than n, then the last returned
2073 iterables will be empty:
2075 >>> children = divide(5, [1, 2, 3])
2076 >>> [list(c) for c in children]
2077 [[1], [2], [3], [], []]
2079 This function will exhaust the iterable before returning.
2080 If order is not important, see :func:`distribute`, which does not first
2081 pull the iterable into memory.
2083 """
2084 if n < 1:
2085 raise ValueError('n must be at least 1')
2087 try:
2088 iterable[:0]
2089 except TypeError:
2090 seq = tuple(iterable)
2091 else:
2092 seq = iterable
2094 q, r = divmod(len(seq), n)
2096 ret = []
2097 stop = 0
2098 for i in range(1, n + 1):
2099 start = stop
2100 stop += q + 1 if i <= r else q
2101 ret.append(iter(seq[start:stop]))
2103 return ret
2106def always_iterable(obj, base_type=(str, bytes)):
2107 """If *obj* is iterable, return an iterator over its items::
2109 >>> obj = (1, 2, 3)
2110 >>> list(always_iterable(obj))
2111 [1, 2, 3]
2113 If *obj* is not iterable, return a one-item iterable containing *obj*::
2115 >>> obj = 1
2116 >>> list(always_iterable(obj))
2117 [1]
2119 If *obj* is ``None``, return an empty iterable:
2121 >>> obj = None
2122 >>> list(always_iterable(None))
2123 []
2125 By default, binary and text strings are not considered iterable::
2127 >>> obj = 'foo'
2128 >>> list(always_iterable(obj))
2129 ['foo']
2131 If *base_type* is set, objects for which ``isinstance(obj, base_type)``
2132 returns ``True`` won't be considered iterable.
2134 >>> obj = {'a': 1}
2135 >>> list(always_iterable(obj)) # Iterate over the dict's keys
2136 ['a']
2137 >>> list(always_iterable(obj, base_type=dict)) # Treat dicts as a unit
2138 [{'a': 1}]
2140 Set *base_type* to ``None`` to avoid any special handling and treat objects
2141 Python considers iterable as iterable:
2143 >>> obj = 'foo'
2144 >>> list(always_iterable(obj, base_type=None))
2145 ['f', 'o', 'o']
2146 """
2147 if obj is None:
2148 return iter(())
2150 if (base_type is not None) and isinstance(obj, base_type):
2151 return iter((obj,))
2153 try:
2154 return iter(obj)
2155 except TypeError:
2156 return iter((obj,))
2159def adjacent(predicate, iterable, distance=1):
2160 """Return an iterable over `(bool, item)` tuples where the `item` is
2161 drawn from *iterable* and the `bool` indicates whether
2162 that item satisfies the *predicate* or is adjacent to an item that does.
2164 For example, to find whether items are adjacent to a ``3``::
2166 >>> list(adjacent(lambda x: x == 3, range(6)))
2167 [(False, 0), (False, 1), (True, 2), (True, 3), (True, 4), (False, 5)]
2169 Set *distance* to change what counts as adjacent. For example, to find
2170 whether items are two places away from a ``3``:
2172 >>> list(adjacent(lambda x: x == 3, range(6), distance=2))
2173 [(False, 0), (True, 1), (True, 2), (True, 3), (True, 4), (True, 5)]
2175 This is useful for contextualizing the results of a search function.
2176 For example, a code comparison tool might want to identify lines that
2177 have changed, but also surrounding lines to give the viewer of the diff
2178 context.
2180 The predicate function will only be called once for each item in the
2181 iterable.
2183 See also :func:`groupby_transform`, which can be used with this function
2184 to group ranges of items with the same `bool` value.
2186 """
2187 # Allow distance=0 mainly for testing that it reproduces results with map()
2188 if distance < 0:
2189 raise ValueError('distance must be at least 0')
2191 i1, i2 = tee(iterable)
2192 padding = [False] * distance
2193 selected = chain(padding, map(predicate, i1), padding)
2194 adjacent_to_selected = map(any, windowed(selected, 2 * distance + 1))
2195 return zip(adjacent_to_selected, i2)
2198def groupby_transform(iterable, keyfunc=None, valuefunc=None, reducefunc=None):
2199 """An extension of :func:`itertools.groupby` that can apply transformations
2200 to the grouped data.
2202 * *keyfunc* is a function computing a key value for each item in *iterable*
2203 * *valuefunc* is a function that transforms the individual items from
2204 *iterable* after grouping
2205 * *reducefunc* is a function that transforms each group of items
2207 >>> iterable = 'aAAbBBcCC'
2208 >>> keyfunc = lambda k: k.upper()
2209 >>> valuefunc = lambda v: v.lower()
2210 >>> reducefunc = lambda g: ''.join(g)
2211 >>> list(groupby_transform(iterable, keyfunc, valuefunc, reducefunc))
2212 [('A', 'aaa'), ('B', 'bbb'), ('C', 'ccc')]
2214 Each optional argument defaults to an identity function if not specified.
2216 :func:`groupby_transform` is useful when grouping elements of an iterable
2217 using a separate iterable as the key. To do this, :func:`zip` the iterables
2218 and pass a *keyfunc* that extracts the first element and a *valuefunc*
2219 that extracts the second element::
2221 >>> from operator import itemgetter
2222 >>> keys = [0, 0, 1, 1, 1, 2, 2, 2, 3]
2223 >>> values = 'abcdefghi'
2224 >>> iterable = zip(keys, values)
2225 >>> grouper = groupby_transform(iterable, itemgetter(0), itemgetter(1))
2226 >>> [(k, ''.join(g)) for k, g in grouper]
2227 [(0, 'ab'), (1, 'cde'), (2, 'fgh'), (3, 'i')]
2229 Note that the order of items in the iterable is significant.
2230 Only adjacent items are grouped together, so if you don't want any
2231 duplicate groups, you should sort the iterable by the key function
2232 or consider :func:`bucket` or :func:`map_reduce`. :func:`map_reduce`
2233 consumes the iterable immediately and returns a dictionary, while
2234 :func:`bucket` does not.
2236 .. seealso:: :func:`bucket`, :func:`map_reduce`
2238 """
2239 ret = groupby(iterable, keyfunc)
2240 if valuefunc:
2241 ret = ((k, map(valuefunc, g)) for k, g in ret)
2242 if reducefunc:
2243 ret = ((k, reducefunc(g)) for k, g in ret)
2245 return ret
2248class numeric_range(Sequence):
2249 """An extension of the built-in ``range()`` function whose arguments can
2250 be any orderable numeric type.
2252 With only *stop* specified, *start* defaults to ``0`` and *step*
2253 defaults to ``1``. The output items will match the type of *stop*:
2255 >>> list(numeric_range(3.5))
2256 [0.0, 1.0, 2.0, 3.0]
2258 With only *start* and *stop* specified, *step* defaults to ``1``. The
2259 output items will match the type of *start*:
2261 >>> from decimal import Decimal
2262 >>> start = Decimal('2.1')
2263 >>> stop = Decimal('5.1')
2264 >>> list(numeric_range(start, stop))
2265 [Decimal('2.1'), Decimal('3.1'), Decimal('4.1')]
2267 With *start*, *stop*, and *step* specified the output items will match
2268 the type of ``start + step``:
2270 >>> from fractions import Fraction
2271 >>> start = Fraction(1, 2) # Start at 1/2
2272 >>> stop = Fraction(5, 2) # End at 5/2
2273 >>> step = Fraction(1, 2) # Count by 1/2
2274 >>> list(numeric_range(start, stop, step))
2275 [Fraction(1, 2), Fraction(1, 1), Fraction(3, 2), Fraction(2, 1)]
2277 If *step* is zero, ``ValueError`` is raised. Negative steps are supported:
2279 >>> list(numeric_range(3, -1, -1.0))
2280 [3.0, 2.0, 1.0, 0.0]
2282 Be aware of the limitations of floating-point numbers; the representation
2283 of the yielded numbers may be surprising.
2285 ``datetime.datetime`` objects can be used for *start* and *stop*, if *step*
2286 is a ``datetime.timedelta`` object:
2288 >>> import datetime
2289 >>> start = datetime.datetime(2019, 1, 1)
2290 >>> stop = datetime.datetime(2019, 1, 3)
2291 >>> step = datetime.timedelta(days=1)
2292 >>> items = iter(numeric_range(start, stop, step))
2293 >>> next(items)
2294 datetime.datetime(2019, 1, 1, 0, 0)
2295 >>> next(items)
2296 datetime.datetime(2019, 1, 2, 0, 0)
2298 """
2300 _EMPTY_HASH = hash(range(0, 0))
2302 def __init__(self, *args):
2303 argc = len(args)
2304 if argc == 1:
2305 (self._stop,) = args
2306 self._start = type(self._stop)(0)
2307 self._step = type(self._stop - self._start)(1)
2308 elif argc == 2:
2309 self._start, self._stop = args
2310 self._step = type(self._stop - self._start)(1)
2311 elif argc == 3:
2312 self._start, self._stop, self._step = args
2313 elif argc == 0:
2314 raise TypeError(
2315 f'numeric_range expected at least 1 argument, got {argc}'
2316 )
2317 else:
2318 raise TypeError(
2319 f'numeric_range expected at most 3 arguments, got {argc}'
2320 )
2322 self._zero = type(self._step)(0)
2323 if self._step == self._zero:
2324 raise ValueError('numeric_range() arg 3 must not be zero')
2325 self._growing = self._step > self._zero
2327 def __bool__(self):
2328 if self._growing:
2329 return self._start < self._stop
2330 else:
2331 return self._start > self._stop
2333 def __contains__(self, elem):
2334 if self._growing:
2335 if self._start <= elem < self._stop:
2336 return (elem - self._start) % self._step == self._zero
2337 else:
2338 if self._start >= elem > self._stop:
2339 return (self._start - elem) % (-self._step) == self._zero
2341 return False
2343 def __eq__(self, other):
2344 # numeric_range object equality is intended to mirror the built-in range
2345 # object's equality.
2346 # https://github.com/python/cpython/blob/f5c4880151b609e0a0a0b05c292d36b18038c061/Objects/rangeobject.c#L499
2347 if not isinstance(other, numeric_range):
2348 return False
2350 if self is other:
2351 return True
2353 len_self = len(self)
2354 if len_self != len(other):
2355 return False
2357 if not len_self:
2358 return True
2360 if self._start != other._start:
2361 return False
2363 if len_self == 1:
2364 return True
2366 return self._step == other._step
2368 def __getitem__(self, key):
2369 if isinstance(key, int):
2370 return self._get_by_index(key)
2371 elif isinstance(key, slice):
2372 start_idx, stop_idx, step_idx = key.indices(self._len)
2373 return numeric_range(
2374 self._start + start_idx * self._step,
2375 self._start + stop_idx * self._step,
2376 self._step * step_idx,
2377 )
2378 else:
2379 raise TypeError(
2380 'numeric range indices must be '
2381 f'integers or slices, not {type(key).__name__}'
2382 )
2384 def __hash__(self):
2385 # numeric_range hashing is intended to mirror the built-in range object's
2386 # hashing.
2387 # https://github.com/python/cpython/blob/f5c4880151b609e0a0a0b05c292d36b18038c061/Objects/rangeobject.c#L570
2388 len_self = len(self)
2389 if not len_self:
2390 return hash((len_self, None, None))
2391 if len_self == 1:
2392 return hash((len_self, self._start, None))
2393 return hash((len_self, self._start, self._step))
2395 def __iter__(self):
2396 values = (self._start + (n * self._step) for n in count())
2397 if self._growing:
2398 return takewhile(partial(gt, self._stop), values)
2399 else:
2400 return takewhile(partial(lt, self._stop), values)
2402 def __len__(self):
2403 return self._len
2405 @cached_property
2406 def _len(self):
2407 if self._growing:
2408 start = self._start
2409 stop = self._stop
2410 step = self._step
2411 else:
2412 start = self._stop
2413 stop = self._start
2414 step = -self._step
2415 distance = stop - start
2416 if distance <= self._zero:
2417 return 0
2418 else: # distance > 0 and step > 0: regular euclidean division
2419 q, r = divmod(distance, step)
2420 return int(q) + int(r != self._zero)
2422 def __reduce__(self):
2423 return numeric_range, (self._start, self._stop, self._step)
2425 def __repr__(self):
2426 if self._step == 1:
2427 return f"numeric_range({self._start!r}, {self._stop!r})"
2428 return (
2429 f"numeric_range({self._start!r}, {self._stop!r}, {self._step!r})"
2430 )
2432 def __reversed__(self):
2433 # Empty iterator
2434 try:
2435 start = self._get_by_index(-1)
2436 except IndexError:
2437 return iter([])
2439 return iter(
2440 numeric_range(start, self._start - self._step, -self._step)
2441 )
2443 def count(self, value):
2444 return int(value in self)
2446 def index(self, value):
2447 if self._growing:
2448 if self._start <= value < self._stop:
2449 q, r = divmod(value - self._start, self._step)
2450 if r == self._zero:
2451 return int(q)
2452 else:
2453 if self._start >= value > self._stop:
2454 q, r = divmod(self._start - value, -self._step)
2455 if r == self._zero:
2456 return int(q)
2458 raise ValueError(f"{value} is not in numeric range")
2460 def _get_by_index(self, i):
2461 if i < 0:
2462 i += self._len
2463 if i < 0 or i >= self._len:
2464 raise IndexError("numeric range object index out of range")
2465 return self._start + i * self._step
2468def count_cycle(iterable, n=None):
2469 """Cycle through the items from *iterable* up to *n* times, yielding
2470 the number of completed cycles along with each item. If *n* is omitted the
2471 process repeats indefinitely.
2473 >>> list(count_cycle('AB', 3))
2474 [(0, 'A'), (0, 'B'), (1, 'A'), (1, 'B'), (2, 'A'), (2, 'B')]
2476 """
2477 if n is not None:
2478 return product(range(n), iterable)
2479 seq = tuple(iterable)
2480 if not seq:
2481 return iter(())
2482 return zip(repeat_each(count(), len(seq)), cycle(seq))
2485def mark_ends(iterable):
2486 """Yield 3-tuples of the form ``(is_first, is_last, item)``.
2488 >>> list(mark_ends('ABC'))
2489 [(True, False, 'A'), (False, False, 'B'), (False, True, 'C')]
2491 Use this when looping over an iterable to take special action on its first
2492 and/or last items:
2494 >>> iterable = ['Header', 100, 200, 'Footer']
2495 >>> total = 0
2496 >>> for is_first, is_last, item in mark_ends(iterable):
2497 ... if is_first:
2498 ... continue # Skip the header
2499 ... if is_last:
2500 ... continue # Skip the footer
2501 ... total += item
2502 >>> print(total)
2503 300
2504 """
2505 it = iter(iterable)
2506 for a in it:
2507 first = True
2508 for b in it:
2509 yield first, False, a
2510 a = b
2511 first = False
2512 yield first, True, a
2515def locate(iterable, pred=bool, window_size=None):
2516 """Yield the index of each item in *iterable* for which *pred* returns
2517 ``True``.
2519 *pred* defaults to :func:`bool`, which will select truthy items:
2521 >>> list(locate([0, 1, 1, 0, 1, 0, 0]))
2522 [1, 2, 4]
2524 Set *pred* to a custom function to, e.g., find the indexes for a particular
2525 item.
2527 >>> list(locate(['a', 'b', 'c', 'b'], lambda x: x == 'b'))
2528 [1, 3]
2530 If *window_size* is given, then the *pred* function will be called with
2531 the values in each window. This enables searching for sub-sequences.
2532 Note that *pred* may receive fewer than *window_size* arguments at the end of
2533 the iterable.
2535 >>> iterable = [0, 1, 2, 3, 0, 1, 2, 3, 0, 1, 2, 3]
2536 >>> pred = lambda *args: args == (1, 2, 3)
2537 >>> list(locate(iterable, pred=pred, window_size=3))
2538 [1, 5, 9]
2540 Use with :func:`seekable` to find indexes and then retrieve the associated
2541 items:
2543 >>> from itertools import count
2544 >>> from more_itertools import seekable
2545 >>> source = (3 * n + 1 if (n % 2) else n // 2 for n in count())
2546 >>> it = seekable(source)
2547 >>> pred = lambda x: x > 100
2548 >>> indexes = locate(it, pred=pred)
2549 >>> i = next(indexes)
2550 >>> it.seek(i)
2551 >>> next(it)
2552 106
2554 """
2555 if window_size is None:
2556 return compress(count(), map(pred, iterable))
2558 if window_size < 1:
2559 raise ValueError('window size must be at least 1')
2561 it = windowed(iterable, window_size, fillvalue=_marker)
2562 return compress(
2563 count(),
2564 (pred(*(x for x in w if x is not _marker)) for w in it),
2565 )
2568def longest_common_prefix(iterables):
2569 """Yield elements of the longest common prefix among given *iterables*.
2571 >>> ''.join(longest_common_prefix(['abcd', 'abc', 'abf']))
2572 'ab'
2574 """
2575 return (c[0] for c in takewhile(all_equal, zip(*iterables)))
2578def lstrip(iterable, pred):
2579 """Yield the items from *iterable*, but strip any from the beginning
2580 for which *pred* returns ``True``.
2582 For example, to remove a set of items from the start of an iterable:
2584 >>> iterable = (None, False, None, 1, 2, None, 3, False, None)
2585 >>> pred = lambda x: x in {None, False, ''}
2586 >>> list(lstrip(iterable, pred))
2587 [1, 2, None, 3, False, None]
2589 This function is analogous to :func:`str.lstrip`, and is essentially
2590 a wrapper for :func:`itertools.dropwhile`.
2592 """
2593 return dropwhile(pred, iterable)
2596def rstrip(iterable, pred):
2597 """Yield the items from *iterable*, but strip any from the end
2598 for which *pred* returns ``True``.
2600 For example, to remove a set of items from the end of an iterable:
2602 >>> iterable = (None, False, None, 1, 2, None, 3, False, None)
2603 >>> pred = lambda x: x in {None, False, ''}
2604 >>> list(rstrip(iterable, pred))
2605 [None, False, None, 1, 2, None, 3]
2607 This function is analogous to :func:`str.rstrip`.
2609 """
2610 cache = []
2611 cache_append = cache.append
2612 cache_clear = cache.clear
2613 for x in iterable:
2614 if pred(x):
2615 cache_append(x)
2616 else:
2617 yield from cache
2618 cache_clear()
2619 yield x
2622def strip(iterable, pred):
2623 """Yield the items from *iterable*, but strip any from the
2624 beginning and end for which *pred* returns ``True``.
2626 For example, to remove a set of items from both ends of an iterable:
2628 >>> iterable = (None, False, None, 1, 2, None, 3, False, None)
2629 >>> pred = lambda x: x in {None, False, ''}
2630 >>> list(strip(iterable, pred))
2631 [1, 2, None, 3]
2633 This function is analogous to :func:`str.strip`.
2635 """
2636 return rstrip(lstrip(iterable, pred), pred)
2639class islice_extended:
2640 """An extension of :func:`itertools.islice` that supports negative values
2641 for *stop*, *start*, and *step*.
2643 >>> iterator = iter('abcdefgh')
2644 >>> list(islice_extended(iterator, -4, -1))
2645 ['e', 'f', 'g']
2647 Slices with negative values require some caching of *iterable*, but this
2648 function takes care to minimize the amount of memory required.
2650 For example, you can use a negative step with an infinite iterator:
2652 >>> from itertools import count
2653 >>> list(islice_extended(count(), 110, 99, -2))
2654 [110, 108, 106, 104, 102, 100]
2656 You can also use slice notation directly:
2658 >>> iterator = map(str, count())
2659 >>> it = islice_extended(iterator)[10:20:2]
2660 >>> list(it)
2661 ['10', '12', '14', '16', '18']
2663 """
2665 def __init__(self, iterable, *args):
2666 it = iter(iterable)
2667 if args:
2668 self._iterator = _islice_helper(it, slice(*args))
2669 else:
2670 self._iterator = it
2672 def __iter__(self):
2673 return self
2675 def __next__(self):
2676 return next(self._iterator)
2678 def __getitem__(self, key):
2679 if isinstance(key, slice):
2680 return islice_extended(_islice_helper(self._iterator, key))
2682 raise TypeError('islice_extended.__getitem__ argument must be a slice')
2685def _islice_helper(it, s):
2686 start = s.start
2687 stop = s.stop
2688 if s.step == 0:
2689 raise ValueError('step argument must be a non-zero integer or None.')
2690 step = s.step or 1
2692 if step > 0:
2693 start = 0 if (start is None) else start
2695 if start < 0:
2696 # Consume all but the last -start items
2697 counter = count(1)
2698 wrapper = compress(it, counter)
2699 cache = deque(wrapper, maxlen=-start)
2700 len_iter = next(counter) - 1
2702 # Adjust start to be positive
2703 i = max(len_iter + start, 0)
2705 # Adjust stop to be positive
2706 if stop is None:
2707 j = len_iter
2708 elif stop >= 0:
2709 j = min(stop, len_iter)
2710 else:
2711 j = max(len_iter + stop, 0)
2713 # Slice the cache
2714 n = j - i
2715 if n <= 0:
2716 return
2718 for index in range(n):
2719 if index % step == 0:
2720 # pop and yield the item.
2721 # We don't want to use an intermediate variable
2722 # it would extend the lifetime of the current item
2723 yield cache.popleft()
2724 else:
2725 # just pop and discard the item
2726 cache.popleft()
2727 elif (stop is not None) and (stop < 0):
2728 # Advance to the start position
2729 next(islice(it, start, start), None)
2731 # When stop is negative, we have to carry -stop items while
2732 # iterating
2733 cache = deque(islice(it, -stop), maxlen=-stop)
2735 for index, item in enumerate(it):
2736 if index % step == 0:
2737 # pop and yield the item.
2738 # We don't want to use an intermediate variable
2739 # it would extend the lifetime of the current item
2740 yield cache.popleft()
2741 else:
2742 # just pop and discard the item
2743 cache.popleft()
2744 cache.append(item)
2745 else:
2746 # When both start and stop are positive we have the normal case
2747 yield from islice(it, start, stop, step)
2748 else:
2749 start = -1 if (start is None) else start
2751 if (stop is not None) and (stop < 0):
2752 # Consume all but the last items
2753 n = -stop - 1
2754 counter = count(1)
2755 wrapper = compress(it, counter)
2756 cache = deque(wrapper, maxlen=n)
2757 len_iter = next(counter) - 1
2759 # If start and stop are both negative they are comparable and
2760 # we can just slice. Otherwise we can adjust start to be negative
2761 # and then slice.
2762 if start < 0:
2763 i, j = start, stop
2764 else:
2765 i, j = min(start - len_iter, -1), None
2767 yield from list(cache)[i:j:step]
2768 else:
2769 # Advance to the stop position
2770 if stop is not None:
2771 m = stop + 1
2772 next(islice(it, m, m), None)
2774 # stop is positive, so if start is negative they are not comparable
2775 # and we need the rest of the items.
2776 if start < 0:
2777 i = start
2778 n = None
2779 # stop is None and start is positive, so we just need items up to
2780 # the start index.
2781 elif stop is None:
2782 i = None
2783 n = start + 1
2784 # Both stop and start are positive, so they are comparable.
2785 else:
2786 i = None
2787 n = start - stop
2788 if n <= 0:
2789 return
2791 cache = list(islice(it, n))
2793 yield from cache[i::step]
2796def always_reversible(iterable):
2797 """An extension of :func:`reversed` that supports all iterables, not
2798 just those which implement the ``Reversible`` or ``Sequence`` protocols.
2800 >>> print(*always_reversible(x for x in range(3)))
2801 2 1 0
2803 If the iterable is already reversible, this function returns the
2804 result of :func:`reversed()`. If the iterable is not reversible,
2805 this function will cache the remaining items in the iterable and
2806 yield them in reverse order, which may require significant storage.
2807 """
2808 try:
2809 return reversed(iterable)
2810 except TypeError:
2811 return reversed(list(iterable))
2814def consecutive_groups(iterable, ordering=None):
2815 """Yield groups of consecutive items using :func:`itertools.groupby`.
2816 The *ordering* function determines whether two items are adjacent by
2817 returning their position.
2819 By default, the ordering function is the identity function. This is
2820 suitable for finding runs of numbers:
2822 >>> iterable = [1, 10, 11, 12, 20, 30, 31, 32, 33, 40]
2823 >>> for group in consecutive_groups(iterable):
2824 ... print(list(group))
2825 [1]
2826 [10, 11, 12]
2827 [20]
2828 [30, 31, 32, 33]
2829 [40]
2831 To find runs of adjacent letters, apply :func:`ord` function
2832 to convert letters to ordinals.
2834 >>> iterable = 'abcdfgilmnop'
2835 >>> ordering = ord
2836 >>> for group in consecutive_groups(iterable, ordering):
2837 ... print(list(group))
2838 ['a', 'b', 'c', 'd']
2839 ['f', 'g']
2840 ['i']
2841 ['l', 'm', 'n', 'o', 'p']
2843 Each group of consecutive items is an iterator that shares its source with
2844 *iterable*. When an output group is advanced, the previous group is
2845 no longer available unless its elements are copied (e.g., into a ``list``).
2847 >>> iterable = [1, 2, 11, 12, 21, 22]
2848 >>> saved_groups = []
2849 >>> for group in consecutive_groups(iterable):
2850 ... saved_groups.append(list(group)) # Copy group elements
2851 >>> saved_groups
2852 [[1, 2], [11, 12], [21, 22]]
2854 """
2855 if ordering is None:
2856 key = lambda x: x[0] - x[1]
2857 else:
2858 key = lambda x: x[0] - ordering(x[1])
2860 for k, g in groupby(enumerate(iterable), key=key):
2861 yield map(itemgetter(1), g)
2864def difference(iterable, func=sub, *, initial=None):
2865 """This function is the inverse of :func:`itertools.accumulate`. By default
2866 it will compute the first difference of *iterable* using
2867 :func:`operator.sub`:
2869 >>> from itertools import accumulate
2870 >>> iterable = accumulate([0, 1, 2, 3, 4]) # produces 0, 1, 3, 6, 10
2871 >>> list(difference(iterable))
2872 [0, 1, 2, 3, 4]
2874 *func* defaults to :func:`operator.sub`, but other functions can be
2875 specified. They will be applied as follows::
2877 A, B, C, D, ... --> A, func(B, A), func(C, B), func(D, C), ...
2879 For example, to do progressive division:
2881 >>> iterable = [1, 2, 6, 24, 120]
2882 >>> func = lambda x, y: x // y
2883 >>> list(difference(iterable, func))
2884 [1, 2, 3, 4, 5]
2886 If the *initial* keyword is set, the first element will be skipped when
2887 computing successive differences.
2889 >>> it = [10, 11, 13, 16] # from accumulate([1, 2, 3], initial=10)
2890 >>> list(difference(it, initial=10))
2891 [1, 2, 3]
2893 """
2894 a, b = tee(iterable)
2895 try:
2896 first = [next(b)]
2897 except StopIteration:
2898 return iter([])
2900 if initial is not None:
2901 return map(func, b, a)
2903 return chain(first, map(func, b, a))
2906class SequenceView(Sequence):
2907 """Return a read-only view of the sequence object *target*.
2909 :class:`SequenceView` objects are analogous to Python's built-in
2910 "dictionary view" types. They provide a dynamic view of a sequence's items,
2911 meaning that when the sequence updates, so does the view.
2913 >>> seq = ['0', '1', '2']
2914 >>> view = SequenceView(seq)
2915 >>> view
2916 SequenceView(['0', '1', '2'])
2917 >>> seq.append('3')
2918 >>> view
2919 SequenceView(['0', '1', '2', '3'])
2921 Sequence views support indexing, slicing, and length queries. They act
2922 like the underlying sequence, except they don't allow assignment:
2924 >>> view[1]
2925 '1'
2926 >>> view[1:-1]
2927 ['1', '2']
2928 >>> len(view)
2929 4
2931 Sequence views are useful as an alternative to copying, as they don't
2932 require (much) extra storage.
2934 """
2936 def __init__(self, target):
2937 if not isinstance(target, Sequence):
2938 raise TypeError
2939 self._target = target
2941 def __getitem__(self, index):
2942 return self._target[index]
2944 def __len__(self):
2945 return len(self._target)
2947 def __repr__(self):
2948 return f'{self.__class__.__name__}({self._target!r})'
2951class seekable:
2952 """Wrap an iterator to allow for seeking backward and forward. This
2953 progressively caches the items in the source iterable so they can be
2954 re-visited.
2956 Call :meth:`seek` with an index to seek to that position in the source
2957 iterable.
2959 To "reset" an iterator, seek to ``0``:
2961 >>> from itertools import count
2962 >>> it = seekable((str(n) for n in count()))
2963 >>> next(it), next(it), next(it)
2964 ('0', '1', '2')
2965 >>> it.seek(0)
2966 >>> next(it), next(it), next(it)
2967 ('0', '1', '2')
2969 You can also seek forward:
2971 >>> it = seekable((str(n) for n in range(20)))
2972 >>> it.seek(10)
2973 >>> next(it)
2974 '10'
2975 >>> it.seek(20) # Seeking past the end of the source isn't a problem
2976 >>> list(it)
2977 []
2978 >>> it.seek(0) # Resetting works even after hitting the end
2979 >>> next(it)
2980 '0'
2982 Call :meth:`relative_seek` to seek relative to the source iterator's
2983 current position.
2985 >>> it = seekable((str(n) for n in range(20)))
2986 >>> next(it), next(it), next(it)
2987 ('0', '1', '2')
2988 >>> it.relative_seek(2)
2989 >>> next(it)
2990 '5'
2991 >>> it.relative_seek(-3) # Source is at '6', we move back to '3'
2992 >>> next(it)
2993 '3'
2994 >>> it.relative_seek(-3) # Source is at '4', we move back to '1'
2995 >>> next(it)
2996 '1'
2999 Call :meth:`peek` to look ahead one item without advancing the iterator:
3001 >>> it = seekable('1234')
3002 >>> it.peek()
3003 '1'
3004 >>> list(it)
3005 ['1', '2', '3', '4']
3006 >>> it.peek(default='empty')
3007 'empty'
3009 Before the iterator is at its end, calling :func:`bool` on it will return
3010 ``True``. After it will return ``False``:
3012 >>> it = seekable('5678')
3013 >>> bool(it)
3014 True
3015 >>> list(it)
3016 ['5', '6', '7', '8']
3017 >>> bool(it)
3018 False
3020 You may view the contents of the cache with the :meth:`elements` method.
3021 That returns a :class:`SequenceView`, a view that updates automatically:
3023 >>> it = seekable((str(n) for n in range(10)))
3024 >>> next(it), next(it), next(it)
3025 ('0', '1', '2')
3026 >>> elements = it.elements()
3027 >>> elements
3028 SequenceView(['0', '1', '2'])
3029 >>> next(it)
3030 '3'
3031 >>> elements
3032 SequenceView(['0', '1', '2', '3'])
3034 Indexing the :class:`seekable` directly returns items from the cache:
3036 >>> it = seekable((str(n) for n in range(10)))
3037 >>> next(it), next(it), next(it)
3038 ('0', '1', '2')
3039 >>> it[-1]
3040 '2'
3041 >>> it[0]
3042 '0'
3044 By default, the cache grows as the source iterable progresses, so beware of
3045 wrapping very large or infinite iterables. Supply *maxlen* to limit the
3046 size of the cache (this of course limits how far back you can seek).
3048 >>> from itertools import count
3049 >>> it = seekable((str(n) for n in count()), maxlen=2)
3050 >>> next(it), next(it), next(it), next(it)
3051 ('0', '1', '2', '3')
3052 >>> list(it.elements())
3053 ['2', '3']
3054 >>> it.seek(0)
3055 >>> next(it), next(it), next(it), next(it)
3056 ('2', '3', '4', '5')
3057 >>> next(it)
3058 '6'
3060 """
3062 def __init__(self, iterable, maxlen=None):
3063 self._source = iter(iterable)
3064 if maxlen is None:
3065 self._cache = []
3066 else:
3067 self._cache = deque([], maxlen)
3068 self._index = None
3070 def __iter__(self):
3071 return self
3073 def __next__(self):
3074 if self._index is not None:
3075 try:
3076 item = self._cache[self._index]
3077 except IndexError:
3078 self._index = None
3079 else:
3080 self._index += 1
3081 return item
3083 item = next(self._source)
3084 self._cache.append(item)
3085 return item
3087 def __bool__(self):
3088 try:
3089 self.peek()
3090 except StopIteration:
3091 return False
3092 return True
3094 def peek(self, default=_marker):
3095 try:
3096 peeked = next(self)
3097 except StopIteration:
3098 if default is _marker:
3099 raise
3100 return default
3101 if self._index is None:
3102 self._index = len(self._cache)
3103 self._index -= 1
3104 return peeked
3106 def elements(self):
3107 return SequenceView(self._cache)
3109 def seek(self, index):
3110 self._index = index
3111 remainder = index - len(self._cache)
3112 if remainder > 0:
3113 consume(self, remainder)
3115 def relative_seek(self, count):
3116 if self._index is None:
3117 self._index = len(self._cache)
3119 self.seek(max(self._index + count, 0))
3121 def __getitem__(self, index):
3122 return self._cache[index]
3125class run_length:
3126 """
3127 :func:`run_length.encode` compresses an iterable with run-length encoding.
3128 It yields groups of repeated items with the count of how many times they
3129 were repeated:
3131 >>> uncompressed = 'abbcccdddd'
3132 >>> list(run_length.encode(uncompressed))
3133 [('a', 1), ('b', 2), ('c', 3), ('d', 4)]
3135 :func:`run_length.decode` decompresses an iterable that was previously
3136 compressed with run-length encoding. It yields the items of the
3137 decompressed iterable:
3139 >>> compressed = [('a', 1), ('b', 2), ('c', 3), ('d', 4)]
3140 >>> list(run_length.decode(compressed))
3141 ['a', 'b', 'b', 'c', 'c', 'c', 'd', 'd', 'd', 'd']
3143 """
3145 @staticmethod
3146 def encode(iterable):
3147 return ((k, ilen(g)) for k, g in groupby(iterable))
3149 @staticmethod
3150 def decode(iterable):
3151 return chain.from_iterable(starmap(repeat, iterable))
3154def exactly_n(iterable, n, predicate=bool):
3155 """Return ``True`` if exactly ``n`` items in the iterable are ``True``
3156 according to the *predicate* function.
3158 >>> exactly_n([True, True, False], 2)
3159 True
3160 >>> exactly_n([True, True, False], 1)
3161 False
3162 >>> exactly_n([0, 1, 2, 3, 4, 5], 3, lambda x: x < 3)
3163 True
3165 The iterable will be advanced until ``n + 1`` truthy items are encountered,
3166 so avoid calling it on infinite iterables.
3168 """
3169 iterator = filter(predicate, iterable)
3170 if n <= 0:
3171 if n < 0:
3172 return False
3173 for _ in iterator:
3174 return False
3175 return True
3177 iterator = islice(iterator, n - 1, None)
3178 for _ in iterator:
3179 for _ in iterator:
3180 return False
3181 return True
3182 return False
3185def circular_shifts(iterable, steps=1):
3186 """Yield the circular shifts of *iterable*.
3188 >>> list(circular_shifts(range(4)))
3189 [(0, 1, 2, 3), (1, 2, 3, 0), (2, 3, 0, 1), (3, 0, 1, 2)]
3191 Set *steps* to the number of places to rotate to the left
3192 (or to the right if negative). Defaults to 1.
3194 >>> list(circular_shifts(range(4), 2))
3195 [(0, 1, 2, 3), (2, 3, 0, 1)]
3197 >>> list(circular_shifts(range(4), -1))
3198 [(0, 1, 2, 3), (3, 0, 1, 2), (2, 3, 0, 1), (1, 2, 3, 0)]
3200 """
3201 buffer = deque(iterable)
3202 if steps == 0:
3203 raise ValueError('Steps should be a non-zero integer')
3205 buffer.rotate(steps)
3206 steps = -steps
3207 n = len(buffer)
3208 n //= math.gcd(n, steps)
3210 for _ in repeat(None, n):
3211 buffer.rotate(steps)
3212 yield tuple(buffer)
3215def make_decorator(wrapping_func, result_index=0):
3216 """Return a decorator version of *wrapping_func*, which is a function that
3217 modifies an iterable. *result_index* is the position in that function's
3218 signature where the iterable goes.
3220 This lets you use itertools on the "production end," i.e. at function
3221 definition. This can augment what the function returns without changing the
3222 function's code.
3224 For example, to produce a decorator version of :func:`chunked`:
3226 >>> from more_itertools import chunked
3227 >>> chunker = make_decorator(chunked, result_index=0)
3228 >>> @chunker(3)
3229 ... def iter_range(n):
3230 ... return iter(range(n))
3231 ...
3232 >>> list(iter_range(9))
3233 [[0, 1, 2], [3, 4, 5], [6, 7, 8]]
3235 To only allow truthy items to be returned:
3237 >>> truth_serum = make_decorator(filter, result_index=1)
3238 >>> @truth_serum(bool)
3239 ... def boolean_test():
3240 ... return [0, 1, '', ' ', False, True]
3241 ...
3242 >>> list(boolean_test())
3243 [1, ' ', True]
3245 The :func:`peekable` and :func:`seekable` wrappers make for practical
3246 decorators:
3248 >>> from more_itertools import peekable
3249 >>> peekable_function = make_decorator(peekable)
3250 >>> @peekable_function()
3251 ... def str_range(*args):
3252 ... return (str(x) for x in range(*args))
3253 ...
3254 >>> it = str_range(1, 20, 2)
3255 >>> next(it), next(it), next(it)
3256 ('1', '3', '5')
3257 >>> it.peek()
3258 '7'
3259 >>> next(it)
3260 '7'
3262 """
3264 # See https://sites.google.com/site/bbayles/index/decorator_factory for
3265 # notes on how this works.
3266 def decorator(*wrapping_args, **wrapping_kwargs):
3267 def outer_wrapper(f):
3268 def inner_wrapper(*args, **kwargs):
3269 result = f(*args, **kwargs)
3270 wrapping_args_ = list(wrapping_args)
3271 wrapping_args_.insert(result_index, result)
3272 return wrapping_func(*wrapping_args_, **wrapping_kwargs)
3274 return inner_wrapper
3276 return outer_wrapper
3278 return decorator
3281def map_reduce(iterable, keyfunc, valuefunc=None, reducefunc=None):
3282 """Return a dictionary that maps the items in *iterable* to categories
3283 defined by *keyfunc*, transforms them with *valuefunc*, and
3284 then summarizes them by category with *reducefunc*.
3286 *valuefunc* defaults to the identity function if it is unspecified.
3287 If *reducefunc* is unspecified, no summarization takes place:
3289 >>> keyfunc = lambda x: x.upper()
3290 >>> result = map_reduce('abbccc', keyfunc)
3291 >>> sorted(result.items())
3292 [('A', ['a']), ('B', ['b', 'b']), ('C', ['c', 'c', 'c'])]
3294 Specifying *valuefunc* transforms the categorized items:
3296 >>> keyfunc = lambda x: x.upper()
3297 >>> valuefunc = lambda x: 1
3298 >>> result = map_reduce('abbccc', keyfunc, valuefunc)
3299 >>> sorted(result.items())
3300 [('A', [1]), ('B', [1, 1]), ('C', [1, 1, 1])]
3302 Specifying *reducefunc* summarizes the categorized items:
3304 >>> keyfunc = lambda x: x.upper()
3305 >>> valuefunc = lambda x: 1
3306 >>> reducefunc = sum
3307 >>> result = map_reduce('abbccc', keyfunc, valuefunc, reducefunc)
3308 >>> sorted(result.items())
3309 [('A', 1), ('B', 2), ('C', 3)]
3311 You may want to filter the input iterable before applying the map/reduce
3312 procedure:
3314 >>> all_items = range(30)
3315 >>> items = [x for x in all_items if 10 <= x <= 20] # Filter
3316 >>> keyfunc = lambda x: x % 2 # Evens map to 0; odds to 1
3317 >>> categories = map_reduce(items, keyfunc=keyfunc)
3318 >>> sorted(categories.items())
3319 [(0, [10, 12, 14, 16, 18, 20]), (1, [11, 13, 15, 17, 19])]
3320 >>> summaries = map_reduce(items, keyfunc=keyfunc, reducefunc=sum)
3321 >>> sorted(summaries.items())
3322 [(0, 90), (1, 75)]
3324 Note that all items in the iterable are gathered into a list before the
3325 summarization step, which may require significant storage.
3327 The returned object is a :obj:`collections.defaultdict` with the
3328 ``default_factory`` set to ``None``, such that it behaves like a normal
3329 dictionary.
3331 .. seealso:: :func:`bucket`, :func:`groupby_transform`
3333 If storage is a concern, :func:`bucket` can be used without consuming the
3334 entire iterable right away. If the elements with the same key are already
3335 adjacent, :func:`groupby_transform` or :func:`itertools.groupby` can be
3336 used without any caching overhead.
3338 """
3340 ret = defaultdict(list)
3342 if valuefunc is None:
3343 for item in iterable:
3344 key = keyfunc(item)
3345 ret[key].append(item)
3347 else:
3348 for item in iterable:
3349 key = keyfunc(item)
3350 value = valuefunc(item)
3351 ret[key].append(value)
3353 if reducefunc is not None:
3354 for key, value_list in ret.items():
3355 ret[key] = reducefunc(value_list)
3357 ret.default_factory = None
3358 return ret
3361def rlocate(iterable, pred=bool, window_size=None):
3362 """Yield the index of each item in *iterable* for which *pred* returns
3363 ``True``, starting from the right and moving left.
3365 *pred* defaults to :func:`bool`, which will select truthy items:
3367 >>> list(rlocate([0, 1, 1, 0, 1, 0, 0])) # Truthy at 1, 2, and 4
3368 [4, 2, 1]
3370 Set *pred* to a custom function to, e.g., find the indexes for a particular
3371 item:
3373 >>> iterator = iter('abcb')
3374 >>> pred = lambda x: x == 'b'
3375 >>> list(rlocate(iterator, pred))
3376 [3, 1]
3378 If *window_size* is given, then the *pred* function will be called with
3379 that many items. This enables searching for sub-sequences:
3381 >>> iterable = [0, 1, 2, 3, 0, 1, 2, 3, 0, 1, 2, 3]
3382 >>> pred = lambda *args: args == (1, 2, 3)
3383 >>> list(rlocate(iterable, pred=pred, window_size=3))
3384 [9, 5, 1]
3386 Beware, this function won't return anything for infinite iterables.
3387 If *iterable* is reversible, ``rlocate`` will reverse it and search from
3388 the right. Otherwise, it will search from the left and return the results
3389 in reverse order.
3391 See :func:`locate` to for other example applications.
3393 """
3394 if window_size is None:
3395 try:
3396 len_iter = len(iterable)
3397 return (len_iter - i - 1 for i in locate(reversed(iterable), pred))
3398 except TypeError:
3399 pass
3401 return reversed(list(locate(iterable, pred, window_size)))
3404def replace(iterable, pred, substitutes, count=None, window_size=1):
3405 """Yield the items from *iterable*, replacing the items for which *pred*
3406 returns ``True`` with the items from the iterable *substitutes*.
3408 >>> iterable = [1, 1, 0, 1, 1, 0, 1, 1]
3409 >>> pred = lambda x: x == 0
3410 >>> substitutes = (2, 3)
3411 >>> list(replace(iterable, pred, substitutes))
3412 [1, 1, 2, 3, 1, 1, 2, 3, 1, 1]
3414 If *count* is given, the number of replacements will be limited:
3416 >>> iterable = [1, 1, 0, 1, 1, 0, 1, 1, 0]
3417 >>> pred = lambda x: x == 0
3418 >>> substitutes = [None]
3419 >>> list(replace(iterable, pred, substitutes, count=2))
3420 [1, 1, None, 1, 1, None, 1, 1, 0]
3422 Use *window_size* to control the number of items passed as arguments to
3423 *pred*. This allows for locating and replacing subsequences.
3425 >>> iterable = [0, 1, 2, 5, 0, 1, 2, 5]
3426 >>> window_size = 3
3427 >>> pred = lambda *args: args == (0, 1, 2) # 3 items passed to pred
3428 >>> substitutes = [3, 4] # Splice in these items
3429 >>> list(replace(iterable, pred, substitutes, window_size=window_size))
3430 [3, 4, 5, 3, 4, 5]
3432 *pred* may receive fewer than *window_size* arguments at the end of
3433 the iterable and should be able to handle this.
3435 """
3436 if window_size < 1:
3437 raise ValueError('window_size must be at least 1')
3439 # Save the substitutes iterable, since it's used more than once
3440 substitutes = tuple(substitutes)
3442 # Add padding such that the number of windows matches the length of the
3443 # iterable
3444 it = chain(iterable, repeat(_marker, window_size - 1))
3445 windows = windowed(it, window_size)
3447 n = 0
3448 for w in windows:
3449 # Strip any _marker padding so pred never sees internal sentinels.
3450 # Near the end of the iterable, pred will receive fewer arguments.
3451 args = tuple(x for x in w if x is not _marker)
3453 # If the current window matches our predicate (and we haven't hit
3454 # our maximum number of replacements), splice in the substitutes
3455 # and then consume the following windows that overlap with this one.
3456 # For example, if the iterable is (0, 1, 2, 3, 4...)
3457 # and the window size is 2, we have (0, 1), (1, 2), (2, 3)...
3458 # If the predicate matches on (0, 1), we need to zap (0, 1) and (1, 2)
3459 if args and pred(*args):
3460 if (count is None) or (n < count):
3461 n += 1
3462 yield from substitutes
3463 consume(windows, window_size - 1)
3464 continue
3466 # If there was no match (or we've reached the replacement limit),
3467 # yield the first item from the window.
3468 if args:
3469 yield args[0]
3472def partitions(iterable):
3473 """Yield all possible order-preserving partitions of *iterable*.
3475 >>> iterable = 'abc'
3476 >>> for part in partitions(iterable):
3477 ... print([''.join(p) for p in part])
3478 ['abc']
3479 ['a', 'bc']
3480 ['ab', 'c']
3481 ['a', 'b', 'c']
3483 This is unrelated to :func:`partition`.
3485 """
3486 sequence = list(iterable)
3487 n = len(sequence)
3488 for i in powerset(range(1, n)):
3489 yield [sequence[i:j] for i, j in zip((0,) + i, i + (n,))]
3492def set_partitions(iterable, k=None, min_size=None, max_size=None):
3493 """
3494 Yield the set partitions of *iterable* into *k* parts. Set partitions are
3495 not order-preserving.
3497 >>> iterable = 'abc'
3498 >>> for part in set_partitions(iterable, 2):
3499 ... print([''.join(p) for p in part])
3500 ['a', 'bc']
3501 ['ab', 'c']
3502 ['b', 'ac']
3505 If *k* is not given, every set partition is generated.
3507 >>> iterable = 'abc'
3508 >>> for part in set_partitions(iterable):
3509 ... print([''.join(p) for p in part])
3510 ['abc']
3511 ['a', 'bc']
3512 ['ab', 'c']
3513 ['b', 'ac']
3514 ['a', 'b', 'c']
3516 if *min_size* and/or *max_size* are given, the minimum and/or maximum size
3517 per block in partition is set.
3519 >>> iterable = 'abc'
3520 >>> for part in set_partitions(iterable, min_size=2):
3521 ... print([''.join(p) for p in part])
3522 ['abc']
3523 >>> for part in set_partitions(iterable, max_size=2):
3524 ... print([''.join(p) for p in part])
3525 ['a', 'bc']
3526 ['ab', 'c']
3527 ['b', 'ac']
3528 ['a', 'b', 'c']
3530 """
3531 L = list(iterable)
3532 n = len(L)
3533 if k is not None:
3534 if k < 1:
3535 raise ValueError(
3536 "Can't partition in a negative or zero number of groups"
3537 )
3538 elif k > n:
3539 return
3541 min_size = min_size if min_size is not None else 0
3542 max_size = max_size if max_size is not None else n
3543 if min_size > max_size:
3544 return
3546 def set_partitions_helper(L, k):
3547 n = len(L)
3548 if k == 1:
3549 yield [L]
3550 elif n == k:
3551 yield [[s] for s in L]
3552 else:
3553 e, *M = L
3554 for p in set_partitions_helper(M, k - 1):
3555 yield [[e], *p]
3556 for p in set_partitions_helper(M, k):
3557 for i in range(len(p)):
3558 yield p[:i] + [[e] + p[i]] + p[i + 1 :]
3560 if k is None:
3561 for k in range(1, n + 1):
3562 yield from filter(
3563 lambda z: all(min_size <= len(bk) <= max_size for bk in z),
3564 set_partitions_helper(L, k),
3565 )
3566 else:
3567 yield from filter(
3568 lambda z: all(min_size <= len(bk) <= max_size for bk in z),
3569 set_partitions_helper(L, k),
3570 )
3573class time_limited:
3574 """
3575 Yield items from *iterable* until *limit_seconds* have passed.
3576 If the time limit expires before all items have been yielded, the
3577 ``timed_out`` parameter will be set to ``True``.
3579 >>> from time import sleep
3580 >>> def generator():
3581 ... yield 1
3582 ... yield 2
3583 ... sleep(0.2)
3584 ... yield 3
3585 >>> iterable = time_limited(0.1, generator())
3586 >>> list(iterable)
3587 [1, 2]
3588 >>> iterable.timed_out
3589 True
3591 Note that the time is checked before each item is yielded, and iteration
3592 stops if the time elapsed is greater than *limit_seconds*. If your time
3593 limit is 1 second, but it takes 2 seconds to generate the first item from
3594 the iterable, the function will run for 2 seconds and not yield anything.
3595 As a special case, when *limit_seconds* is zero, the iterator never
3596 returns anything.
3598 """
3600 def __init__(self, limit_seconds, iterable):
3601 if limit_seconds < 0:
3602 raise ValueError('limit_seconds must be positive')
3603 self.limit_seconds = limit_seconds
3604 self._iterator = iter(iterable)
3605 self._start_time = monotonic()
3606 self.timed_out = False
3608 def __iter__(self):
3609 return self
3611 def __next__(self):
3612 if self.limit_seconds == 0:
3613 self.timed_out = True
3614 raise StopIteration
3615 item = next(self._iterator)
3616 if monotonic() - self._start_time > self.limit_seconds:
3617 self.timed_out = True
3618 raise StopIteration
3620 return item
3623def only(iterable, default=None, too_long=None):
3624 """If *iterable* has only one item, return it.
3625 If it has zero items, return *default*.
3626 If it has more than one item, raise the exception given by *too_long*,
3627 which is ``ValueError`` by default.
3629 >>> only([], default='missing')
3630 'missing'
3631 >>> only([1])
3632 1
3633 >>> only([1, 2]) # doctest: +IGNORE_EXCEPTION_DETAIL
3634 Traceback (most recent call last):
3635 ...
3636 ValueError: Expected exactly one item in iterable, but got 1, 2,
3637 and perhaps more.'
3638 >>> only([1, 2], too_long=TypeError) # doctest: +IGNORE_EXCEPTION_DETAIL
3639 Traceback (most recent call last):
3640 ...
3641 TypeError
3643 Note that :func:`only` attempts to advance *iterable* twice to ensure there
3644 is only one item. See :func:`spy` or :func:`peekable` to check
3645 iterable contents less destructively.
3647 """
3648 iterator = iter(iterable)
3649 for first in iterator:
3650 for second in iterator:
3651 msg = (
3652 f'Expected exactly one item in iterable, but got {first!r}, '
3653 f'{second!r}, and perhaps more.'
3654 )
3655 raise too_long or ValueError(msg)
3656 return first
3657 return default
3660def ichunked(iterable, n):
3661 """Break *iterable* into sub-iterables with *n* elements each.
3662 :func:`ichunked` is like :func:`chunked`, but it yields iterables
3663 instead of lists.
3665 If the sub-iterables are read in order, the elements of *iterable*
3666 won't be stored in memory.
3667 If they are read out of order, :func:`itertools.tee` is used to cache
3668 elements as necessary.
3670 >>> from itertools import count
3671 >>> all_chunks = ichunked(count(), 4)
3672 >>> c_1, c_2, c_3 = next(all_chunks), next(all_chunks), next(all_chunks)
3673 >>> list(c_2) # c_1's elements have been cached; c_3's haven't been
3674 [4, 5, 6, 7]
3675 >>> list(c_1)
3676 [0, 1, 2, 3]
3677 >>> list(c_3)
3678 [8, 9, 10, 11]
3680 """
3681 iterator = iter(iterable)
3682 for first in iterator:
3683 rest = islice(iterator, n - 1)
3684 cache, cacher = tee(rest)
3685 yield chain([first], rest, cache)
3686 consume(cacher)
3689def iequals(*iterables):
3690 """Return ``True`` if all given *iterables* are equal to each other,
3691 which means that they contain the same elements in the same order.
3693 The function is useful for comparing iterables of different data types
3694 or iterables that do not support equality checks.
3696 >>> iequals("abc", ['a', 'b', 'c'], ('a', 'b', 'c'), iter("abc"))
3697 True
3699 >>> iequals("abc", "acb")
3700 False
3702 Not to be confused with :func:`all_equal`, which checks whether all
3703 elements of iterable are equal to each other.
3705 """
3706 try:
3707 return all(map(all_equal, zip(*iterables, strict=True)))
3708 except ValueError:
3709 return False
3712def distinct_combinations(iterable, r):
3713 """Yield the distinct combinations of *r* items taken from *iterable*.
3715 >>> list(distinct_combinations([0, 0, 1], 2))
3716 [(0, 0), (0, 1)]
3718 Equivalent to ``set(combinations(iterable))``, except duplicates are not
3719 generated and thrown away. For larger input sequences this is much more
3720 efficient.
3722 """
3723 if r < 0:
3724 raise ValueError('r must be non-negative')
3725 elif r == 0:
3726 yield ()
3727 return
3728 pool = tuple(iterable)
3729 generators = [unique_everseen(enumerate(pool), key=itemgetter(1))]
3730 current_combo = [None] * r
3731 level = 0
3732 while generators:
3733 try:
3734 cur_idx, p = next(generators[-1])
3735 except StopIteration:
3736 generators.pop()
3737 level -= 1
3738 continue
3739 current_combo[level] = p
3740 if level + 1 == r:
3741 yield tuple(current_combo)
3742 else:
3743 generators.append(
3744 unique_everseen(
3745 enumerate(pool[cur_idx + 1 :], cur_idx + 1),
3746 key=itemgetter(1),
3747 )
3748 )
3749 level += 1
3752def filter_except(validator, iterable, *exceptions):
3753 """Yield the items from *iterable* for which the *validator* function does
3754 not raise one of the specified *exceptions*.
3756 *validator* is called for each item in *iterable*.
3757 It should be a function that accepts one argument and raises an exception
3758 if that item is not valid.
3760 >>> iterable = ['1', '2', 'three', '4', None]
3761 >>> list(filter_except(int, iterable, ValueError, TypeError))
3762 ['1', '2', '4']
3764 If an exception other than one given by *exceptions* is raised by
3765 *validator*, it is raised like normal.
3766 """
3767 for item in iterable:
3768 try:
3769 validator(item)
3770 except exceptions:
3771 pass
3772 else:
3773 yield item
3776def map_except(function, iterable, *exceptions):
3777 """Transform each item from *iterable* with *function* and yield the
3778 result, unless *function* raises one of the specified *exceptions*.
3780 *function* is called to transform each item in *iterable*.
3781 It should accept one argument.
3783 >>> iterable = ['1', '2', 'three', '4', None]
3784 >>> list(map_except(int, iterable, ValueError, TypeError))
3785 [1, 2, 4]
3787 If an exception other than one given by *exceptions* is raised by
3788 *function*, it is raised like normal.
3789 """
3790 for item in iterable:
3791 try:
3792 yield function(item)
3793 except exceptions:
3794 pass
3797def map_if(iterable, pred, func, func_else=None):
3798 """Evaluate each item from *iterable* using *pred*. If the result is
3799 equivalent to ``True``, transform the item with *func* and yield it.
3800 Otherwise, transform the item with *func_else* and yield it.
3802 *pred*, *func*, and *func_else* should each be functions that accept
3803 one argument. By default, *func_else* is the identity function.
3805 >>> from math import sqrt
3806 >>> iterable = list(range(-5, 5))
3807 >>> iterable
3808 [-5, -4, -3, -2, -1, 0, 1, 2, 3, 4]
3809 >>> list(map_if(iterable, lambda x: x > 3, lambda x: 'toobig'))
3810 [-5, -4, -3, -2, -1, 0, 1, 2, 3, 'toobig']
3811 >>> list(map_if(iterable, lambda x: x >= 0,
3812 ... lambda x: f'{sqrt(x):.2f}', lambda x: None))
3813 [None, None, None, None, None, '0.00', '1.00', '1.41', '1.73', '2.00']
3814 """
3816 if func_else is None:
3817 for item in iterable:
3818 yield func(item) if pred(item) else item
3820 else:
3821 for item in iterable:
3822 yield func(item) if pred(item) else func_else(item)
3825def _sample_unweighted(iterator, k, strict):
3826 # Algorithm L in the 1994 paper by Kim-Hung Li:
3827 # "Reservoir-Sampling Algorithms of Time Complexity O(n(1+log(N/n)))".
3829 reservoir = list(islice(iterator, k))
3830 if strict and len(reservoir) < k:
3831 raise ValueError('Sample larger than population')
3832 W = 1.0
3834 with suppress(StopIteration):
3835 while True:
3836 W *= random() ** (1 / k)
3837 skip = floor(log(random()) / log1p(-W))
3838 element = next(islice(iterator, skip, None))
3839 reservoir[randrange(k)] = element
3841 shuffle(reservoir)
3842 return reservoir
3845def _sample_weighted(iterator, k, weights, strict):
3846 # Implementation of "A-ExpJ" from the 2006 paper by Efraimidis et al. :
3847 # "Weighted random sampling with a reservoir".
3849 # Log-transform for numerical stability for weights that are small/large
3850 weight_keys = (log(random()) / weight for weight in weights)
3852 # Fill up the reservoir (collection of samples) with the first `k`
3853 # weight-keys and elements, then heapify the list.
3854 reservoir = take(k, zip(weight_keys, iterator))
3855 if strict and len(reservoir) < k:
3856 raise ValueError('Sample larger than population')
3858 heapify(reservoir)
3860 # The number of jumps before changing the reservoir is a random variable
3861 # with an exponential distribution. Sample it using random() and logs.
3862 smallest_weight_key, _ = reservoir[0]
3863 weights_to_skip = log(random()) / smallest_weight_key
3865 for weight, element in zip(weights, iterator):
3866 if weight >= weights_to_skip:
3867 # The notation here is consistent with the paper, but we store
3868 # the weight-keys in log-space for better numerical stability.
3869 smallest_weight_key, _ = reservoir[0]
3870 t_w = exp(weight * smallest_weight_key)
3871 r_2 = uniform(t_w, 1) # generate U(t_w, 1)
3872 weight_key = log(r_2) / weight
3873 heapreplace(reservoir, (weight_key, element))
3874 smallest_weight_key, _ = reservoir[0]
3875 weights_to_skip = log(random()) / smallest_weight_key
3876 else:
3877 weights_to_skip -= weight
3879 ret = [element for weight_key, element in reservoir]
3880 shuffle(ret)
3881 return ret
3884def _sample_counted(population, k, counts, strict):
3885 element = None
3886 remaining = 0
3888 def feed(i):
3889 # Advance *i* steps ahead and consume an element
3890 nonlocal element, remaining
3892 while i + 1 > remaining:
3893 i = i - remaining
3894 element = next(population)
3895 remaining = next(counts)
3896 remaining -= i + 1
3897 return element
3899 with suppress(StopIteration):
3900 reservoir = []
3901 for _ in range(k):
3902 reservoir.append(feed(0))
3904 if strict and len(reservoir) < k:
3905 raise ValueError('Sample larger than population')
3907 with suppress(StopIteration):
3908 W = 1.0
3909 while True:
3910 W *= random() ** (1 / k)
3911 skip = floor(log(random()) / log1p(-W))
3912 element = feed(skip)
3913 reservoir[randrange(k)] = element
3915 shuffle(reservoir)
3916 return reservoir
3919def sample(iterable, k, weights=None, *, counts=None, strict=False):
3920 """Return a *k*-length list of elements chosen (without replacement)
3921 from the *iterable*.
3923 Similar to :func:`random.sample`, but works on inputs that aren't
3924 indexable (such as sets and dictionaries) and on inputs where the
3925 size isn't known in advance (such as generators).
3927 >>> iterable = range(100)
3928 >>> sample(iterable, 5) # doctest: +SKIP
3929 [81, 60, 96, 16, 4]
3931 For iterables with repeated elements, you may supply *counts* to
3932 indicate the repeats.
3934 >>> iterable = ['a', 'b']
3935 >>> counts = [3, 4] # Equivalent to 'a', 'a', 'a', 'b', 'b', 'b', 'b'
3936 >>> sample(iterable, k=3, counts=counts) # doctest: +SKIP
3937 ['a', 'a', 'b']
3939 An iterable with *weights* may be given:
3941 >>> iterable = range(100)
3942 >>> weights = (i * i + 1 for i in range(100))
3943 >>> sampled = sample(iterable, 5, weights=weights) # doctest: +SKIP
3944 [79, 67, 74, 66, 78]
3946 Weighted selections are made without replacement.
3947 After an element is selected, it is removed from the pool and the
3948 relative weights of the other elements increase (this
3949 does not match the behavior of :func:`random.sample`'s *counts*
3950 parameter). Note that *weights* may not be used with *counts*.
3952 If the length of *iterable* is less than *k*,
3953 ``ValueError`` is raised if *strict* is ``True`` and
3954 all elements are returned (in shuffled order) if *strict* is ``False``.
3956 By default, the `Algorithm L <https://w.wiki/ANrM>`__ reservoir sampling
3957 technique is used. When *weights* are provided,
3958 `Algorithm A-ExpJ <https://w.wiki/ANrS>`__ is used instead.
3960 Notes on reproducibility:
3962 * The algorithms rely on inexact floating-point functions provided
3963 by the underlying math library (e.g. ``log``, ``log1p``, and ``pow``).
3964 Those functions can `produce slightly different results
3965 <https://members.loria.fr/PZimmermann/papers/accuracy.pdf>`_ on
3966 different builds. Accordingly, selections can vary across builds
3967 even for the same seed.
3969 * The algorithms loop over the input and make selections based on
3970 ordinal position, so selections from unordered collections (such as
3971 sets) won't reproduce across sessions on the same platform using the
3972 same seed. For example, this won't reproduce::
3974 >> seed(8675309)
3975 >> sample(set('abcdefghijklmnopqrstuvwxyz'), 10)
3976 ['c', 'p', 'e', 'w', 's', 'a', 'j', 'd', 'n', 't']
3978 """
3979 iterator = iter(iterable)
3981 if k < 0:
3982 raise ValueError('k must be non-negative')
3984 if k == 0:
3985 return []
3987 if weights is not None and counts is not None:
3988 raise TypeError('weights and counts are mutually exclusive')
3990 elif weights is not None:
3991 weights = iter(weights)
3992 return _sample_weighted(iterator, k, weights, strict)
3994 elif counts is not None:
3995 counts = iter(counts)
3996 return _sample_counted(iterator, k, counts, strict)
3998 else:
3999 return _sample_unweighted(iterator, k, strict)
4002def is_sorted(iterable, key=None, reverse=False, strict=False):
4003 """Returns ``True`` if the items of iterable are in sorted order, and
4004 ``False`` otherwise. *key* and *reverse* have the same meaning that they do
4005 in the built-in :func:`sorted` function.
4007 >>> is_sorted(['1', '2', '3', '4', '5'], key=int)
4008 True
4009 >>> is_sorted([5, 4, 3, 1, 2], reverse=True)
4010 False
4012 If *strict*, tests for strict sorting, that is, returns ``False`` if equal
4013 elements are found:
4015 >>> is_sorted([1, 2, 2])
4016 True
4017 >>> is_sorted([1, 2, 2], strict=True)
4018 False
4020 The function returns ``False`` after encountering the first out-of-order
4021 item, which means it may produce results that differ from the built-in
4022 :func:`sorted` function for objects with unusual comparison dynamics
4023 (like ``math.nan``). If there are no out-of-order items, the iterable is
4024 exhausted.
4025 """
4026 it = iterable if (key is None) else map(key, iterable)
4027 a, b = tee(it)
4028 next(b, None)
4029 if reverse:
4030 b, a = a, b
4031 return all(map(lt, a, b)) if strict else not any(map(lt, b, a))
4034class AbortThread(BaseException):
4035 pass
4038class callback_iter:
4039 """Convert a function that uses callbacks to an iterator.
4041 .. deprecated:: 11.0.0
4042 Will be removed in a future major release.
4044 Let *func* be a function that takes a `callback` keyword argument.
4045 For example:
4047 >>> def func(callback=None):
4048 ... for i, c in [(1, 'a'), (2, 'b'), (3, 'c')]:
4049 ... if callback:
4050 ... callback(i, c)
4051 ... return 4
4054 Use ``with callback_iter(func)`` to get an iterator over the parameters
4055 that are delivered to the callback.
4057 >>> with callback_iter(func) as it:
4058 ... for args, kwargs in it:
4059 ... print(args)
4060 (1, 'a')
4061 (2, 'b')
4062 (3, 'c')
4064 The function will be called in a background thread. The ``done`` property
4065 indicates whether it has completed execution.
4067 >>> it.done
4068 True
4070 If it completes successfully, its return value will be available
4071 in the ``result`` property.
4073 >>> it.result
4074 4
4076 Notes:
4078 * If the function uses some keyword argument besides ``callback``, supply
4079 *callback_kwd*.
4080 * If it finished executing, but raised an exception, accessing the
4081 ``result`` property will raise the same exception.
4082 * If it hasn't finished executing, accessing the ``result``
4083 property from within the ``with`` block will raise ``RuntimeError``.
4084 * If it hasn't finished executing, accessing the ``result`` property from
4085 outside the ``with`` block will raise a
4086 ``more_itertools.AbortThread`` exception.
4087 * Provide *wait_seconds* to adjust how frequently the it is polled for
4088 output.
4090 """
4092 def __init__(self, func, callback_kwd='callback', wait_seconds=0.1):
4093 self._func = func
4094 self._callback_kwd = callback_kwd
4095 self._aborted = False
4096 self._future = None
4097 self._wait_seconds = wait_seconds
4099 # Lazily import concurrent.future
4100 self._module = __import__('concurrent.futures').futures
4101 self._executor = self._module.ThreadPoolExecutor(max_workers=1)
4102 self._iterator = self._reader()
4104 def __enter__(self):
4105 return self
4107 def __exit__(self, exc_type, exc_value, traceback):
4108 self._aborted = True
4109 self._executor.shutdown()
4111 def __iter__(self):
4112 return self
4114 def __next__(self):
4115 return next(self._iterator)
4117 @property
4118 def done(self):
4119 if self._future is None:
4120 return False
4121 return self._future.done()
4123 @property
4124 def result(self):
4125 if self._future:
4126 try:
4127 return self._future.result(timeout=0)
4128 except self._module.TimeoutError:
4129 pass
4131 raise RuntimeError('Function has not yet completed')
4133 def _reader(self):
4134 q = Queue()
4136 def callback(*args, **kwargs):
4137 if self._aborted:
4138 raise AbortThread('canceled by user')
4140 q.put((args, kwargs))
4142 self._future = self._executor.submit(
4143 self._func, **{self._callback_kwd: callback}
4144 )
4146 while True:
4147 try:
4148 item = q.get(timeout=self._wait_seconds)
4149 except Empty:
4150 pass
4151 else:
4152 q.task_done()
4153 yield item
4155 if self._future.done():
4156 break
4158 remaining = []
4159 while True:
4160 try:
4161 item = q.get_nowait()
4162 except Empty:
4163 break
4164 else:
4165 q.task_done()
4166 remaining.append(item)
4167 q.join()
4168 yield from remaining
4171def windowed_complete(iterable, n):
4172 """
4173 Yield ``(beginning, middle, end)`` tuples, where:
4175 * Each ``middle`` has *n* items from *iterable*
4176 * Each ``beginning`` has the items before the ones in ``middle``
4177 * Each ``end`` has the items after the ones in ``middle``
4179 >>> iterable = range(7)
4180 >>> n = 3
4181 >>> for beginning, middle, end in windowed_complete(iterable, n):
4182 ... print(beginning, middle, end)
4183 () (0, 1, 2) (3, 4, 5, 6)
4184 (0,) (1, 2, 3) (4, 5, 6)
4185 (0, 1) (2, 3, 4) (5, 6)
4186 (0, 1, 2) (3, 4, 5) (6,)
4187 (0, 1, 2, 3) (4, 5, 6) ()
4189 Note that *n* must be at least 0 and most equal to the length of
4190 *iterable*.
4192 This function will exhaust the iterable and may require significant
4193 storage.
4194 """
4195 if n < 0:
4196 raise ValueError('n must be >= 0')
4198 seq = tuple(iterable)
4199 size = len(seq)
4201 if n > size:
4202 raise ValueError('n must be <= len(seq)')
4204 for i in range(size - n + 1):
4205 beginning = seq[:i]
4206 middle = seq[i : i + n]
4207 end = seq[i + n :]
4208 yield beginning, middle, end
4211def all_unique(iterable, key=None):
4212 """
4213 Returns ``True`` if all the elements of *iterable* are unique (no two
4214 elements are equal).
4216 >>> all_unique('ABCB')
4217 False
4219 If a *key* function is specified, it will be used to make comparisons.
4221 >>> all_unique('ABCb')
4222 True
4223 >>> all_unique('ABCb', str.lower)
4224 False
4226 The function returns as soon as the first non-unique element is
4227 encountered. Iterables with a mix of hashable and unhashable items can
4228 be used, but the function will be slower for unhashable items.
4229 """
4230 seenset = set()
4231 seenset_add = seenset.add
4232 seenlist = []
4233 seenlist_add = seenlist.append
4234 for element in map(key, iterable) if key else iterable:
4235 try:
4236 if element in seenset:
4237 return False
4238 seenset_add(element)
4239 except TypeError:
4240 if element in seenlist:
4241 return False
4242 seenlist_add(element)
4243 return True
4246def nth_product(index, *iterables, repeat=1):
4247 """Equivalent to ``list(product(*iterables, repeat=repeat))[index]``.
4249 The products of *iterables* can be ordered lexicographically.
4250 :func:`nth_product` computes the product at sort position *index* without
4251 computing the previous products.
4253 >>> nth_product(8, range(2), range(2), range(2), range(2))
4254 (1, 0, 0, 0)
4256 The *repeat* keyword argument specifies the number of repetitions
4257 of the iterables. The above example is equivalent to::
4259 >>> nth_product(8, range(2), repeat=4)
4260 (1, 0, 0, 0)
4262 ``IndexError`` will be raised if the given *index* is invalid.
4263 """
4264 pools = tuple(map(tuple, reversed(iterables))) * repeat
4265 ns = tuple(map(len, pools))
4267 c = prod(ns)
4269 if index < 0:
4270 index += c
4271 if not 0 <= index < c:
4272 raise IndexError
4274 result = []
4275 for pool, n in zip(pools, ns):
4276 result.append(pool[index % n])
4277 index //= n
4279 return tuple(reversed(result))
4282def nth_permutation(iterable, r, index):
4283 """Equivalent to ``list(permutations(iterable, r))[index]```
4285 The subsequences of *iterable* that are of length *r* where order is
4286 important can be ordered lexicographically. :func:`nth_permutation`
4287 computes the subsequence at sort position *index* directly, without
4288 computing the previous subsequences.
4290 >>> nth_permutation('ghijk', 2, 5)
4291 ('h', 'i')
4293 ``ValueError`` will be raised If *r* is negative.
4294 ``IndexError`` will be raised if the given *index* is invalid.
4295 """
4296 pool = list(iterable)
4297 n = len(pool)
4298 if r is None:
4299 r = n
4300 c = perm(n, r)
4302 if index < 0:
4303 index += c
4304 if not 0 <= index < c:
4305 raise IndexError
4307 result = [0] * r
4308 q = index
4309 for d in range(n - r + 1, n + 1):
4310 q, i = divmod(q, d)
4311 result[n - d] = i
4312 if q == 0:
4313 break
4315 return tuple(map(pool.pop, result))
4318def nth_combination_with_replacement(iterable, r, index):
4319 """Equivalent to
4320 ``list(combinations_with_replacement(iterable, r))[index]``.
4323 The subsequences with repetition of *iterable* that are of length *r* can
4324 be ordered lexicographically. :func:`nth_combination_with_replacement`
4325 computes the subsequence at sort position *index* directly, without
4326 computing the previous subsequences with replacement.
4328 >>> nth_combination_with_replacement(range(5), 3, 5)
4329 (0, 1, 1)
4331 ``ValueError`` will be raised If *r* is negative.
4332 ``IndexError`` will be raised if the given *index* is invalid.
4333 """
4334 pool = tuple(iterable)
4335 n = len(pool)
4336 if r < 0:
4337 raise ValueError
4338 c = comb(n + r - 1, r) if n else 0 if r else 1
4340 if index < 0:
4341 index += c
4342 if not 0 <= index < c:
4343 raise IndexError
4345 result = []
4346 i = 0
4347 while r:
4348 r -= 1
4349 while n >= 0:
4350 num_combs = comb(n + r - 1, r)
4351 if index < num_combs:
4352 break
4353 n -= 1
4354 i += 1
4355 index -= num_combs
4356 result.append(pool[i])
4358 return tuple(result)
4361def value_chain(*args):
4362 """Yield all arguments passed to the function in the same order in which
4363 they were passed. If an argument itself is iterable then iterate over its
4364 values.
4366 >>> list(value_chain(1, 2, 3, [4, 5, 6]))
4367 [1, 2, 3, 4, 5, 6]
4369 Binary and text strings are not considered iterable and are emitted
4370 as-is:
4372 >>> list(value_chain('12', '34', ['56', '78']))
4373 ['12', '34', '56', '78']
4375 Pre- or postpend a single element to an iterable:
4377 >>> list(value_chain(1, [2, 3, 4, 5, 6]))
4378 [1, 2, 3, 4, 5, 6]
4379 >>> list(value_chain([1, 2, 3, 4, 5], 6))
4380 [1, 2, 3, 4, 5, 6]
4382 Multiple levels of nesting are not flattened.
4384 """
4385 scalar_types = (str, bytes)
4386 for value in args:
4387 if isinstance(value, scalar_types):
4388 yield value
4389 continue
4390 try:
4391 yield from value
4392 except TypeError:
4393 yield value
4396def product_index(element, *iterables, repeat=1):
4397 """Equivalent to ``list(product(*iterables, repeat=repeat)).index(tuple(element))``
4399 The products of *iterables* can be ordered lexicographically.
4400 :func:`product_index` computes the first index of *element* without
4401 computing the previous products.
4403 >>> product_index([8, 2], range(10), range(5))
4404 42
4406 The *repeat* keyword argument specifies the number of repetitions
4407 of the iterables::
4409 >>> product_index([8, 0, 7], range(10), repeat=3)
4410 807
4412 ``ValueError`` will be raised if the given *element* isn't in the product
4413 of *args*.
4414 """
4415 elements = tuple(element)
4416 pools = tuple(map(tuple, iterables)) * repeat
4417 if len(elements) != len(pools):
4418 raise ValueError('element is not a product of args')
4420 index = 0
4421 for elem, pool in zip(elements, pools):
4422 index = index * len(pool) + pool.index(elem)
4423 return index
4426def combination_index(element, iterable):
4427 """Equivalent to ``list(combinations(iterable, r)).index(element)``
4429 The subsequences of *iterable* that are of length *r* can be ordered
4430 lexicographically. :func:`combination_index` computes the index of the
4431 first *element*, without computing the previous combinations.
4433 >>> combination_index('adf', 'abcdefg')
4434 10
4436 ``ValueError`` will be raised if the given *element* isn't one of the
4437 combinations of *iterable*.
4438 """
4439 element = enumerate(element)
4440 k, y = next(element, (None, None))
4441 if k is None:
4442 return 0
4444 indexes = []
4445 pool = enumerate(iterable)
4446 for n, x in pool:
4447 if x == y:
4448 indexes.append(n)
4449 tmp, y = next(element, (None, None))
4450 if tmp is None:
4451 break
4452 else:
4453 k = tmp
4454 else:
4455 raise ValueError('element is not a combination of iterable')
4457 n, _ = last(pool, default=(n, None))
4459 index = 1
4460 for i, j in enumerate(reversed(indexes), start=1):
4461 j = n - j
4462 if i <= j:
4463 index += comb(j, i)
4465 return comb(n + 1, k + 1) - index
4468def combination_with_replacement_index(element, iterable):
4469 """Equivalent to
4470 ``list(combinations_with_replacement(iterable, r)).index(element)``
4472 The subsequences with repetition of *iterable* that are of length *r* can
4473 be ordered lexicographically. :func:`combination_with_replacement_index`
4474 computes the index of the first *element*, without computing the previous
4475 combinations with replacement.
4477 >>> combination_with_replacement_index('adf', 'abcdefg')
4478 20
4480 ``ValueError`` will be raised if the given *element* isn't one of the
4481 combinations with replacement of *iterable*.
4482 """
4483 element = tuple(element)
4484 l = len(element)
4485 element = enumerate(element)
4487 k, y = next(element, (None, None))
4488 if k is None:
4489 return 0
4491 indexes = []
4492 pool = tuple(iterable)
4493 for n, x in enumerate(pool):
4494 while x == y:
4495 indexes.append(n)
4496 tmp, y = next(element, (None, None))
4497 if tmp is None:
4498 break
4499 else:
4500 k = tmp
4501 if y is None:
4502 break
4503 else:
4504 raise ValueError(
4505 'element is not a combination with replacement of iterable'
4506 )
4508 n = len(pool)
4509 occupations = [0] * n
4510 for p in indexes:
4511 occupations[p] += 1
4513 index = 0
4514 cumulative_sum = 0
4515 for k in range(1, n):
4516 cumulative_sum += occupations[k - 1]
4517 j = l + n - 1 - k - cumulative_sum
4518 i = n - k
4519 if i <= j:
4520 index += comb(j, i)
4522 return index
4525def permutation_index(element, iterable):
4526 """Equivalent to ``list(permutations(iterable, r)).index(element)```
4528 The subsequences of *iterable* that are of length *r* where order is
4529 important can be ordered lexicographically. :func:`permutation_index`
4530 computes the index of the first *element* directly, without computing
4531 the previous permutations.
4533 >>> permutation_index([1, 3, 2], range(5))
4534 19
4536 ``ValueError`` will be raised if the given *element* isn't one of the
4537 permutations of *iterable*.
4538 """
4539 index = 0
4540 pool = list(iterable)
4541 for i, x in zip(range(len(pool), -1, -1), element):
4542 r = pool.index(x)
4543 index = index * i + r
4544 del pool[r]
4546 return index
4549class countable:
4550 """Wrap *iterable* and keep a count of how many items have been consumed.
4552 The ``items_seen`` attribute starts at ``0`` and increments as the iterable
4553 is consumed:
4555 >>> iterable = map(str, range(10))
4556 >>> it = countable(iterable)
4557 >>> it.items_seen
4558 0
4559 >>> next(it), next(it)
4560 ('0', '1')
4561 >>> list(it)
4562 ['2', '3', '4', '5', '6', '7', '8', '9']
4563 >>> it.items_seen
4564 10
4565 """
4567 def __init__(self, iterable):
4568 self._iterator = iter(iterable)
4569 self.items_seen = 0
4571 def __iter__(self):
4572 return self
4574 def __next__(self):
4575 item = next(self._iterator)
4576 self.items_seen += 1
4578 return item
4581def chunked_even(iterable, n):
4582 """Break *iterable* into lists of approximately length *n*.
4583 Items are distributed such the lengths of the lists differ by at most
4584 1 item.
4586 >>> iterable = [1, 2, 3, 4, 5, 6, 7]
4587 >>> n = 3
4588 >>> list(chunked_even(iterable, n)) # List lengths: 3, 2, 2
4589 [[1, 2, 3], [4, 5], [6, 7]]
4590 >>> list(chunked(iterable, n)) # List lengths: 3, 3, 1
4591 [[1, 2, 3], [4, 5, 6], [7]]
4593 """
4594 iterator = iter(iterable)
4596 # Initialize a buffer to process the chunks while keeping
4597 # some back to fill any underfilled chunks
4598 min_buffer = (n - 1) * (n - 2)
4599 buffer = list(islice(iterator, min_buffer))
4601 # Append items until we have a completed chunk
4602 for _ in islice(map(buffer.append, iterator), n, None, n):
4603 yield buffer[:n]
4604 del buffer[:n]
4606 # Check if any chunks need addition processing
4607 if not buffer:
4608 return
4609 length = len(buffer)
4611 # Chunks are either size `full_size <= n` or `partial_size = full_size - 1`
4612 q, r = divmod(length, n)
4613 num_lists = q + (1 if r > 0 else 0)
4614 q, r = divmod(length, num_lists)
4615 full_size = q + (1 if r > 0 else 0)
4616 partial_size = full_size - 1
4617 num_full = length - partial_size * num_lists
4619 # Yield chunks of full size
4620 partial_start_idx = num_full * full_size
4621 if full_size > 0:
4622 for i in range(0, partial_start_idx, full_size):
4623 yield buffer[i : i + full_size]
4625 # Yield chunks of partial size
4626 if partial_size > 0:
4627 for i in range(partial_start_idx, length, partial_size):
4628 yield buffer[i : i + partial_size]
4631def zip_broadcast(*objects, scalar_types=(str, bytes), strict=False):
4632 """A version of :func:`zip` that "broadcasts" any scalar
4633 (i.e., non-iterable) items into output tuples.
4635 >>> iterable_1 = [1, 2, 3]
4636 >>> iterable_2 = ['a', 'b', 'c']
4637 >>> scalar = '_'
4638 >>> list(zip_broadcast(iterable_1, iterable_2, scalar))
4639 [(1, 'a', '_'), (2, 'b', '_'), (3, 'c', '_')]
4641 The *scalar_types* keyword argument determines what types are considered
4642 scalar. It is set to ``(str, bytes)`` by default. Set it to ``None`` to
4643 treat strings and byte strings as iterable:
4645 >>> list(zip_broadcast('abc', 0, 'xyz', scalar_types=None))
4646 [('a', 0, 'x'), ('b', 0, 'y'), ('c', 0, 'z')]
4648 If the *strict* keyword argument is ``True``, then
4649 ``ValueError`` will be raised if any of the iterables have
4650 different lengths.
4651 """
4653 def is_scalar(obj):
4654 if scalar_types and isinstance(obj, scalar_types):
4655 return True
4656 try:
4657 iter(obj)
4658 except TypeError:
4659 return True
4660 else:
4661 return False
4663 size = len(objects)
4664 if not size:
4665 return
4667 new_item = [None] * size
4668 iterables, iterable_positions = [], []
4669 for i, obj in enumerate(objects):
4670 if is_scalar(obj):
4671 new_item[i] = obj
4672 else:
4673 iterables.append(iter(obj))
4674 iterable_positions.append(i)
4676 if not iterables:
4677 yield tuple(objects)
4678 return
4680 for item in zip(*iterables, strict=strict):
4681 for i, new_item[i] in zip(iterable_positions, item):
4682 pass
4683 yield tuple(new_item)
4686def unique_in_window(iterable, n, key=None):
4687 """Yield the items from *iterable* that haven't been seen recently.
4688 *n* is the size of the sliding window.
4690 >>> iterable = [0, 1, 0, 2, 3, 0]
4691 >>> n = 3
4692 >>> list(unique_in_window(iterable, n))
4693 [0, 1, 2, 3, 0]
4695 The *key* function, if provided, will be used to determine uniqueness:
4697 >>> list(unique_in_window('abAcda', 3, key=lambda x: x.lower()))
4698 ['a', 'b', 'c', 'd', 'a']
4700 Updates a sliding window no larger than n and yields a value
4701 if the item only occurs once in the updated window.
4703 When `n == 1`, *unique_in_window* is memoryless:
4705 >>> list(unique_in_window('aab', n=1))
4706 ['a', 'a', 'b']
4708 The items in *iterable* must be hashable.
4710 """
4711 if n <= 0:
4712 raise ValueError('n must be greater than 0')
4714 window = deque(maxlen=n)
4715 counts = Counter()
4716 use_key = key is not None
4718 for item in iterable:
4719 if len(window) == n:
4720 to_discard = window[0]
4721 if counts[to_discard] == 1:
4722 del counts[to_discard]
4723 else:
4724 counts[to_discard] -= 1
4726 k = key(item) if use_key else item
4727 if k not in counts:
4728 yield item
4729 counts[k] += 1
4730 window.append(k)
4733def duplicates_everseen(iterable, key=None):
4734 """Yield duplicate elements after their first appearance.
4736 >>> list(duplicates_everseen('mississippi'))
4737 ['s', 'i', 's', 's', 'i', 'p', 'i']
4738 >>> list(duplicates_everseen('AaaBbbCccAaa', str.lower))
4739 ['a', 'a', 'b', 'b', 'c', 'c', 'A', 'a', 'a']
4741 This function is analogous to :func:`unique_everseen` and is subject to
4742 the same performance considerations.
4744 If you would like each duplicate to only appear once, much like ``uniq -d``
4745 in the Unix shell or ``Itertools::duplicates`` from the Rust ``itertools``
4746 crate, pass the return value of this function into :func:`unique_everseen`
4747 with the same ``key``.
4749 """
4750 seen_set = set()
4751 seen_list = []
4752 use_key = key is not None
4754 for element in iterable:
4755 k = key(element) if use_key else element
4756 try:
4757 if k not in seen_set:
4758 seen_set.add(k)
4759 else:
4760 yield element
4761 except TypeError:
4762 if k not in seen_list:
4763 seen_list.append(k)
4764 else:
4765 yield element
4768def duplicates_justseen(iterable, key=None):
4769 """Yields serially-duplicate elements after their first appearance.
4771 >>> list(duplicates_justseen('mississippi'))
4772 ['s', 's', 'p']
4773 >>> list(duplicates_justseen('AaaBbbCccAaa', str.lower))
4774 ['a', 'a', 'b', 'b', 'c', 'c', 'a', 'a']
4776 This function is analogous to :func:`unique_justseen`.
4778 """
4779 return flatten(g for _, g in groupby(iterable, key) for _ in g)
4782def classify_unique(iterable, key=None):
4783 """Classify each element in terms of its uniqueness.
4785 For each element in the input iterable, return a 3-tuple consisting of:
4787 1. The element itself
4788 2. ``False`` if the element is equal to the one preceding it in the input,
4789 ``True`` otherwise (i.e. the equivalent of :func:`unique_justseen`)
4790 3. ``False`` if this element has been seen anywhere in the input before,
4791 ``True`` otherwise (i.e. the equivalent of :func:`unique_everseen`)
4793 >>> list(classify_unique('otto')) # doctest: +NORMALIZE_WHITESPACE
4794 [('o', True, True),
4795 ('t', True, True),
4796 ('t', False, False),
4797 ('o', True, False)]
4799 This function is analogous to :func:`unique_everseen` and is subject to
4800 the same performance considerations.
4802 """
4803 seen_set = set()
4804 seen_list = []
4805 use_key = key is not None
4806 previous = None
4808 for i, element in enumerate(iterable):
4809 k = key(element) if use_key else element
4810 is_unique_justseen = not i or previous != k
4811 previous = k
4812 is_unique_everseen = False
4813 try:
4814 if k not in seen_set:
4815 seen_set.add(k)
4816 is_unique_everseen = True
4817 except TypeError:
4818 if k not in seen_list:
4819 seen_list.append(k)
4820 is_unique_everseen = True
4821 yield element, is_unique_justseen, is_unique_everseen
4824def minmax(iterable_or_value, *others, key=None, default=_marker):
4825 """Returns both the smallest and largest items from an iterable
4826 or from two or more arguments.
4828 >>> minmax([3, 1, 5])
4829 (1, 5)
4831 >>> minmax(4, 2, 6)
4832 (2, 6)
4834 If a *key* function is provided, it will be used to transform the input
4835 items for comparison.
4837 >>> minmax([5, 30], key=str) # '30' sorts before '5'
4838 (30, 5)
4840 If a *default* value is provided, it will be returned if there are no
4841 input items.
4843 >>> minmax([], default=(0, 0))
4844 (0, 0)
4846 Otherwise ``ValueError`` is raised.
4848 This function makes a single pass over the input elements and takes care to
4849 minimize the number of comparisons made during processing.
4851 Note that unlike the builtin ``max`` function, which always returns the first
4852 item with the maximum value, this function may return another item when there are
4853 ties.
4855 This function is based on the
4856 `recipe <https://code.activestate.com/recipes/577916-fast-minmax-function>`__ by
4857 Raymond Hettinger.
4858 """
4859 iterable = (iterable_or_value, *others) if others else iterable_or_value
4861 it = iter(iterable)
4863 try:
4864 lo = hi = next(it)
4865 except StopIteration as exc:
4866 if default is _marker:
4867 raise ValueError(
4868 '`minmax()` argument is an empty iterable. '
4869 'Provide a `default` value to suppress this error.'
4870 ) from exc
4871 return default
4873 # Different branches depending on the presence of key. This saves a lot
4874 # of unimportant copies which would slow the "key=None" branch
4875 # significantly down.
4876 if key is None:
4877 for x, y in zip_longest(it, it, fillvalue=lo):
4878 if y < x:
4879 if y < lo:
4880 lo = y
4881 if hi < x:
4882 hi = x
4883 else:
4884 if x < lo:
4885 lo = x
4886 if hi < y:
4887 hi = y
4889 else:
4890 lo_key = hi_key = key(lo)
4892 for x, y in zip_longest(it, it, fillvalue=lo):
4893 x_key, y_key = key(x), key(y)
4895 if y_key < x_key:
4896 if y_key < lo_key:
4897 lo, lo_key = y, y_key
4898 if hi_key < x_key:
4899 hi, hi_key = x, x_key
4900 else:
4901 if x_key < lo_key:
4902 lo, lo_key = x, x_key
4903 if hi_key < y_key:
4904 hi, hi_key = y, y_key
4906 return lo, hi
4909def constrained_batches(
4910 iterable, max_size, max_count=None, get_len=len, strict=True
4911):
4912 """Yield batches of items from *iterable* with a combined size limited by
4913 *max_size*.
4915 >>> iterable = [b'12345', b'123', b'12345678', b'1', b'1', b'12', b'1']
4916 >>> list(constrained_batches(iterable, 10))
4917 [(b'12345', b'123'), (b'12345678', b'1', b'1'), (b'12', b'1')]
4919 If a *max_count* is supplied, the number of items per batch is also
4920 limited:
4922 >>> iterable = [b'12345', b'123', b'12345678', b'1', b'1', b'12', b'1']
4923 >>> list(constrained_batches(iterable, 10, max_count = 2))
4924 [(b'12345', b'123'), (b'12345678', b'1'), (b'1', b'12'), (b'1',)]
4926 If a *get_len* function is supplied, use that instead of :func:`len` to
4927 determine item size.
4929 If *strict* is ``True``, raise ``ValueError`` if any single item is bigger
4930 than *max_size*. Otherwise, allow single items to exceed *max_size*.
4931 """
4932 if max_size <= 0:
4933 raise ValueError('maximum size must be greater than zero')
4935 batch = []
4936 batch_size = 0
4937 batch_count = 0
4938 for item in iterable:
4939 item_len = get_len(item)
4940 if strict and item_len > max_size:
4941 raise ValueError('item size exceeds maximum size')
4943 reached_count = batch_count == max_count
4944 reached_size = item_len + batch_size > max_size
4945 if batch_count and (reached_size or reached_count):
4946 yield tuple(batch)
4947 batch.clear()
4948 batch_size = 0
4949 batch_count = 0
4951 batch.append(item)
4952 batch_size += item_len
4953 batch_count += 1
4955 if batch:
4956 yield tuple(batch)
4959def gray_product(*iterables, repeat=1):
4960 """Like :func:`itertools.product`, but return tuples in an order such
4961 that only one element in the generated tuple changes from one iteration
4962 to the next.
4964 >>> list(gray_product('AB','CD'))
4965 [('A', 'C'), ('B', 'C'), ('B', 'D'), ('A', 'D')]
4967 The *repeat* keyword argument specifies the number of repetitions
4968 of the iterables. For example, ``gray_product('AB', repeat=3)`` is
4969 equivalent to ``gray_product('AB', 'AB', 'AB')``.
4971 This function consumes all of the input iterables before producing output.
4972 If any of the input iterables have fewer than two items, ``ValueError``
4973 is raised.
4975 For information on the algorithm, see
4976 `this section <https://www-cs-faculty.stanford.edu/~knuth/fasc2a.ps.gz>`__
4977 of Donald Knuth's *The Art of Computer Programming*.
4978 """
4979 all_iterables = tuple(map(tuple, iterables)) * repeat
4980 iterable_count = len(all_iterables)
4981 for iterable in all_iterables:
4982 if len(iterable) < 2:
4983 raise ValueError("each iterable must have two or more items")
4985 # This is based on "Algorithm H" from section 7.2.1.1, page 20.
4986 # a holds the indexes of the source iterables for the n-tuple to be yielded
4987 # f is the array of "focus pointers"
4988 # o is the array of "directions"
4989 a = [0] * iterable_count
4990 f = list(range(iterable_count + 1))
4991 o = [1] * iterable_count
4992 while True:
4993 yield tuple(all_iterables[i][a[i]] for i in range(iterable_count))
4994 j = f[0]
4995 f[0] = 0
4996 if j == iterable_count:
4997 break
4998 a[j] = a[j] + o[j]
4999 if a[j] == 0 or a[j] == len(all_iterables[j]) - 1:
5000 o[j] = -o[j]
5001 f[j] = f[j + 1]
5002 f[j + 1] = j + 1
5005def partial_product(*iterables, repeat=1):
5006 """Yields tuples containing one item from each iterator, with subsequent
5007 tuples changing a single item at a time by advancing each iterator until it
5008 is exhausted. This sequence guarantees every value in each iterable is
5009 output at least once without generating all possible combinations.
5011 This may be useful, for example, when testing an expensive function.
5013 >>> list(partial_product('AB', 'C', 'DEF'))
5014 [('A', 'C', 'D'), ('B', 'C', 'D'), ('B', 'C', 'E'), ('B', 'C', 'F')]
5016 The *repeat* keyword argument specifies the number of repetitions
5017 of the iterables. For example, ``partial_product('AB', repeat=3)`` is
5018 equivalent to ``partial_product('AB', 'AB', 'AB')``.
5019 """
5021 all_iterables = tuple(map(tuple, iterables)) * repeat
5022 iterators = tuple(map(iter, all_iterables))
5024 try:
5025 prod = [next(it) for it in iterators]
5026 except StopIteration:
5027 return
5028 yield tuple(prod)
5030 for i, it in enumerate(iterators):
5031 for prod[i] in it:
5032 yield tuple(prod)
5035def takewhile_inclusive(predicate, iterable):
5036 """A variant of :func:`takewhile` that yields one additional element.
5038 >>> list(takewhile_inclusive(lambda x: x < 5, [1, 4, 6, 4, 1]))
5039 [1, 4, 6]
5041 :func:`takewhile` would return ``[1, 4]``.
5042 """
5043 for x in iterable:
5044 yield x
5045 if not predicate(x):
5046 break
5049def outer_product(func, xs, ys, *args, **kwargs):
5050 """A generalized outer product that applies a binary function to all
5051 pairs of items. Returns a 2D matrix with ``len(xs)`` rows and ``len(ys)``
5052 columns.
5053 Also accepts ``*args`` and ``**kwargs`` that are passed to ``func``.
5055 Multiplication table:
5057 >>> from operator import mul
5058 >>> list(outer_product(mul, range(1, 4), range(1, 6)))
5059 [(1, 2, 3, 4, 5), (2, 4, 6, 8, 10), (3, 6, 9, 12, 15)]
5061 Cross tabulation:
5063 >>> xs = ['A', 'B', 'A', 'A', 'B', 'B', 'A', 'A', 'B', 'B']
5064 >>> ys = ['X', 'X', 'X', 'Y', 'Z', 'Z', 'Y', 'Y', 'Z', 'Z']
5065 >>> pair_counts = Counter(zip(xs, ys))
5066 >>> count_rows = lambda x, y: pair_counts[x, y]
5067 >>> list(outer_product(count_rows, sorted(set(xs)), sorted(set(ys))))
5068 [(2, 3, 0), (1, 0, 4)]
5070 Usage with ``*args`` and ``**kwargs``:
5072 >>> animals = ['cat', 'wolf', 'mouse']
5073 >>> list(outer_product(min, animals, animals, key=len))
5074 [('cat', 'cat', 'cat'), ('cat', 'wolf', 'wolf'), ('cat', 'wolf', 'mouse')]
5075 """
5076 ys = tuple(ys)
5077 return batched(
5078 starmap(lambda x, y: func(x, y, *args, **kwargs), product(xs, ys)),
5079 n=len(ys),
5080 )
5083def iter_suppress(iterable, *exceptions):
5084 """Yield each of the items from *iterable*. If the iteration raises one of
5085 the specified *exceptions*, that exception will be suppressed and iteration
5086 will stop.
5088 >>> from itertools import chain
5089 >>> def breaks_at_five(x):
5090 ... while True:
5091 ... if x >= 5:
5092 ... raise RuntimeError
5093 ... yield x
5094 ... x += 1
5095 >>> it_1 = iter_suppress(breaks_at_five(1), RuntimeError)
5096 >>> it_2 = iter_suppress(breaks_at_five(2), RuntimeError)
5097 >>> list(chain(it_1, it_2))
5098 [1, 2, 3, 4, 2, 3, 4]
5099 """
5100 try:
5101 yield from iterable
5102 except exceptions:
5103 return
5106def filter_map(func, iterable):
5107 """Apply *func* to every element of *iterable*, yielding only those which
5108 are not ``None``.
5110 >>> elems = ['1', 'a', '2', 'b', '3']
5111 >>> list(filter_map(lambda s: int(s) if s.isnumeric() else None, elems))
5112 [1, 2, 3]
5113 """
5114 for x in iterable:
5115 y = func(x)
5116 if y is not None:
5117 yield y
5120def powerset_of_sets(iterable, *, baseset=set):
5121 """Yields all possible subsets of the iterable.
5123 >>> list(powerset_of_sets([1, 2, 3])) # doctest: +SKIP
5124 [set(), {1}, {2}, {3}, {1, 2}, {1, 3}, {2, 3}, {1, 2, 3}]
5125 >>> list(powerset_of_sets([1, 1, 0])) # doctest: +SKIP
5126 [set(), {1}, {0}, {0, 1}]
5128 :func:`powerset_of_sets` takes care to minimize the number
5129 of hash operations performed.
5131 The *baseset* parameter determines what kind of sets are
5132 constructed, either *set* or *frozenset*.
5133 """
5134 sets = tuple(dict.fromkeys(map(frozenset, zip(iterable))))
5135 union = baseset().union
5136 return chain.from_iterable(
5137 starmap(union, combinations(sets, r)) for r in range(len(sets) + 1)
5138 )
5141def join_mappings(**field_to_map):
5142 """
5143 Joins multiple mappings together using their common keys.
5145 >>> user_scores = {'elliot': 50, 'claris': 60}
5146 >>> user_times = {'elliot': 30, 'claris': 40}
5147 >>> join_mappings(score=user_scores, time=user_times)
5148 {'elliot': {'score': 50, 'time': 30}, 'claris': {'score': 60, 'time': 40}}
5149 """
5150 ret = defaultdict(dict)
5152 for field_name, mapping in field_to_map.items():
5153 for key, value in mapping.items():
5154 ret[key][field_name] = value
5156 return dict(ret)
5159def _complex_sumprod(v1, v2):
5160 """High precision sumprod() for complex numbers.
5161 Used by :func:`dft` and :func:`idft`.
5162 """
5164 real = attrgetter('real')
5165 imag = attrgetter('imag')
5166 r1 = chain(map(real, v1), map(neg, map(imag, v1)))
5167 r2 = chain(map(real, v2), map(imag, v2))
5168 i1 = chain(map(real, v1), map(imag, v1))
5169 i2 = chain(map(imag, v2), map(real, v2))
5170 return complex(_fsumprod(r1, r2), _fsumprod(i1, i2))
5173def dft(xarr):
5174 """Discrete Fourier Transform. *xarr* is a sequence of complex numbers.
5175 Yields the components of the corresponding transformed output vector.
5177 >>> import cmath
5178 >>> xarr = [1, 2-1j, -1j, -1+2j] # time domain
5179 >>> Xarr = [2, -2-2j, -2j, 4+4j] # frequency domain
5180 >>> magnitudes, phases = zip(*map(cmath.polar, Xarr))
5181 >>> all(map(cmath.isclose, dft(xarr), Xarr))
5182 True
5184 Inputs are restricted to numeric types that can add and multiply
5185 with a complex number. This includes int, float, complex, and
5186 Fraction, but excludes Decimal.
5188 See :func:`idft` for the inverse Discrete Fourier Transform.
5189 """
5190 N = len(xarr)
5191 roots_of_unity = [e ** (n / N * tau * -1j) for n in range(N)]
5192 for k in range(N):
5193 coeffs = [roots_of_unity[k * n % N] for n in range(N)]
5194 yield _complex_sumprod(xarr, coeffs)
5197def idft(Xarr):
5198 """Inverse Discrete Fourier Transform. *Xarr* is a sequence of
5199 complex numbers. Yields the components of the corresponding
5200 inverse-transformed output vector.
5202 >>> import cmath
5203 >>> xarr = [1, 2-1j, -1j, -1+2j] # time domain
5204 >>> Xarr = [2, -2-2j, -2j, 4+4j] # frequency domain
5205 >>> all(map(cmath.isclose, idft(Xarr), xarr))
5206 True
5208 Inputs are restricted to numeric types that can add and multiply
5209 with a complex number. This includes int, float, complex, and
5210 Fraction, but excludes Decimal.
5212 See :func:`dft` for the Discrete Fourier Transform.
5213 """
5214 N = len(Xarr)
5215 roots_of_unity = [e ** (n / N * tau * 1j) for n in range(N)]
5216 for k in range(N):
5217 coeffs = [roots_of_unity[k * n % N] for n in range(N)]
5218 yield _complex_sumprod(Xarr, coeffs) / N
5221def doublestarmap(func, iterable):
5222 """Apply *func* to every item of *iterable* by dictionary unpacking
5223 the item into *func*.
5225 The difference between :func:`itertools.starmap` and :func:`doublestarmap`
5226 parallels the distinction between ``func(*a)`` and ``func(**a)``.
5228 >>> iterable = [{'a': 1, 'b': 2}, {'a': 40, 'b': 60}]
5229 >>> list(doublestarmap(lambda a, b: a + b, iterable))
5230 [3, 100]
5232 ``TypeError`` will be raised if *func*'s signature doesn't match the
5233 mapping contained in *iterable* or if *iterable* does not contain mappings.
5234 """
5235 for item in iterable:
5236 yield func(**item)
5239def _nth_prime_bounds(n):
5240 """Bounds for the nth prime (counting from 1): lb < p_n < ub."""
5241 # At and above 688,383, the lb/ub spread is under 0.003 * p_n.
5243 if n < 1:
5244 raise ValueError
5246 if n < 6:
5247 return (n, 2.25 * n)
5249 # https://en.wikipedia.org/wiki/Prime-counting_function#Inequalities
5250 upper_bound = n * log(n * log(n))
5251 lower_bound = upper_bound - n
5252 if n >= 688_383:
5253 upper_bound -= n * (1.0 - (log(log(n)) - 2.0) / log(n))
5255 return lower_bound, upper_bound
5258def nth_prime(n, *, approximate=False):
5259 """Return the nth prime (counting from 0).
5261 >>> nth_prime(0)
5262 2
5263 >>> nth_prime(100)
5264 547
5266 If *approximate* is set to True, will return a prime close
5267 to the nth prime. The estimation is much faster than computing
5268 an exact result.
5270 >>> nth_prime(200_000_000, approximate=True) # Exact result is 4222234763
5271 4217820427
5273 """
5274 lb, ub = _nth_prime_bounds(n + 1)
5276 if not approximate or n <= 1_000_000:
5277 return nth(sieve(ceil(ub)), n)
5279 # Search from the midpoint and return the first odd prime
5280 odd = floor((lb + ub) / 2) | 1
5281 return first_true(count(odd, step=2), pred=is_prime)
5284def argmin(iterable, *, key=None):
5285 """
5286 Index of the first occurrence of a minimum value in an iterable.
5288 >>> argmin('efghabcdijkl')
5289 4
5290 >>> argmin([3, 2, 1, 0, 4, 2, 1, 0])
5291 3
5293 For example, look up a label corresponding to the position
5294 of a value that minimizes a cost function::
5296 >>> def cost(x):
5297 ... "Days for a wound to heal given a subject's age."
5298 ... return x**2 - 20*x + 150
5299 ...
5300 >>> labels = ['homer', 'marge', 'bart', 'lisa', 'maggie']
5301 >>> ages = [ 35, 30, 10, 9, 1 ]
5303 # Fastest healing family member
5304 >>> labels[argmin(ages, key=cost)]
5305 'bart'
5307 # Age with fastest healing
5308 >>> min(ages, key=cost)
5309 10
5311 """
5312 if key is not None:
5313 iterable = map(key, iterable)
5314 return min(enumerate(iterable), key=itemgetter(1))[0]
5317def argmax(iterable, *, key=None):
5318 """
5319 Index of the first occurrence of a maximum value in an iterable.
5321 >>> argmax('abcdefghabcd')
5322 7
5323 >>> argmax([0, 1, 2, 3, 3, 2, 1, 0])
5324 3
5326 For example, identify the best machine learning model::
5328 >>> models = ['svm', 'random forest', 'knn', 'naïve bayes']
5329 >>> accuracy = [ 68, 61, 84, 72 ]
5331 # Most accurate model
5332 >>> models[argmax(accuracy)]
5333 'knn'
5335 # Best accuracy
5336 >>> max(accuracy)
5337 84
5339 """
5340 if key is not None:
5341 iterable = map(key, iterable)
5342 return max(enumerate(iterable), key=itemgetter(1))[0]
5345def _extract_monotonic(iterator, indices):
5346 'Non-decreasing indices, lazily consumed'
5347 num_read = 0
5348 for index in indices:
5349 advance = index - num_read
5350 try:
5351 value = next(islice(iterator, advance, None))
5352 except ValueError:
5353 if advance != -1 or index < 0:
5354 raise ValueError(f'Invalid index: {index}') from None
5355 except StopIteration:
5356 raise IndexError(index) from None
5357 else:
5358 num_read += advance + 1
5359 yield value
5362def _extract_buffered(iterator, index_and_position):
5363 'Arbitrary index order, greedily consumed'
5364 buffer = {}
5365 iterator_position = -1
5366 next_to_emit = 0
5368 for index, order in index_and_position:
5369 advance = index - iterator_position
5370 if advance:
5371 try:
5372 value = next(islice(iterator, advance - 1, None))
5373 except StopIteration:
5374 raise IndexError(index) from None
5375 iterator_position = index
5377 buffer[order] = value
5379 while next_to_emit in buffer:
5380 yield buffer.pop(next_to_emit)
5381 next_to_emit += 1
5384def extract(iterable, indices, *, monotonic=False):
5385 """Yield values at the specified indices.
5387 Example:
5389 >>> data = 'abcdefghijklmnopqrstuvwxyz'
5390 >>> list(extract(data, [7, 4, 11, 11, 14]))
5391 ['h', 'e', 'l', 'l', 'o']
5393 The *iterable* is consumed lazily and can be infinite.
5395 When *monotonic* is false, the *indices* are consumed immediately
5396 and must be finite. When *monotonic* is true, *indices* are consumed
5397 lazily and can be infinite but must be non-decreasing.
5399 Raises ``IndexError`` if an index lies beyond the iterable.
5400 Raises ``ValueError`` for a negative index or for a decreasing
5401 index when *monotonic* is true.
5402 """
5404 iterator = iter(iterable)
5405 indices = iter(indices)
5407 if monotonic:
5408 return _extract_monotonic(iterator, indices)
5410 index_and_position = sorted(zip(indices, count()))
5411 if index_and_position and index_and_position[0][0] < 0:
5412 raise ValueError('Indices must be non-negative')
5413 return _extract_buffered(iterator, index_and_position)
5416class serialize:
5417 """Wrap a non-concurrent iterator with a lock to enforce sequential access.
5419 Applies a non-reentrant lock around calls to ``__next__``, allowing
5420 iterator and generator instances to be shared by multiple consumer
5421 threads.
5422 """
5424 __slots__ = ('_iterator', '_lock')
5426 def __init__(self, iterable):
5427 self._iterator = iter(iterable)
5428 self._lock = Lock()
5430 def __iter__(self):
5431 return self
5433 def __next__(self):
5434 with self._lock:
5435 return next(self._iterator)
5437 def send(self, value, /):
5438 """Send a value to a generator.
5440 Raises AttributeError if not a generator.
5441 """
5442 with self._lock:
5443 return self._iterator.send(value)
5445 def throw(self, *args):
5446 """Call throw() on a generator.
5448 Raises AttributeError if not a generator.
5449 """
5450 with self._lock:
5451 return self._iterator.throw(*args)
5453 def close(self):
5454 """Call close() on a generator.
5456 Raises AttributeError if not a generator.
5457 """
5458 with self._lock:
5459 return self._iterator.close()
5462def synchronized(func):
5463 """Wrap an iterator-returning callable to make its iterators thread-safe.
5465 Existing itertools and more-itertools can be wrapped so that their
5466 iterator instances are serialized.
5468 For example, ``itertools.count`` does not make thread-safe instances,
5469 but that is easily fixed with::
5471 atomic_counter = synchronized(itertools.count)
5473 Can also be used as a decorator for generator functions definitions
5474 so that the generator instances are serialized::
5476 @synchronized
5477 def enumerate_and_timestamp(iterable):
5478 for count, value in enumerate(iterable):
5479 yield count, time_ns(), value
5481 """
5483 @wraps(func)
5484 def inner(*args, **kwargs):
5485 iterator = func(*args, **kwargs)
5486 return serialize(iterator)
5488 return inner
5491def concurrent_tee(iterable, n=2):
5492 """Variant of itertools.tee() but with guaranteed threading semantics.
5494 Takes a non-threadsafe iterator as an input and creates concurrent
5495 tee objects for other threads to have reliable independent copies of
5496 the data stream.
5498 The new iterators are only thread-safe if consumed within a single thread.
5499 To share just one of the new iterators across multiple threads, wrap it
5500 with :func:`serialize`.
5501 """
5503 if n < 0:
5504 raise ValueError
5505 if n == 0:
5506 return ()
5507 iterator = _concurrent_tee(iterable)
5508 result = [iterator]
5509 for _ in range(n - 1):
5510 result.append(_concurrent_tee(iterator))
5511 return tuple(result)
5514class _concurrent_tee:
5515 __slots__ = ('iterator', 'link', 'lock')
5517 def __init__(self, iterable):
5518 if isinstance(iterable, _concurrent_tee):
5519 self.iterator = iterable.iterator
5520 self.link = iterable.link
5521 self.lock = iterable.lock
5522 else:
5523 self.iterator = iter(iterable)
5524 self.link = [None, None]
5525 self.lock = Lock()
5527 def __iter__(self):
5528 return self
5530 def __next__(self):
5531 link = self.link
5532 if link[1] is None:
5533 with self.lock:
5534 if link[1] is None:
5535 link[0] = next(self.iterator)
5536 link[1] = [None, None]
5537 value, self.link = link
5538 return value
5541def subfactorial(n):
5542 """Number of permutations of *n* elements with no fixed points.
5544 The :func:`subfactorial` function computes the length of
5545 :func:`derangements`. For example, there are 1,854 ways to
5546 rearrange the letters in word "epsilon" without leaving any
5547 letter in its original position:
5549 >>> from more_itertools import derangements, ilen
5550 >>> ilen(derangements('epsilon'))
5551 1854
5552 >>> subfactorial(len('epsilon'))
5553 1854
5555 Reference: https://oeis.org/A000166
5557 """
5558 if n < 0:
5559 raise ValueError
5560 sf = adj = 1
5561 for i in range(n + 1):
5562 sf = sf * i + adj
5563 adj = -adj
5564 return sf