1import warnings 2 3from collections import Counter, defaultdict, deque, abc 4from collections.abc import Sequence 5from concurrent.futures import ThreadPoolExecutor 6from functools import partial, reduce, wraps 7from heapq import merge, heapify, heapreplace, heappop 8from itertools import ( 9 chain, 10 compress, 11 count, 12 cycle, 13 dropwhile, 14 groupby, 15 islice, 16 repeat, 17 starmap, 18 takewhile, 19 tee, 20 zip_longest, 21) 22from math import exp, factorial, floor, log 23from queue import Empty, Queue 24from random import random, randrange, uniform 25from operator import itemgetter, mul, sub, gt, lt, ge, le 26from sys import hexversion, maxsize 27from time import monotonic 28 29from .recipes import ( 30 consume, 31 flatten, 32 pairwise, 33 powerset, 34 take, 35 unique_everseen, 36) 37 38__all__ = [ 39 'AbortThread', 40 'SequenceView', 41 'UnequalIterablesError', 42 'adjacent', 43 'all_unique', 44 'always_iterable', 45 'always_reversible', 46 'bucket', 47 'callback_iter', 48 'chunked', 49 'chunked_even', 50 'circular_shifts', 51 'collapse', 52 'collate', 53 'combination_index', 54 'consecutive_groups', 55 'consumer', 56 'count_cycle', 57 'countable', 58 'difference', 59 'distinct_combinations', 60 'distinct_permutations', 61 'distribute', 62 'divide', 63 'duplicates_everseen', 64 'duplicates_justseen', 65 'exactly_n', 66 'filter_except', 67 'first', 68 'groupby_transform', 69 'ichunked', 70 'ilen', 71 'interleave', 72 'interleave_evenly', 73 'interleave_longest', 74 'intersperse', 75 'is_sorted', 76 'islice_extended', 77 'iterate', 78 'last', 79 'locate', 80 'lstrip', 81 'make_decorator', 82 'map_except', 83 'map_if', 84 'map_reduce', 85 'mark_ends', 86 'minmax', 87 'nth_or_last', 88 'nth_permutation', 89 'nth_product', 90 'numeric_range', 91 'one', 92 'only', 93 'padded', 94 'partitions', 95 'peekable', 96 'permutation_index', 97 'product_index', 98 'raise_', 99 'repeat_each', 100 'repeat_last', 101 'replace', 102 'rlocate', 103 'rstrip', 104 'run_length', 105 'sample', 106 'seekable', 107 'set_partitions', 108 'side_effect', 109 'sliced', 110 'sort_together', 111 'split_after', 112 'split_at', 113 'split_before', 114 'split_into', 115 'split_when', 116 'spy', 117 'stagger', 118 'strip', 119 'strictly_n', 120 'substrings', 121 'substrings_indexes', 122 'time_limited', 123 'unique_in_window', 124 'unique_to_each', 125 'unzip', 126 'value_chain', 127 'windowed', 128 'windowed_complete', 129 'with_iter', 130 'zip_broadcast', 131 'zip_equal', 132 'zip_offset', 133] 134 135 136_marker = object() 137 138 139def chunked(iterable, n, strict=False): 140 """Break *iterable* into lists of length *n*: 141 142 >>> list(chunked([1, 2, 3, 4, 5, 6], 3)) 143 [[1, 2, 3], [4, 5, 6]] 144 145 By the default, the last yielded list will have fewer than *n* elements 146 if the length of *iterable* is not divisible by *n*: 147 148 >>> list(chunked([1, 2, 3, 4, 5, 6, 7, 8], 3)) 149 [[1, 2, 3], [4, 5, 6], [7, 8]] 150 151 To use a fill-in value instead, see the :func:`grouper` recipe. 152 153 If the length of *iterable* is not divisible by *n* and *strict* is 154 ``True``, then ``ValueError`` will be raised before the last 155 list is yielded. 156 157 """ 158 iterator = iter(partial(take, n, iter(iterable)), []) 159 if strict: 160 if n is None: 161 raise ValueError('n must not be None when using strict mode.') 162 163 def ret(): 164 for chunk in iterator: 165 if len(chunk) != n: 166 raise ValueError('iterable is not divisible by n.') 167 yield chunk 168 169 return iter(ret()) 170 else: 171 return iterator 172 173 174def first(iterable, default=_marker): 175 """Return the first item of *iterable*, or *default* if *iterable* is 176 empty. 177 178 >>> first([0, 1, 2, 3]) 179 0 180 >>> first([], 'some default') 181 'some default' 182 183 If *default* is not provided and there are no items in the iterable, 184 raise ``ValueError``. 185 186 :func:`first` is useful when you have a generator of expensive-to-retrieve 187 values and want any arbitrary one. It is marginally shorter than 188 ``next(iter(iterable), default)``. 189 190 """ 191 try: 192 return next(iter(iterable)) 193 except StopIteration as e: 194 if default is _marker: 195 raise ValueError( 196 'first() was called on an empty iterable, and no ' 197 'default value was provided.' 198 ) from e 199 return default 200 201 202def last(iterable, default=_marker): 203 """Return the last item of *iterable*, or *default* if *iterable* is 204 empty. 205 206 >>> last([0, 1, 2, 3]) 207 3 208 >>> last([], 'some default') 209 'some default' 210 211 If *default* is not provided and there are no items in the iterable, 212 raise ``ValueError``. 213 """ 214 try: 215 if isinstance(iterable, Sequence): 216 return iterable[-1] 217 # Work around https://bugs.python.org/issue38525 218 elif hasattr(iterable, '__reversed__') and (hexversion != 0x030800F0): 219 return next(reversed(iterable)) 220 else: 221 return deque(iterable, maxlen=1)[-1] 222 except (IndexError, TypeError, StopIteration): 223 if default is _marker: 224 raise ValueError( 225 'last() was called on an empty iterable, and no default was ' 226 'provided.' 227 ) 228 return default 229 230 231def nth_or_last(iterable, n, default=_marker): 232 """Return the nth or the last item of *iterable*, 233 or *default* if *iterable* is empty. 234 235 >>> nth_or_last([0, 1, 2, 3], 2) 236 2 237 >>> nth_or_last([0, 1], 2) 238 1 239 >>> nth_or_last([], 0, 'some default') 240 'some default' 241 242 If *default* is not provided and there are no items in the iterable, 243 raise ``ValueError``. 244 """ 245 return last(islice(iterable, n + 1), default=default) 246 247 248class peekable: 249 """Wrap an iterator to allow lookahead and prepending elements. 250 251 Call :meth:`peek` on the result to get the value that will be returned 252 by :func:`next`. This won't advance the iterator: 253 254 >>> p = peekable(['a', 'b']) 255 >>> p.peek() 256 'a' 257 >>> next(p) 258 'a' 259 260 Pass :meth:`peek` a default value to return that instead of raising 261 ``StopIteration`` when the iterator is exhausted. 262 263 >>> p = peekable([]) 264 >>> p.peek('hi') 265 'hi' 266 267 peekables also offer a :meth:`prepend` method, which "inserts" items 268 at the head of the iterable: 269 270 >>> p = peekable([1, 2, 3]) 271 >>> p.prepend(10, 11, 12) 272 >>> next(p) 273 10 274 >>> p.peek() 275 11 276 >>> list(p) 277 [11, 12, 1, 2, 3] 278 279 peekables can be indexed. Index 0 is the item that will be returned by 280 :func:`next`, index 1 is the item after that, and so on: 281 The values up to the given index will be cached. 282 283 >>> p = peekable(['a', 'b', 'c', 'd']) 284 >>> p[0] 285 'a' 286 >>> p[1] 287 'b' 288 >>> next(p) 289 'a' 290 291 Negative indexes are supported, but be aware that they will cache the 292 remaining items in the source iterator, which may require significant 293 storage. 294 295 To check whether a peekable is exhausted, check its truth value: 296 297 >>> p = peekable(['a', 'b']) 298 >>> if p: # peekable has items 299 ... list(p) 300 ['a', 'b'] 301 >>> if not p: # peekable is exhausted 302 ... list(p) 303 [] 304 305 """ 306 307 def __init__(self, iterable): 308 self._it = iter(iterable) 309 self._cache = deque() 310 311 def __iter__(self): 312 return self 313 314 def __bool__(self): 315 try: 316 self.peek() 317 except StopIteration: 318 return False 319 return True 320 321 def peek(self, default=_marker): 322 """Return the item that will be next returned from ``next()``. 323 324 Return ``default`` if there are no items left. If ``default`` is not 325 provided, raise ``StopIteration``. 326 327 """ 328 if not self._cache: 329 try: 330 self._cache.append(next(self._it)) 331 except StopIteration: 332 if default is _marker: 333 raise 334 return default 335 return self._cache[0] 336 337 def prepend(self, *items): 338 """Stack up items to be the next ones returned from ``next()`` or 339 ``self.peek()``. The items will be returned in 340 first in, first out order:: 341 342 >>> p = peekable([1, 2, 3]) 343 >>> p.prepend(10, 11, 12) 344 >>> next(p) 345 10 346 >>> list(p) 347 [11, 12, 1, 2, 3] 348 349 It is possible, by prepending items, to "resurrect" a peekable that 350 previously raised ``StopIteration``. 351 352 >>> p = peekable([]) 353 >>> next(p) 354 Traceback (most recent call last): 355 ... 356 StopIteration 357 >>> p.prepend(1) 358 >>> next(p) 359 1 360 >>> next(p) 361 Traceback (most recent call last): 362 ... 363 StopIteration 364 365 """ 366 self._cache.extendleft(reversed(items)) 367 368 def __next__(self): 369 if self._cache: 370 return self._cache.popleft() 371 372 return next(self._it) 373 374 def _get_slice(self, index): 375 # Normalize the slice's arguments 376 step = 1 if (index.step is None) else index.step 377 if step > 0: 378 start = 0 if (index.start is None) else index.start 379 stop = maxsize if (index.stop is None) else index.stop 380 elif step < 0: 381 start = -1 if (index.start is None) else index.start 382 stop = (-maxsize - 1) if (index.stop is None) else index.stop 383 else: 384 raise ValueError('slice step cannot be zero') 385 386 # If either the start or stop index is negative, we'll need to cache 387 # the rest of the iterable in order to slice from the right side. 388 if (start < 0) or (stop < 0): 389 self._cache.extend(self._it) 390 # Otherwise we'll need to find the rightmost index and cache to that 391 # point. 392 else: 393 n = min(max(start, stop) + 1, maxsize) 394 cache_len = len(self._cache) 395 if n >= cache_len: 396 self._cache.extend(islice(self._it, n - cache_len)) 397 398 return list(self._cache)[index] 399 400 def __getitem__(self, index): 401 if isinstance(index, slice): 402 return self._get_slice(index) 403 404 cache_len = len(self._cache) 405 if index < 0: 406 self._cache.extend(self._it) 407 elif index >= cache_len: 408 self._cache.extend(islice(self._it, index + 1 - cache_len)) 409 410 return self._cache[index] 411 412 413def collate(*iterables, **kwargs): 414 """Return a sorted merge of the items from each of several already-sorted 415 *iterables*. 416 417 >>> list(collate('ACDZ', 'AZ', 'JKL')) 418 ['A', 'A', 'C', 'D', 'J', 'K', 'L', 'Z', 'Z'] 419 420 Works lazily, keeping only the next value from each iterable in memory. Use 421 :func:`collate` to, for example, perform a n-way mergesort of items that 422 don't fit in memory. 423 424 If a *key* function is specified, the iterables will be sorted according 425 to its result: 426 427 >>> key = lambda s: int(s) # Sort by numeric value, not by string 428 >>> list(collate(['1', '10'], ['2', '11'], key=key)) 429 ['1', '2', '10', '11'] 430 431 432 If the *iterables* are sorted in descending order, set *reverse* to 433 ``True``: 434 435 >>> list(collate([5, 3, 1], [4, 2, 0], reverse=True)) 436 [5, 4, 3, 2, 1, 0] 437 438 If the elements of the passed-in iterables are out of order, you might get 439 unexpected results. 440 441 On Python 3.5+, this function is an alias for :func:`heapq.merge`. 442 443 """ 444 warnings.warn( 445 "collate is no longer part of more_itertools, use heapq.merge", 446 DeprecationWarning, 447 ) 448 return merge(*iterables, **kwargs) 449 450 451def consumer(func): 452 """Decorator that automatically advances a PEP-342-style "reverse iterator" 453 to its first yield point so you don't have to call ``next()`` on it 454 manually. 455 456 >>> @consumer 457 ... def tally(): 458 ... i = 0 459 ... while True: 460 ... print('Thing number %s is %s.' % (i, (yield))) 461 ... i += 1 462 ... 463 >>> t = tally() 464 >>> t.send('red') 465 Thing number 0 is red. 466 >>> t.send('fish') 467 Thing number 1 is fish. 468 469 Without the decorator, you would have to call ``next(t)`` before 470 ``t.send()`` could be used. 471 472 """ 473 474 @wraps(func) 475 def wrapper(*args, **kwargs): 476 gen = func(*args, **kwargs) 477 next(gen) 478 return gen 479 480 return wrapper 481 482 483def ilen(iterable): 484 """Return the number of items in *iterable*. 485 486 >>> ilen(x for x in range(1000000) if x % 3 == 0) 487 333334 488 489 This consumes the iterable, so handle with care. 490 491 """ 492 # This approach was selected because benchmarks showed it's likely the 493 # fastest of the known implementations at the time of writing. 494 # See GitHub tracker: #236, #230. 495 counter = count() 496 deque(zip(iterable, counter), maxlen=0) 497 return next(counter) 498 499 500def iterate(func, start): 501 """Return ``start``, ``func(start)``, ``func(func(start))``, ... 502 503 >>> from itertools import islice 504 >>> list(islice(iterate(lambda x: 2*x, 1), 10)) 505 [1, 2, 4, 8, 16, 32, 64, 128, 256, 512] 506 507 """ 508 while True: 509 yield start 510 start = func(start) 511 512 513def with_iter(context_manager): 514 """Wrap an iterable in a ``with`` statement, so it closes once exhausted. 515 516 For example, this will close the file when the iterator is exhausted:: 517 518 upper_lines = (line.upper() for line in with_iter(open('foo'))) 519 520 Any context manager which returns an iterable is a candidate for 521 ``with_iter``. 522 523 """ 524 with context_manager as iterable: 525 yield from iterable 526 527 528def one(iterable, too_short=None, too_long=None): 529 """Return the first item from *iterable*, which is expected to contain only 530 that item. Raise an exception if *iterable* is empty or has more than one 531 item. 532 533 :func:`one` is useful for ensuring that an iterable contains only one item. 534 For example, it can be used to retrieve the result of a database query 535 that is expected to return a single row. 536 537 If *iterable* is empty, ``ValueError`` will be raised. You may specify a 538 different exception with the *too_short* keyword: 539 540 >>> it = [] 541 >>> one(it) # doctest: +IGNORE_EXCEPTION_DETAIL 542 Traceback (most recent call last): 543 ... 544 ValueError: too many items in iterable (expected 1)' 545 >>> too_short = IndexError('too few items') 546 >>> one(it, too_short=too_short) # doctest: +IGNORE_EXCEPTION_DETAIL 547 Traceback (most recent call last): 548 ... 549 IndexError: too few items 550 551 Similarly, if *iterable* contains more than one item, ``ValueError`` will 552 be raised. You may specify a different exception with the *too_long* 553 keyword: 554 555 >>> it = ['too', 'many'] 556 >>> one(it) # doctest: +IGNORE_EXCEPTION_DETAIL 557 Traceback (most recent call last): 558 ... 559 ValueError: Expected exactly one item in iterable, but got 'too', 560 'many', and perhaps more. 561 >>> too_long = RuntimeError 562 >>> one(it, too_long=too_long) # doctest: +IGNORE_EXCEPTION_DETAIL 563 Traceback (most recent call last): 564 ... 565 RuntimeError 566 567 Note that :func:`one` attempts to advance *iterable* twice to ensure there 568 is only one item. See :func:`spy` or :func:`peekable` to check iterable 569 contents less destructively. 570 571 """ 572 it = iter(iterable) 573 574 try: 575 first_value = next(it) 576 except StopIteration as e: 577 raise ( 578 too_short or ValueError('too few items in iterable (expected 1)') 579 ) from e 580 581 try: 582 second_value = next(it) 583 except StopIteration: 584 pass 585 else: 586 msg = ( 587 'Expected exactly one item in iterable, but got {!r}, {!r}, ' 588 'and perhaps more.'.format(first_value, second_value) 589 ) 590 raise too_long or ValueError(msg) 591 592 return first_value 593 594 595def raise_(exception, *args): 596 raise exception(*args) 597 598 599def strictly_n(iterable, n, too_short=None, too_long=None): 600 """Validate that *iterable* has exactly *n* items and return them if 601 it does. If it has fewer than *n* items, call function *too_short* 602 with those items. If it has more than *n* items, call function 603 *too_long* with the first ``n + 1`` items. 604 605 >>> iterable = ['a', 'b', 'c', 'd'] 606 >>> n = 4 607 >>> list(strictly_n(iterable, n)) 608 ['a', 'b', 'c', 'd'] 609 610 By default, *too_short* and *too_long* are functions that raise 611 ``ValueError``. 612 613 >>> list(strictly_n('ab', 3)) # doctest: +IGNORE_EXCEPTION_DETAIL 614 Traceback (most recent call last): 615 ... 616 ValueError: too few items in iterable (got 2) 617 618 >>> list(strictly_n('abc', 2)) # doctest: +IGNORE_EXCEPTION_DETAIL 619 Traceback (most recent call last): 620 ... 621 ValueError: too many items in iterable (got at least 3) 622 623 You can instead supply functions that do something else. 624 *too_short* will be called with the number of items in *iterable*. 625 *too_long* will be called with `n + 1`. 626 627 >>> def too_short(item_count): 628 ... raise RuntimeError 629 >>> it = strictly_n('abcd', 6, too_short=too_short) 630 >>> list(it) # doctest: +IGNORE_EXCEPTION_DETAIL 631 Traceback (most recent call last): 632 ... 633 RuntimeError 634 635 >>> def too_long(item_count): 636 ... print('The boss is going to hear about this') 637 >>> it = strictly_n('abcdef', 4, too_long=too_long) 638 >>> list(it) 639 The boss is going to hear about this 640 ['a', 'b', 'c', 'd'] 641 642 """ 643 if too_short is None: 644 too_short = lambda item_count: raise_( 645 ValueError, 646 'Too few items in iterable (got {})'.format(item_count), 647 ) 648 649 if too_long is None: 650 too_long = lambda item_count: raise_( 651 ValueError, 652 'Too many items in iterable (got at least {})'.format(item_count), 653 ) 654 655 it = iter(iterable) 656 for i in range(n): 657 try: 658 item = next(it) 659 except StopIteration: 660 too_short(i) 661 return 662 else: 663 yield item 664 665 try: 666 next(it) 667 except StopIteration: 668 pass 669 else: 670 too_long(n + 1) 671 672 673def distinct_permutations(iterable, r=None): 674 """Yield successive distinct permutations of the elements in *iterable*. 675 676 >>> sorted(distinct_permutations([1, 0, 1])) 677 [(0, 1, 1), (1, 0, 1), (1, 1, 0)] 678 679 Equivalent to ``set(permutations(iterable))``, except duplicates are not 680 generated and thrown away. For larger input sequences this is much more 681 efficient. 682 683 Duplicate permutations arise when there are duplicated elements in the 684 input iterable. The number of items returned is 685 `n! / (x_1! * x_2! * ... * x_n!)`, where `n` is the total number of 686 items input, and each `x_i` is the count of a distinct item in the input 687 sequence. 688 689 If *r* is given, only the *r*-length permutations are yielded. 690 691 >>> sorted(distinct_permutations([1, 0, 1], r=2)) 692 [(0, 1), (1, 0), (1, 1)] 693 >>> sorted(distinct_permutations(range(3), r=2)) 694 [(0, 1), (0, 2), (1, 0), (1, 2), (2, 0), (2, 1)] 695 696 """ 697 # Algorithm: https://w.wiki/Qai 698 def _full(A): 699 while True: 700 # Yield the permutation we have 701 yield tuple(A) 702 703 # Find the largest index i such that A[i] < A[i + 1] 704 for i in range(size - 2, -1, -1): 705 if A[i] < A[i + 1]: 706 break 707 # If no such index exists, this permutation is the last one 708 else: 709 return 710 711 # Find the largest index j greater than j such that A[i] < A[j] 712 for j in range(size - 1, i, -1): 713 if A[i] < A[j]: 714 break 715 716 # Swap the value of A[i] with that of A[j], then reverse the 717 # sequence from A[i + 1] to form the new permutation 718 A[i], A[j] = A[j], A[i] 719 A[i + 1 :] = A[: i - size : -1] # A[i + 1:][::-1] 720 721 # Algorithm: modified from the above 722 def _partial(A, r): 723 # Split A into the first r items and the last r items 724 head, tail = A[:r], A[r:] 725 right_head_indexes = range(r - 1, -1, -1) 726 left_tail_indexes = range(len(tail)) 727 728 while True: 729 # Yield the permutation we have 730 yield tuple(head) 731 732 # Starting from the right, find the first index of the head with 733 # value smaller than the maximum value of the tail - call it i. 734 pivot = tail[-1] 735 for i in right_head_indexes: 736 if head[i] < pivot: 737 break 738 pivot = head[i] 739 else: 740 return 741 742 # Starting from the left, find the first value of the tail 743 # with a value greater than head[i] and swap. 744 for j in left_tail_indexes: 745 if tail[j] > head[i]: 746 head[i], tail[j] = tail[j], head[i] 747 break 748 # If we didn't find one, start from the right and find the first 749 # index of the head with a value greater than head[i] and swap. 750 else: 751 for j in right_head_indexes: 752 if head[j] > head[i]: 753 head[i], head[j] = head[j], head[i] 754 break 755 756 # Reverse head[i + 1:] and swap it with tail[:r - (i + 1)] 757 tail += head[: i - r : -1] # head[i + 1:][::-1] 758 i += 1 759 head[i:], tail[:] = tail[: r - i], tail[r - i :] 760 761 items = sorted(iterable) 762 763 size = len(items) 764 if r is None: 765 r = size 766 767 if 0 < r <= size: 768 return _full(items) if (r == size) else _partial(items, r) 769 770 return iter(() if r else ((),)) 771 772 773def intersperse(e, iterable, n=1): 774 """Intersperse filler element *e* among the items in *iterable*, leaving 775 *n* items between each filler element. 776 777 >>> list(intersperse('!', [1, 2, 3, 4, 5])) 778 [1, '!', 2, '!', 3, '!', 4, '!', 5] 779 780 >>> list(intersperse(None, [1, 2, 3, 4, 5], n=2)) 781 [1, 2, None, 3, 4, None, 5] 782 783 """ 784 if n == 0: 785 raise ValueError('n must be > 0') 786 elif n == 1: 787 # interleave(repeat(e), iterable) -> e, x_0, e, x_1, e, x_2... 788 # islice(..., 1, None) -> x_0, e, x_1, e, x_2... 789 return islice(interleave(repeat(e), iterable), 1, None) 790 else: 791 # interleave(filler, chunks) -> [e], [x_0, x_1], [e], [x_2, x_3]... 792 # islice(..., 1, None) -> [x_0, x_1], [e], [x_2, x_3]... 793 # flatten(...) -> x_0, x_1, e, x_2, x_3... 794 filler = repeat([e]) 795 chunks = chunked(iterable, n) 796 return flatten(islice(interleave(filler, chunks), 1, None)) 797 798 799def unique_to_each(*iterables): 800 """Return the elements from each of the input iterables that aren't in the 801 other input iterables. 802 803 For example, suppose you have a set of packages, each with a set of 804 dependencies:: 805 806 {'pkg_1': {'A', 'B'}, 'pkg_2': {'B', 'C'}, 'pkg_3': {'B', 'D'}} 807 808 If you remove one package, which dependencies can also be removed? 809 810 If ``pkg_1`` is removed, then ``A`` is no longer necessary - it is not 811 associated with ``pkg_2`` or ``pkg_3``. Similarly, ``C`` is only needed for 812 ``pkg_2``, and ``D`` is only needed for ``pkg_3``:: 813 814 >>> unique_to_each({'A', 'B'}, {'B', 'C'}, {'B', 'D'}) 815 [['A'], ['C'], ['D']] 816 817 If there are duplicates in one input iterable that aren't in the others 818 they will be duplicated in the output. Input order is preserved:: 819 820 >>> unique_to_each("mississippi", "missouri") 821 [['p', 'p'], ['o', 'u', 'r']] 822 823 It is assumed that the elements of each iterable are hashable. 824 825 """ 826 pool = [list(it) for it in iterables] 827 counts = Counter(chain.from_iterable(map(set, pool))) 828 uniques = {element for element in counts if counts[element] == 1} 829 return [list(filter(uniques.__contains__, it)) for it in pool] 830 831 832def windowed(seq, n, fillvalue=None, step=1): 833 """Return a sliding window of width *n* over the given iterable. 834 835 >>> all_windows = windowed([1, 2, 3, 4, 5], 3) 836 >>> list(all_windows) 837 [(1, 2, 3), (2, 3, 4), (3, 4, 5)] 838 839 When the window is larger than the iterable, *fillvalue* is used in place 840 of missing values: 841 842 >>> list(windowed([1, 2, 3], 4)) 843 [(1, 2, 3, None)] 844 845 Each window will advance in increments of *step*: 846 847 >>> list(windowed([1, 2, 3, 4, 5, 6], 3, fillvalue='!', step=2)) 848 [(1, 2, 3), (3, 4, 5), (5, 6, '!')] 849 850 To slide into the iterable's items, use :func:`chain` to add filler items 851 to the left: 852 853 >>> iterable = [1, 2, 3, 4] 854 >>> n = 3 855 >>> padding = [None] * (n - 1) 856 >>> list(windowed(chain(padding, iterable), 3)) 857 [(None, None, 1), (None, 1, 2), (1, 2, 3), (2, 3, 4)] 858 """ 859 if n < 0: 860 raise ValueError('n must be >= 0') 861 if n == 0: 862 yield tuple() 863 return 864 if step < 1: 865 raise ValueError('step must be >= 1') 866 867 window = deque(maxlen=n) 868 i = n 869 for _ in map(window.append, seq): 870 i -= 1 871 if not i: 872 i = step 873 yield tuple(window) 874 875 size = len(window) 876 if size < n: 877 yield tuple(chain(window, repeat(fillvalue, n - size))) 878 elif 0 < i < min(step, n): 879 window += (fillvalue,) * i 880 yield tuple(window) 881 882 883def substrings(iterable): 884 """Yield all of the substrings of *iterable*. 885 886 >>> [''.join(s) for s in substrings('more')] 887 ['m', 'o', 'r', 'e', 'mo', 'or', 're', 'mor', 'ore', 'more'] 888 889 Note that non-string iterables can also be subdivided. 890 891 >>> list(substrings([0, 1, 2])) 892 [(0,), (1,), (2,), (0, 1), (1, 2), (0, 1, 2)] 893 894 """ 895 # The length-1 substrings 896 seq = [] 897 for item in iter(iterable): 898 seq.append(item) 899 yield (item,) 900 seq = tuple(seq) 901 item_count = len(seq) 902 903 # And the rest 904 for n in range(2, item_count + 1): 905 for i in range(item_count - n + 1): 906 yield seq[i : i + n] 907 908 909def substrings_indexes(seq, reverse=False): 910 """Yield all substrings and their positions in *seq* 911 912 The items yielded will be a tuple of the form ``(substr, i, j)``, where 913 ``substr == seq[i:j]``. 914 915 This function only works for iterables that support slicing, such as 916 ``str`` objects. 917 918 >>> for item in substrings_indexes('more'): 919 ... print(item) 920 ('m', 0, 1) 921 ('o', 1, 2) 922 ('r', 2, 3) 923 ('e', 3, 4) 924 ('mo', 0, 2) 925 ('or', 1, 3) 926 ('re', 2, 4) 927 ('mor', 0, 3) 928 ('ore', 1, 4) 929 ('more', 0, 4) 930 931 Set *reverse* to ``True`` to yield the same items in the opposite order. 932 933 934 """ 935 r = range(1, len(seq) + 1) 936 if reverse: 937 r = reversed(r) 938 return ( 939 (seq[i : i + L], i, i + L) for L in r for i in range(len(seq) - L + 1) 940 ) 941 942 943class bucket: 944 """Wrap *iterable* and return an object that buckets it iterable into 945 child iterables based on a *key* function. 946 947 >>> iterable = ['a1', 'b1', 'c1', 'a2', 'b2', 'c2', 'b3'] 948 >>> s = bucket(iterable, key=lambda x: x[0]) # Bucket by 1st character 949 >>> sorted(list(s)) # Get the keys 950 ['a', 'b', 'c'] 951 >>> a_iterable = s['a'] 952 >>> next(a_iterable) 953 'a1' 954 >>> next(a_iterable) 955 'a2' 956 >>> list(s['b']) 957 ['b1', 'b2', 'b3'] 958 959 The original iterable will be advanced and its items will be cached until 960 they are used by the child iterables. This may require significant storage. 961 962 By default, attempting to select a bucket to which no items belong will 963 exhaust the iterable and cache all values. 964 If you specify a *validator* function, selected buckets will instead be 965 checked against it. 966 967 >>> from itertools import count 968 >>> it = count(1, 2) # Infinite sequence of odd numbers 969 >>> key = lambda x: x % 10 # Bucket by last digit 970 >>> validator = lambda x: x in {1, 3, 5, 7, 9} # Odd digits only 971 >>> s = bucket(it, key=key, validator=validator) 972 >>> 2 in s 973 False 974 >>> list(s[2]) 975 [] 976 977 """ 978 979 def __init__(self, iterable, key, validator=None): 980 self._it = iter(iterable) 981 self._key = key 982 self._cache = defaultdict(deque) 983 self._validator = validator or (lambda x: True) 984 985 def __contains__(self, value): 986 if not self._validator(value): 987 return False 988 989 try: 990 item = next(self[value]) 991 except StopIteration: 992 return False 993 else: 994 self._cache[value].appendleft(item) 995 996 return True 997 998 def _get_values(self, value): 999 """ 1000 Helper to yield items from the parent iterator that match *value*. 1001 Items that don't match are stored in the local cache as they 1002 are encountered. 1003 """ 1004 while True: 1005 # If we've cached some items that match the target value, emit 1006 # the first one and evict it from the cache. 1007 if self._cache[value]: 1008 yield self._cache[value].popleft() 1009 # Otherwise we need to advance the parent iterator to search for 1010 # a matching item, caching the rest. 1011 else: 1012 while True: 1013 try: 1014 item = next(self._it) 1015 except StopIteration: 1016 return 1017 item_value = self._key(item) 1018 if item_value == value: 1019 yield item 1020 break 1021 elif self._validator(item_value): 1022 self._cache[item_value].append(item) 1023 1024 def __iter__(self): 1025 for item in self._it: 1026 item_value = self._key(item) 1027 if self._validator(item_value): 1028 self._cache[item_value].append(item) 1029 1030 yield from self._cache.keys() 1031 1032 def __getitem__(self, value): 1033 if not self._validator(value): 1034 return iter(()) 1035 1036 return self._get_values(value) 1037 1038 1039def spy(iterable, n=1): 1040 """Return a 2-tuple with a list containing the first *n* elements of 1041 *iterable*, and an iterator with the same items as *iterable*. 1042 This allows you to "look ahead" at the items in the iterable without 1043 advancing it. 1044 1045 There is one item in the list by default: 1046 1047 >>> iterable = 'abcdefg' 1048 >>> head, iterable = spy(iterable) 1049 >>> head 1050 ['a'] 1051 >>> list(iterable) 1052 ['a', 'b', 'c', 'd', 'e', 'f', 'g'] 1053 1054 You may use unpacking to retrieve items instead of lists: 1055 1056 >>> (head,), iterable = spy('abcdefg') 1057 >>> head 1058 'a' 1059 >>> (first, second), iterable = spy('abcdefg', 2) 1060 >>> first 1061 'a' 1062 >>> second 1063 'b' 1064 1065 The number of items requested can be larger than the number of items in 1066 the iterable: 1067 1068 >>> iterable = [1, 2, 3, 4, 5] 1069 >>> head, iterable = spy(iterable, 10) 1070 >>> head 1071 [1, 2, 3, 4, 5] 1072 >>> list(iterable) 1073 [1, 2, 3, 4, 5] 1074 1075 """ 1076 it = iter(iterable) 1077 head = take(n, it) 1078 1079 return head.copy(), chain(head, it) 1080 1081 1082def interleave(*iterables): 1083 """Return a new iterable yielding from each iterable in turn, 1084 until the shortest is exhausted. 1085 1086 >>> list(interleave([1, 2, 3], [4, 5], [6, 7, 8])) 1087 [1, 4, 6, 2, 5, 7] 1088 1089 For a version that doesn't terminate after the shortest iterable is 1090 exhausted, see :func:`interleave_longest`. 1091 1092 """ 1093 return chain.from_iterable(zip(*iterables)) 1094 1095 1096def interleave_longest(*iterables): 1097 """Return a new iterable yielding from each iterable in turn, 1098 skipping any that are exhausted. 1099 1100 >>> list(interleave_longest([1, 2, 3], [4, 5], [6, 7, 8])) 1101 [1, 4, 6, 2, 5, 7, 3, 8] 1102 1103 This function produces the same output as :func:`roundrobin`, but may 1104 perform better for some inputs (in particular when the number of iterables 1105 is large). 1106 1107 """ 1108 i = chain.from_iterable(zip_longest(*iterables, fillvalue=_marker)) 1109 return (x for x in i if x is not _marker) 1110 1111 1112def interleave_evenly(iterables, lengths=None): 1113 """ 1114 Interleave multiple iterables so that their elements are evenly distributed 1115 throughout the output sequence. 1116 1117 >>> iterables = [1, 2, 3, 4, 5], ['a', 'b'] 1118 >>> list(interleave_evenly(iterables)) 1119 [1, 2, 'a', 3, 4, 'b', 5] 1120 1121 >>> iterables = [[1, 2, 3], [4, 5], [6, 7, 8]] 1122 >>> list(interleave_evenly(iterables)) 1123 [1, 6, 4, 2, 7, 3, 8, 5] 1124 1125 This function requires iterables of known length. Iterables without 1126 ``__len__()`` can be used by manually specifying lengths with *lengths*: 1127 1128 >>> from itertools import combinations, repeat 1129 >>> iterables = [combinations(range(4), 2), ['a', 'b', 'c']] 1130 >>> lengths = [4 * (4 - 1) // 2, 3] 1131 >>> list(interleave_evenly(iterables, lengths=lengths)) 1132 [(0, 1), (0, 2), 'a', (0, 3), (1, 2), 'b', (1, 3), (2, 3), 'c'] 1133 1134 Based on Bresenham's algorithm. 1135 """ 1136 if lengths is None: 1137 try: 1138 lengths = [len(it) for it in iterables] 1139 except TypeError: 1140 raise ValueError( 1141 'Iterable lengths could not be determined automatically. ' 1142 'Specify them with the lengths keyword.' 1143 ) 1144 elif len(iterables) != len(lengths): 1145 raise ValueError('Mismatching number of iterables and lengths.') 1146 1147 dims = len(lengths) 1148 1149 # sort iterables by length, descending 1150 lengths_permute = sorted( 1151 range(dims), key=lambda i: lengths[i], reverse=True 1152 ) 1153 lengths_desc = [lengths[i] for i in lengths_permute] 1154 iters_desc = [iter(iterables[i]) for i in lengths_permute] 1155 1156 # the longest iterable is the primary one (Bresenham: the longest 1157 # distance along an axis) 1158 delta_primary, deltas_secondary = lengths_desc[0], lengths_desc[1:] 1159 iter_primary, iters_secondary = iters_desc[0], iters_desc[1:] 1160 errors = [delta_primary // dims] * len(deltas_secondary) 1161 1162 to_yield = sum(lengths) 1163 while to_yield: 1164 yield next(iter_primary) 1165 to_yield -= 1 1166 # update errors for each secondary iterable 1167 errors = [e - delta for e, delta in zip(errors, deltas_secondary)] 1168 1169 # those iterables for which the error is negative are yielded 1170 # ("diagonal step" in Bresenham) 1171 for i, e in enumerate(errors): 1172 if e < 0: 1173 yield next(iters_secondary[i]) 1174 to_yield -= 1 1175 errors[i] += delta_primary 1176 1177 1178def collapse(iterable, base_type=None, levels=None): 1179 """Flatten an iterable with multiple levels of nesting (e.g., a list of 1180 lists of tuples) into non-iterable types. 1181 1182 >>> iterable = [(1, 2), ([3, 4], [[5], [6]])] 1183 >>> list(collapse(iterable)) 1184 [1, 2, 3, 4, 5, 6] 1185 1186 Binary and text strings are not considered iterable and 1187 will not be collapsed. 1188 1189 To avoid collapsing other types, specify *base_type*: 1190 1191 >>> iterable = ['ab', ('cd', 'ef'), ['gh', 'ij']] 1192 >>> list(collapse(iterable, base_type=tuple)) 1193 ['ab', ('cd', 'ef'), 'gh', 'ij'] 1194 1195 Specify *levels* to stop flattening after a certain level: 1196 1197 >>> iterable = [('a', ['b']), ('c', ['d'])] 1198 >>> list(collapse(iterable)) # Fully flattened 1199 ['a', 'b', 'c', 'd'] 1200 >>> list(collapse(iterable, levels=1)) # Only one level flattened 1201 ['a', ['b'], 'c', ['d']] 1202 1203 """ 1204 1205 def walk(node, level): 1206 if ( 1207 ((levels is not None) and (level > levels)) 1208 or isinstance(node, (str, bytes)) 1209 or ((base_type is not None) and isinstance(node, base_type)) 1210 ): 1211 yield node 1212 return 1213 1214 try: 1215 tree = iter(node) 1216 except TypeError: 1217 yield node 1218 return 1219 else: 1220 for child in tree: 1221 yield from walk(child, level + 1) 1222 1223 yield from walk(iterable, 0) 1224 1225 1226def side_effect(func, iterable, chunk_size=None, before=None, after=None): 1227 """Invoke *func* on each item in *iterable* (or on each *chunk_size* group 1228 of items) before yielding the item. 1229 1230 `func` must be a function that takes a single argument. Its return value 1231 will be discarded. 1232 1233 *before* and *after* are optional functions that take no arguments. They 1234 will be executed before iteration starts and after it ends, respectively. 1235 1236 `side_effect` can be used for logging, updating progress bars, or anything 1237 that is not functionally "pure." 1238 1239 Emitting a status message: 1240 1241 >>> from more_itertools import consume 1242 >>> func = lambda item: print('Received {}'.format(item)) 1243 >>> consume(side_effect(func, range(2))) 1244 Received 0 1245 Received 1 1246 1247 Operating on chunks of items: 1248 1249 >>> pair_sums = [] 1250 >>> func = lambda chunk: pair_sums.append(sum(chunk)) 1251 >>> list(side_effect(func, [0, 1, 2, 3, 4, 5], 2)) 1252 [0, 1, 2, 3, 4, 5] 1253 >>> list(pair_sums) 1254 [1, 5, 9] 1255 1256 Writing to a file-like object: 1257 1258 >>> from io import StringIO 1259 >>> from more_itertools import consume 1260 >>> f = StringIO() 1261 >>> func = lambda x: print(x, file=f) 1262 >>> before = lambda: print(u'HEADER', file=f) 1263 >>> after = f.close 1264 >>> it = [u'a', u'b', u'c'] 1265 >>> consume(side_effect(func, it, before=before, after=after)) 1266 >>> f.closed 1267 True 1268 1269 """ 1270 try: 1271 if before is not None: 1272 before() 1273 1274 if chunk_size is None: 1275 for item in iterable: 1276 func(item) 1277 yield item 1278 else: 1279 for chunk in chunked(iterable, chunk_size): 1280 func(chunk) 1281 yield from chunk 1282 finally: 1283 if after is not None: 1284 after() 1285 1286 1287def sliced(seq, n, strict=False): 1288 """Yield slices of length *n* from the sequence *seq*. 1289 1290 >>> list(sliced((1, 2, 3, 4, 5, 6), 3)) 1291 [(1, 2, 3), (4, 5, 6)] 1292 1293 By the default, the last yielded slice will have fewer than *n* elements 1294 if the length of *seq* is not divisible by *n*: 1295 1296 >>> list(sliced((1, 2, 3, 4, 5, 6, 7, 8), 3)) 1297 [(1, 2, 3), (4, 5, 6), (7, 8)] 1298 1299 If the length of *seq* is not divisible by *n* and *strict* is 1300 ``True``, then ``ValueError`` will be raised before the last 1301 slice is yielded. 1302 1303 This function will only work for iterables that support slicing. 1304 For non-sliceable iterables, see :func:`chunked`. 1305 1306 """ 1307 iterator = takewhile(len, (seq[i : i + n] for i in count(0, n))) 1308 if strict: 1309 1310 def ret(): 1311 for _slice in iterator: 1312 if len(_slice) != n: 1313 raise ValueError("seq is not divisible by n.") 1314 yield _slice 1315 1316 return iter(ret()) 1317 else: 1318 return iterator 1319 1320 1321def split_at(iterable, pred, maxsplit=-1, keep_separator=False): 1322 """Yield lists of items from *iterable*, where each list is delimited by 1323 an item where callable *pred* returns ``True``. 1324 1325 >>> list(split_at('abcdcba', lambda x: x == 'b')) 1326 [['a'], ['c', 'd', 'c'], ['a']] 1327 1328 >>> list(split_at(range(10), lambda n: n % 2 == 1)) 1329 [[0], [2], [4], [6], [8], []] 1330 1331 At most *maxsplit* splits are done. If *maxsplit* is not specified or -1, 1332 then there is no limit on the number of splits: 1333 1334 >>> list(split_at(range(10), lambda n: n % 2 == 1, maxsplit=2)) 1335 [[0], [2], [4, 5, 6, 7, 8, 9]] 1336 1337 By default, the delimiting items are not included in the output. 1338 The include them, set *keep_separator* to ``True``. 1339 1340 >>> list(split_at('abcdcba', lambda x: x == 'b', keep_separator=True)) 1341 [['a'], ['b'], ['c', 'd', 'c'], ['b'], ['a']] 1342 1343 """ 1344 if maxsplit == 0: 1345 yield list(iterable) 1346 return 1347 1348 buf = [] 1349 it = iter(iterable) 1350 for item in it: 1351 if pred(item): 1352 yield buf 1353 if keep_separator: 1354 yield [item] 1355 if maxsplit == 1: 1356 yield list(it) 1357 return 1358 buf = [] 1359 maxsplit -= 1 1360 else: 1361 buf.append(item) 1362 yield buf 1363 1364 1365def split_before(iterable, pred, maxsplit=-1): 1366 """Yield lists of items from *iterable*, where each list ends just before 1367 an item for which callable *pred* returns ``True``: 1368 1369 >>> list(split_before('OneTwo', lambda s: s.isupper())) 1370 [['O', 'n', 'e'], ['T', 'w', 'o']] 1371 1372 >>> list(split_before(range(10), lambda n: n % 3 == 0)) 1373 [[0, 1, 2], [3, 4, 5], [6, 7, 8], [9]] 1374 1375 At most *maxsplit* splits are done. If *maxsplit* is not specified or -1, 1376 then there is no limit on the number of splits: 1377 1378 >>> list(split_before(range(10), lambda n: n % 3 == 0, maxsplit=2)) 1379 [[0, 1, 2], [3, 4, 5], [6, 7, 8, 9]] 1380 """ 1381 if maxsplit == 0: 1382 yield list(iterable) 1383 return 1384 1385 buf = [] 1386 it = iter(iterable) 1387 for item in it: 1388 if pred(item) and buf: 1389 yield buf 1390 if maxsplit == 1: 1391 yield [item] + list(it) 1392 return 1393 buf = [] 1394 maxsplit -= 1 1395 buf.append(item) 1396 if buf: 1397 yield buf 1398 1399 1400def split_after(iterable, pred, maxsplit=-1): 1401 """Yield lists of items from *iterable*, where each list ends with an 1402 item where callable *pred* returns ``True``: 1403 1404 >>> list(split_after('one1two2', lambda s: s.isdigit())) 1405 [['o', 'n', 'e', '1'], ['t', 'w', 'o', '2']] 1406 1407 >>> list(split_after(range(10), lambda n: n % 3 == 0)) 1408 [[0], [1, 2, 3], [4, 5, 6], [7, 8, 9]] 1409 1410 At most *maxsplit* splits are done. If *maxsplit* is not specified or -1, 1411 then there is no limit on the number of splits: 1412 1413 >>> list(split_after(range(10), lambda n: n % 3 == 0, maxsplit=2)) 1414 [[0], [1, 2, 3], [4, 5, 6, 7, 8, 9]] 1415 1416 """ 1417 if maxsplit == 0: 1418 yield list(iterable) 1419 return 1420 1421 buf = [] 1422 it = iter(iterable) 1423 for item in it: 1424 buf.append(item) 1425 if pred(item) and buf: 1426 yield buf 1427 if maxsplit == 1: 1428 yield list(it) 1429 return 1430 buf = [] 1431 maxsplit -= 1 1432 if buf: 1433 yield buf 1434 1435 1436def split_when(iterable, pred, maxsplit=-1): 1437 """Split *iterable* into pieces based on the output of *pred*. 1438 *pred* should be a function that takes successive pairs of items and 1439 returns ``True`` if the iterable should be split in between them. 1440 1441 For example, to find runs of increasing numbers, split the iterable when 1442 element ``i`` is larger than element ``i + 1``: 1443 1444 >>> list(split_when([1, 2, 3, 3, 2, 5, 2, 4, 2], lambda x, y: x > y)) 1445 [[1, 2, 3, 3], [2, 5], [2, 4], [2]] 1446 1447 At most *maxsplit* splits are done. If *maxsplit* is not specified or -1, 1448 then there is no limit on the number of splits: 1449 1450 >>> list(split_when([1, 2, 3, 3, 2, 5, 2, 4, 2], 1451 ... lambda x, y: x > y, maxsplit=2)) 1452 [[1, 2, 3, 3], [2, 5], [2, 4, 2]] 1453 1454 """ 1455 if maxsplit == 0: 1456 yield list(iterable) 1457 return 1458 1459 it = iter(iterable) 1460 try: 1461 cur_item = next(it) 1462 except StopIteration: 1463 return 1464 1465 buf = [cur_item] 1466 for next_item in it: 1467 if pred(cur_item, next_item): 1468 yield buf 1469 if maxsplit == 1: 1470 yield [next_item] + list(it) 1471 return 1472 buf = [] 1473 maxsplit -= 1 1474 1475 buf.append(next_item) 1476 cur_item = next_item 1477 1478 yield buf 1479 1480 1481def split_into(iterable, sizes): 1482 """Yield a list of sequential items from *iterable* of length 'n' for each 1483 integer 'n' in *sizes*. 1484 1485 >>> list(split_into([1,2,3,4,5,6], [1,2,3])) 1486 [[1], [2, 3], [4, 5, 6]] 1487 1488 If the sum of *sizes* is smaller than the length of *iterable*, then the 1489 remaining items of *iterable* will not be returned. 1490 1491 >>> list(split_into([1,2,3,4,5,6], [2,3])) 1492 [[1, 2], [3, 4, 5]] 1493 1494 If the sum of *sizes* is larger than the length of *iterable*, fewer items 1495 will be returned in the iteration that overruns *iterable* and further 1496 lists will be empty: 1497 1498 >>> list(split_into([1,2,3,4], [1,2,3,4])) 1499 [[1], [2, 3], [4], []] 1500 1501 When a ``None`` object is encountered in *sizes*, the returned list will 1502 contain items up to the end of *iterable* the same way that itertools.slice 1503 does: 1504 1505 >>> list(split_into([1,2,3,4,5,6,7,8,9,0], [2,3,None])) 1506 [[1, 2], [3, 4, 5], [6, 7, 8, 9, 0]] 1507 1508 :func:`split_into` can be useful for grouping a series of items where the 1509 sizes of the groups are not uniform. An example would be where in a row 1510 from a table, multiple columns represent elements of the same feature 1511 (e.g. a point represented by x,y,z) but, the format is not the same for 1512 all columns. 1513 """ 1514 # convert the iterable argument into an iterator so its contents can 1515 # be consumed by islice in case it is a generator 1516 it = iter(iterable) 1517 1518 for size in sizes: 1519 if size is None: 1520 yield list(it) 1521 return 1522 else: 1523 yield list(islice(it, size)) 1524 1525 1526def padded(iterable, fillvalue=None, n=None, next_multiple=False): 1527 """Yield the elements from *iterable*, followed by *fillvalue*, such that 1528 at least *n* items are emitted. 1529 1530 >>> list(padded([1, 2, 3], '?', 5)) 1531 [1, 2, 3, '?', '?'] 1532 1533 If *next_multiple* is ``True``, *fillvalue* will be emitted until the 1534 number of items emitted is a multiple of *n*:: 1535 1536 >>> list(padded([1, 2, 3, 4], n=3, next_multiple=True)) 1537 [1, 2, 3, 4, None, None] 1538 1539 If *n* is ``None``, *fillvalue* will be emitted indefinitely. 1540 1541 """ 1542 it = iter(iterable) 1543 if n is None: 1544 yield from chain(it, repeat(fillvalue)) 1545 elif n < 1: 1546 raise ValueError('n must be at least 1') 1547 else: 1548 item_count = 0 1549 for item in it: 1550 yield item 1551 item_count += 1 1552 1553 remaining = (n - item_count) % n if next_multiple else n - item_count 1554 for _ in range(remaining): 1555 yield fillvalue 1556 1557 1558def repeat_each(iterable, n=2): 1559 """Repeat each element in *iterable* *n* times. 1560 1561 >>> list(repeat_each('ABC', 3)) 1562 ['A', 'A', 'A', 'B', 'B', 'B', 'C', 'C', 'C'] 1563 """ 1564 return chain.from_iterable(map(repeat, iterable, repeat(n))) 1565 1566 1567def repeat_last(iterable, default=None): 1568 """After the *iterable* is exhausted, keep yielding its last element. 1569 1570 >>> list(islice(repeat_last(range(3)), 5)) 1571 [0, 1, 2, 2, 2] 1572 1573 If the iterable is empty, yield *default* forever:: 1574 1575 >>> list(islice(repeat_last(range(0), 42), 5)) 1576 [42, 42, 42, 42, 42] 1577 1578 """ 1579 item = _marker 1580 for item in iterable: 1581 yield item 1582 final = default if item is _marker else item 1583 yield from repeat(final) 1584 1585 1586def distribute(n, iterable): 1587 """Distribute the items from *iterable* among *n* smaller iterables. 1588 1589 >>> group_1, group_2 = distribute(2, [1, 2, 3, 4, 5, 6]) 1590 >>> list(group_1) 1591 [1, 3, 5] 1592 >>> list(group_2) 1593 [2, 4, 6] 1594 1595 If the length of *iterable* is not evenly divisible by *n*, then the 1596 length of the returned iterables will not be identical: 1597 1598 >>> children = distribute(3, [1, 2, 3, 4, 5, 6, 7]) 1599 >>> [list(c) for c in children] 1600 [[1, 4, 7], [2, 5], [3, 6]] 1601 1602 If the length of *iterable* is smaller than *n*, then the last returned 1603 iterables will be empty: 1604 1605 >>> children = distribute(5, [1, 2, 3]) 1606 >>> [list(c) for c in children] 1607 [[1], [2], [3], [], []] 1608 1609 This function uses :func:`itertools.tee` and may require significant 1610 storage. If you need the order items in the smaller iterables to match the 1611 original iterable, see :func:`divide`. 1612 1613 """ 1614 if n < 1: 1615 raise ValueError('n must be at least 1') 1616 1617 children = tee(iterable, n) 1618 return [islice(it, index, None, n) for index, it in enumerate(children)] 1619 1620 1621def stagger(iterable, offsets=(-1, 0, 1), longest=False, fillvalue=None): 1622 """Yield tuples whose elements are offset from *iterable*. 1623 The amount by which the `i`-th item in each tuple is offset is given by 1624 the `i`-th item in *offsets*. 1625 1626 >>> list(stagger([0, 1, 2, 3])) 1627 [(None, 0, 1), (0, 1, 2), (1, 2, 3)] 1628 >>> list(stagger(range(8), offsets=(0, 2, 4))) 1629 [(0, 2, 4), (1, 3, 5), (2, 4, 6), (3, 5, 7)] 1630 1631 By default, the sequence will end when the final element of a tuple is the 1632 last item in the iterable. To continue until the first element of a tuple 1633 is the last item in the iterable, set *longest* to ``True``:: 1634 1635 >>> list(stagger([0, 1, 2, 3], longest=True)) 1636 [(None, 0, 1), (0, 1, 2), (1, 2, 3), (2, 3, None), (3, None, None)] 1637 1638 By default, ``None`` will be used to replace offsets beyond the end of the 1639 sequence. Specify *fillvalue* to use some other value. 1640 1641 """ 1642 children = tee(iterable, len(offsets)) 1643 1644 return zip_offset( 1645 *children, offsets=offsets, longest=longest, fillvalue=fillvalue 1646 ) 1647 1648 1649class UnequalIterablesError(ValueError): 1650 def __init__(self, details=None): 1651 msg = 'Iterables have different lengths' 1652 if details is not None: 1653 msg += (': index 0 has length {}; index {} has length {}').format( 1654 *details 1655 ) 1656 1657 super().__init__(msg) 1658 1659 1660def _zip_equal_generator(iterables): 1661 for combo in zip_longest(*iterables, fillvalue=_marker): 1662 for val in combo: 1663 if val is _marker: 1664 raise UnequalIterablesError() 1665 yield combo 1666 1667 1668def _zip_equal(*iterables): 1669 # Check whether the iterables are all the same size. 1670 try: 1671 first_size = len(iterables[0]) 1672 for i, it in enumerate(iterables[1:], 1): 1673 size = len(it) 1674 if size != first_size: 1675 break 1676 else: 1677 # If we didn't break out, we can use the built-in zip. 1678 return zip(*iterables) 1679 1680 # If we did break out, there was a mismatch. 1681 raise UnequalIterablesError(details=(first_size, i, size)) 1682 # If any one of the iterables didn't have a length, start reading 1683 # them until one runs out. 1684 except TypeError: 1685 return _zip_equal_generator(iterables) 1686 1687 1688def zip_equal(*iterables): 1689 """``zip`` the input *iterables* together, but raise 1690 ``UnequalIterablesError`` if they aren't all the same length. 1691 1692 >>> it_1 = range(3) 1693 >>> it_2 = iter('abc') 1694 >>> list(zip_equal(it_1, it_2)) 1695 [(0, 'a'), (1, 'b'), (2, 'c')] 1696 1697 >>> it_1 = range(3) 1698 >>> it_2 = iter('abcd') 1699 >>> list(zip_equal(it_1, it_2)) # doctest: +IGNORE_EXCEPTION_DETAIL 1700 Traceback (most recent call last): 1701 ... 1702 more_itertools.more.UnequalIterablesError: Iterables have different 1703 lengths 1704 1705 """ 1706 if hexversion >= 0x30A00A6: 1707 warnings.warn( 1708 ( 1709 'zip_equal will be removed in a future version of ' 1710 'more-itertools. Use the builtin zip function with ' 1711 'strict=True instead.' 1712 ), 1713 DeprecationWarning, 1714 ) 1715 1716 return _zip_equal(*iterables) 1717 1718 1719def zip_offset(*iterables, offsets, longest=False, fillvalue=None): 1720 """``zip`` the input *iterables* together, but offset the `i`-th iterable 1721 by the `i`-th item in *offsets*. 1722 1723 >>> list(zip_offset('0123', 'abcdef', offsets=(0, 1))) 1724 [('0', 'b'), ('1', 'c'), ('2', 'd'), ('3', 'e')] 1725 1726 This can be used as a lightweight alternative to SciPy or pandas to analyze 1727 data sets in which some series have a lead or lag relationship. 1728 1729 By default, the sequence will end when the shortest iterable is exhausted. 1730 To continue until the longest iterable is exhausted, set *longest* to 1731 ``True``. 1732 1733 >>> list(zip_offset('0123', 'abcdef', offsets=(0, 1), longest=True)) 1734 [('0', 'b'), ('1', 'c'), ('2', 'd'), ('3', 'e'), (None, 'f')] 1735 1736 By default, ``None`` will be used to replace offsets beyond the end of the 1737 sequence. Specify *fillvalue* to use some other value. 1738 1739 """ 1740 if len(iterables) != len(offsets): 1741 raise ValueError("Number of iterables and offsets didn't match") 1742 1743 staggered = [] 1744 for it, n in zip(iterables, offsets): 1745 if n < 0: 1746 staggered.append(chain(repeat(fillvalue, -n), it)) 1747 elif n > 0: 1748 staggered.append(islice(it, n, None)) 1749 else: 1750 staggered.append(it) 1751 1752 if longest: 1753 return zip_longest(*staggered, fillvalue=fillvalue) 1754 1755 return zip(*staggered) 1756 1757 1758def sort_together(iterables, key_list=(0,), key=None, reverse=False): 1759 """Return the input iterables sorted together, with *key_list* as the 1760 priority for sorting. All iterables are trimmed to the length of the 1761 shortest one. 1762 1763 This can be used like the sorting function in a spreadsheet. If each 1764 iterable represents a column of data, the key list determines which 1765 columns are used for sorting. 1766 1767 By default, all iterables are sorted using the ``0``-th iterable:: 1768 1769 >>> iterables = [(4, 3, 2, 1), ('a', 'b', 'c', 'd')] 1770 >>> sort_together(iterables) 1771 [(1, 2, 3, 4), ('d', 'c', 'b', 'a')] 1772 1773 Set a different key list to sort according to another iterable. 1774 Specifying multiple keys dictates how ties are broken:: 1775 1776 >>> iterables = [(3, 1, 2), (0, 1, 0), ('c', 'b', 'a')] 1777 >>> sort_together(iterables, key_list=(1, 2)) 1778 [(2, 3, 1), (0, 0, 1), ('a', 'c', 'b')] 1779 1780 To sort by a function of the elements of the iterable, pass a *key* 1781 function. Its arguments are the elements of the iterables corresponding to 1782 the key list:: 1783 1784 >>> names = ('a', 'b', 'c') 1785 >>> lengths = (1, 2, 3) 1786 >>> widths = (5, 2, 1) 1787 >>> def area(length, width): 1788 ... return length * width 1789 >>> sort_together([names, lengths, widths], key_list=(1, 2), key=area) 1790 [('c', 'b', 'a'), (3, 2, 1), (1, 2, 5)] 1791 1792 Set *reverse* to ``True`` to sort in descending order. 1793 1794 >>> sort_together([(1, 2, 3), ('c', 'b', 'a')], reverse=True) 1795 [(3, 2, 1), ('a', 'b', 'c')] 1796 1797 """ 1798 if key is None: 1799 # if there is no key function, the key argument to sorted is an 1800 # itemgetter 1801 key_argument = itemgetter(*key_list) 1802 else: 1803 # if there is a key function, call it with the items at the offsets 1804 # specified by the key function as arguments 1805 key_list = list(key_list) 1806 if len(key_list) == 1: 1807 # if key_list contains a single item, pass the item at that offset 1808 # as the only argument to the key function 1809 key_offset = key_list[0] 1810 key_argument = lambda zipped_items: key(zipped_items[key_offset]) 1811 else: 1812 # if key_list contains multiple items, use itemgetter to return a 1813 # tuple of items, which we pass as *args to the key function 1814 get_key_items = itemgetter(*key_list) 1815 key_argument = lambda zipped_items: key( 1816 *get_key_items(zipped_items) 1817 ) 1818 1819 return list( 1820 zip(*sorted(zip(*iterables), key=key_argument, reverse=reverse)) 1821 ) 1822 1823 1824def unzip(iterable): 1825 """The inverse of :func:`zip`, this function disaggregates the elements 1826 of the zipped *iterable*. 1827 1828 The ``i``-th iterable contains the ``i``-th element from each element 1829 of the zipped iterable. The first element is used to to determine the 1830 length of the remaining elements. 1831 1832 >>> iterable = [('a', 1), ('b', 2), ('c', 3), ('d', 4)] 1833 >>> letters, numbers = unzip(iterable) 1834 >>> list(letters) 1835 ['a', 'b', 'c', 'd'] 1836 >>> list(numbers) 1837 [1, 2, 3, 4] 1838 1839 This is similar to using ``zip(*iterable)``, but it avoids reading 1840 *iterable* into memory. Note, however, that this function uses 1841 :func:`itertools.tee` and thus may require significant storage. 1842 1843 """ 1844 head, iterable = spy(iter(iterable)) 1845 if not head: 1846 # empty iterable, e.g. zip([], [], []) 1847 return () 1848 # spy returns a one-length iterable as head 1849 head = head[0] 1850 iterables = tee(iterable, len(head)) 1851 1852 def itemgetter(i): 1853 def getter(obj): 1854 try: 1855 return obj[i] 1856 except IndexError: 1857 # basically if we have an iterable like 1858 # iter([(1, 2, 3), (4, 5), (6,)]) 1859 # the second unzipped iterable would fail at the third tuple 1860 # since it would try to access tup[1] 1861 # same with the third unzipped iterable and the second tuple 1862 # to support these "improperly zipped" iterables, 1863 # we create a custom itemgetter 1864 # which just stops the unzipped iterables 1865 # at first length mismatch 1866 raise StopIteration 1867 1868 return getter 1869 1870 return tuple(map(itemgetter(i), it) for i, it in enumerate(iterables)) 1871 1872 1873def divide(n, iterable): 1874 """Divide the elements from *iterable* into *n* parts, maintaining 1875 order. 1876 1877 >>> group_1, group_2 = divide(2, [1, 2, 3, 4, 5, 6]) 1878 >>> list(group_1) 1879 [1, 2, 3] 1880 >>> list(group_2) 1881 [4, 5, 6] 1882 1883 If the length of *iterable* is not evenly divisible by *n*, then the 1884 length of the returned iterables will not be identical: 1885 1886 >>> children = divide(3, [1, 2, 3, 4, 5, 6, 7]) 1887 >>> [list(c) for c in children] 1888 [[1, 2, 3], [4, 5], [6, 7]] 1889 1890 If the length of the iterable is smaller than n, then the last returned 1891 iterables will be empty: 1892 1893 >>> children = divide(5, [1, 2, 3]) 1894 >>> [list(c) for c in children] 1895 [[1], [2], [3], [], []] 1896 1897 This function will exhaust the iterable before returning and may require 1898 significant storage. If order is not important, see :func:`distribute`, 1899 which does not first pull the iterable into memory. 1900 1901 """ 1902 if n < 1: 1903 raise ValueError('n must be at least 1') 1904 1905 try: 1906 iterable[:0] 1907 except TypeError: 1908 seq = tuple(iterable) 1909 else: 1910 seq = iterable 1911 1912 q, r = divmod(len(seq), n) 1913 1914 ret = [] 1915 stop = 0 1916 for i in range(1, n + 1): 1917 start = stop 1918 stop += q + 1 if i <= r else q 1919 ret.append(iter(seq[start:stop])) 1920 1921 return ret 1922 1923 1924def always_iterable(obj, base_type=(str, bytes)): 1925 """If *obj* is iterable, return an iterator over its items:: 1926 1927 >>> obj = (1, 2, 3) 1928 >>> list(always_iterable(obj)) 1929 [1, 2, 3] 1930 1931 If *obj* is not iterable, return a one-item iterable containing *obj*:: 1932 1933 >>> obj = 1 1934 >>> list(always_iterable(obj)) 1935 [1] 1936 1937 If *obj* is ``None``, return an empty iterable: 1938 1939 >>> obj = None 1940 >>> list(always_iterable(None)) 1941 [] 1942 1943 By default, binary and text strings are not considered iterable:: 1944 1945 >>> obj = 'foo' 1946 >>> list(always_iterable(obj)) 1947 ['foo'] 1948 1949 If *base_type* is set, objects for which ``isinstance(obj, base_type)`` 1950 returns ``True`` won't be considered iterable. 1951 1952 >>> obj = {'a': 1} 1953 >>> list(always_iterable(obj)) # Iterate over the dict's keys 1954 ['a'] 1955 >>> list(always_iterable(obj, base_type=dict)) # Treat dicts as a unit 1956 [{'a': 1}] 1957 1958 Set *base_type* to ``None`` to avoid any special handling and treat objects 1959 Python considers iterable as iterable: 1960 1961 >>> obj = 'foo' 1962 >>> list(always_iterable(obj, base_type=None)) 1963 ['f', 'o', 'o'] 1964 """ 1965 if obj is None: 1966 return iter(()) 1967 1968 if (base_type is not None) and isinstance(obj, base_type): 1969 return iter((obj,)) 1970 1971 try: 1972 return iter(obj) 1973 except TypeError: 1974 return iter((obj,)) 1975 1976 1977def adjacent(predicate, iterable, distance=1): 1978 """Return an iterable over `(bool, item)` tuples where the `item` is 1979 drawn from *iterable* and the `bool` indicates whether 1980 that item satisfies the *predicate* or is adjacent to an item that does. 1981 1982 For example, to find whether items are adjacent to a ``3``:: 1983 1984 >>> list(adjacent(lambda x: x == 3, range(6))) 1985 [(False, 0), (False, 1), (True, 2), (True, 3), (True, 4), (False, 5)] 1986 1987 Set *distance* to change what counts as adjacent. For example, to find 1988 whether items are two places away from a ``3``: 1989 1990 >>> list(adjacent(lambda x: x == 3, range(6), distance=2)) 1991 [(False, 0), (True, 1), (True, 2), (True, 3), (True, 4), (True, 5)] 1992 1993 This is useful for contextualizing the results of a search function. 1994 For example, a code comparison tool might want to identify lines that 1995 have changed, but also surrounding lines to give the viewer of the diff 1996 context. 1997 1998 The predicate function will only be called once for each item in the 1999 iterable. 2000 2001 See also :func:`groupby_transform`, which can be used with this function 2002 to group ranges of items with the same `bool` value. 2003 2004 """ 2005 # Allow distance=0 mainly for testing that it reproduces results with map() 2006 if distance < 0: 2007 raise ValueError('distance must be at least 0') 2008 2009 i1, i2 = tee(iterable) 2010 padding = [False] * distance 2011 selected = chain(padding, map(predicate, i1), padding) 2012 adjacent_to_selected = map(any, windowed(selected, 2 * distance + 1)) 2013 return zip(adjacent_to_selected, i2) 2014 2015 2016def groupby_transform(iterable, keyfunc=None, valuefunc=None, reducefunc=None): 2017 """An extension of :func:`itertools.groupby` that can apply transformations 2018 to the grouped data. 2019 2020 * *keyfunc* is a function computing a key value for each item in *iterable* 2021 * *valuefunc* is a function that transforms the individual items from 2022 *iterable* after grouping 2023 * *reducefunc* is a function that transforms each group of items 2024 2025 >>> iterable = 'aAAbBBcCC' 2026 >>> keyfunc = lambda k: k.upper() 2027 >>> valuefunc = lambda v: v.lower() 2028 >>> reducefunc = lambda g: ''.join(g) 2029 >>> list(groupby_transform(iterable, keyfunc, valuefunc, reducefunc)) 2030 [('A', 'aaa'), ('B', 'bbb'), ('C', 'ccc')] 2031 2032 Each optional argument defaults to an identity function if not specified. 2033 2034 :func:`groupby_transform` is useful when grouping elements of an iterable 2035 using a separate iterable as the key. To do this, :func:`zip` the iterables 2036 and pass a *keyfunc* that extracts the first element and a *valuefunc* 2037 that extracts the second element:: 2038 2039 >>> from operator import itemgetter 2040 >>> keys = [0, 0, 1, 1, 1, 2, 2, 2, 3] 2041 >>> values = 'abcdefghi' 2042 >>> iterable = zip(keys, values) 2043 >>> grouper = groupby_transform(iterable, itemgetter(0), itemgetter(1)) 2044 >>> [(k, ''.join(g)) for k, g in grouper] 2045 [(0, 'ab'), (1, 'cde'), (2, 'fgh'), (3, 'i')] 2046 2047 Note that the order of items in the iterable is significant. 2048 Only adjacent items are grouped together, so if you don't want any 2049 duplicate groups, you should sort the iterable by the key function. 2050 2051 """ 2052 ret = groupby(iterable, keyfunc) 2053 if valuefunc: 2054 ret = ((k, map(valuefunc, g)) for k, g in ret) 2055 if reducefunc: 2056 ret = ((k, reducefunc(g)) for k, g in ret) 2057 2058 return ret 2059 2060 2061class numeric_range(abc.Sequence, abc.Hashable): 2062 """An extension of the built-in ``range()`` function whose arguments can 2063 be any orderable numeric type. 2064 2065 With only *stop* specified, *start* defaults to ``0`` and *step* 2066 defaults to ``1``. The output items will match the type of *stop*: 2067 2068 >>> list(numeric_range(3.5)) 2069 [0.0, 1.0, 2.0, 3.0] 2070 2071 With only *start* and *stop* specified, *step* defaults to ``1``. The 2072 output items will match the type of *start*: 2073 2074 >>> from decimal import Decimal 2075 >>> start = Decimal('2.1') 2076 >>> stop = Decimal('5.1') 2077 >>> list(numeric_range(start, stop)) 2078 [Decimal('2.1'), Decimal('3.1'), Decimal('4.1')] 2079 2080 With *start*, *stop*, and *step* specified the output items will match 2081 the type of ``start + step``: 2082 2083 >>> from fractions import Fraction 2084 >>> start = Fraction(1, 2) # Start at 1/2 2085 >>> stop = Fraction(5, 2) # End at 5/2 2086 >>> step = Fraction(1, 2) # Count by 1/2 2087 >>> list(numeric_range(start, stop, step)) 2088 [Fraction(1, 2), Fraction(1, 1), Fraction(3, 2), Fraction(2, 1)] 2089 2090 If *step* is zero, ``ValueError`` is raised. Negative steps are supported: 2091 2092 >>> list(numeric_range(3, -1, -1.0)) 2093 [3.0, 2.0, 1.0, 0.0] 2094 2095 Be aware of the limitations of floating point numbers; the representation 2096 of the yielded numbers may be surprising. 2097 2098 ``datetime.datetime`` objects can be used for *start* and *stop*, if *step* 2099 is a ``datetime.timedelta`` object: 2100 2101 >>> import datetime 2102 >>> start = datetime.datetime(2019, 1, 1) 2103 >>> stop = datetime.datetime(2019, 1, 3) 2104 >>> step = datetime.timedelta(days=1) 2105 >>> items = iter(numeric_range(start, stop, step)) 2106 >>> next(items) 2107 datetime.datetime(2019, 1, 1, 0, 0) 2108 >>> next(items) 2109 datetime.datetime(2019, 1, 2, 0, 0) 2110 2111 """ 2112 2113 _EMPTY_HASH = hash(range(0, 0)) 2114 2115 def __init__(self, *args): 2116 argc = len(args) 2117 if argc == 1: 2118 (self._stop,) = args 2119 self._start = type(self._stop)(0) 2120 self._step = type(self._stop - self._start)(1) 2121 elif argc == 2: 2122 self._start, self._stop = args 2123 self._step = type(self._stop - self._start)(1) 2124 elif argc == 3: 2125 self._start, self._stop, self._step = args 2126 elif argc == 0: 2127 raise TypeError( 2128 'numeric_range expected at least ' 2129 '1 argument, got {}'.format(argc) 2130 ) 2131 else: 2132 raise TypeError( 2133 'numeric_range expected at most ' 2134 '3 arguments, got {}'.format(argc) 2135 ) 2136 2137 self._zero = type(self._step)(0) 2138 if self._step == self._zero: 2139 raise ValueError('numeric_range() arg 3 must not be zero') 2140 self._growing = self._step > self._zero 2141 self._init_len() 2142 2143 def __bool__(self): 2144 if self._growing: 2145 return self._start < self._stop 2146 else: 2147 return self._start > self._stop 2148 2149 def __contains__(self, elem): 2150 if self._growing: 2151 if self._start <= elem < self._stop: 2152 return (elem - self._start) % self._step == self._zero 2153 else: 2154 if self._start >= elem > self._stop: 2155 return (self._start - elem) % (-self._step) == self._zero 2156 2157 return False 2158 2159 def __eq__(self, other): 2160 if isinstance(other, numeric_range): 2161 empty_self = not bool(self) 2162 empty_other = not bool(other) 2163 if empty_self or empty_other: 2164 return empty_self and empty_other # True if both empty 2165 else: 2166 return ( 2167 self._start == other._start 2168 and self._step == other._step 2169 and self._get_by_index(-1) == other._get_by_index(-1) 2170 ) 2171 else: 2172 return False 2173 2174 def __getitem__(self, key): 2175 if isinstance(key, int): 2176 return self._get_by_index(key) 2177 elif isinstance(key, slice): 2178 step = self._step if key.step is None else key.step * self._step 2179 2180 if key.start is None or key.start <= -self._len: 2181 start = self._start 2182 elif key.start >= self._len: 2183 start = self._stop 2184 else: # -self._len < key.start < self._len 2185 start = self._get_by_index(key.start) 2186 2187 if key.stop is None or key.stop >= self._len: 2188 stop = self._stop 2189 elif key.stop <= -self._len: 2190 stop = self._start 2191 else: # -self._len < key.stop < self._len 2192 stop = self._get_by_index(key.stop) 2193 2194 return numeric_range(start, stop, step) 2195 else: 2196 raise TypeError( 2197 'numeric range indices must be ' 2198 'integers or slices, not {}'.format(type(key).__name__) 2199 ) 2200 2201 def __hash__(self): 2202 if self: 2203 return hash((self._start, self._get_by_index(-1), self._step)) 2204 else: 2205 return self._EMPTY_HASH 2206 2207 def __iter__(self): 2208 values = (self._start + (n * self._step) for n in count()) 2209 if self._growing: 2210 return takewhile(partial(gt, self._stop), values) 2211 else: 2212 return takewhile(partial(lt, self._stop), values) 2213 2214 def __len__(self): 2215 return self._len 2216 2217 def _init_len(self): 2218 if self._growing: 2219 start = self._start 2220 stop = self._stop 2221 step = self._step 2222 else: 2223 start = self._stop 2224 stop = self._start 2225 step = -self._step 2226 distance = stop - start 2227 if distance <= self._zero: 2228 self._len = 0 2229 else: # distance > 0 and step > 0: regular euclidean division 2230 q, r = divmod(distance, step) 2231 self._len = int(q) + int(r != self._zero) 2232 2233 def __reduce__(self): 2234 return numeric_range, (self._start, self._stop, self._step) 2235 2236 def __repr__(self): 2237 if self._step == 1: 2238 return "numeric_range({}, {})".format( 2239 repr(self._start), repr(self._stop) 2240 ) 2241 else: 2242 return "numeric_range({}, {}, {})".format( 2243 repr(self._start), repr(self._stop), repr(self._step) 2244 ) 2245 2246 def __reversed__(self): 2247 return iter( 2248 numeric_range( 2249 self._get_by_index(-1), self._start - self._step, -self._step 2250 ) 2251 ) 2252 2253 def count(self, value): 2254 return int(value in self) 2255 2256 def index(self, value): 2257 if self._growing: 2258 if self._start <= value < self._stop: 2259 q, r = divmod(value - self._start, self._step) 2260 if r == self._zero: 2261 return int(q) 2262 else: 2263 if self._start >= value > self._stop: 2264 q, r = divmod(self._start - value, -self._step) 2265 if r == self._zero: 2266 return int(q) 2267 2268 raise ValueError("{} is not in numeric range".format(value)) 2269 2270 def _get_by_index(self, i): 2271 if i < 0: 2272 i += self._len 2273 if i < 0 or i >= self._len: 2274 raise IndexError("numeric range object index out of range") 2275 return self._start + i * self._step 2276 2277 2278def count_cycle(iterable, n=None): 2279 """Cycle through the items from *iterable* up to *n* times, yielding 2280 the number of completed cycles along with each item. If *n* is omitted the 2281 process repeats indefinitely. 2282 2283 >>> list(count_cycle('AB', 3)) 2284 [(0, 'A'), (0, 'B'), (1, 'A'), (1, 'B'), (2, 'A'), (2, 'B')] 2285 2286 """ 2287 iterable = tuple(iterable) 2288 if not iterable: 2289 return iter(()) 2290 counter = count() if n is None else range(n) 2291 return ((i, item) for i in counter for item in iterable) 2292 2293 2294def mark_ends(iterable): 2295 """Yield 3-tuples of the form ``(is_first, is_last, item)``. 2296 2297 >>> list(mark_ends('ABC')) 2298 [(True, False, 'A'), (False, False, 'B'), (False, True, 'C')] 2299 2300 Use this when looping over an iterable to take special action on its first 2301 and/or last items: 2302 2303 >>> iterable = ['Header', 100, 200, 'Footer'] 2304 >>> total = 0 2305 >>> for is_first, is_last, item in mark_ends(iterable): 2306 ... if is_first: 2307 ... continue # Skip the header 2308 ... if is_last: 2309 ... continue # Skip the footer 2310 ... total += item 2311 >>> print(total) 2312 300 2313 """ 2314 it = iter(iterable) 2315 2316 try: 2317 b = next(it) 2318 except StopIteration: 2319 return 2320 2321 try: 2322 for i in count(): 2323 a = b 2324 b = next(it) 2325 yield i == 0, False, a 2326 2327 except StopIteration: 2328 yield i == 0, True, a 2329 2330 2331def locate(iterable, pred=bool, window_size=None): 2332 """Yield the index of each item in *iterable* for which *pred* returns 2333 ``True``. 2334 2335 *pred* defaults to :func:`bool`, which will select truthy items: 2336 2337 >>> list(locate([0, 1, 1, 0, 1, 0, 0])) 2338 [1, 2, 4] 2339 2340 Set *pred* to a custom function to, e.g., find the indexes for a particular 2341 item. 2342 2343 >>> list(locate(['a', 'b', 'c', 'b'], lambda x: x == 'b')) 2344 [1, 3] 2345 2346 If *window_size* is given, then the *pred* function will be called with 2347 that many items. This enables searching for sub-sequences: 2348 2349 >>> iterable = [0, 1, 2, 3, 0, 1, 2, 3, 0, 1, 2, 3] 2350 >>> pred = lambda *args: args == (1, 2, 3) 2351 >>> list(locate(iterable, pred=pred, window_size=3)) 2352 [1, 5, 9] 2353 2354 Use with :func:`seekable` to find indexes and then retrieve the associated 2355 items: 2356 2357 >>> from itertools import count 2358 >>> from more_itertools import seekable 2359 >>> source = (3 * n + 1 if (n % 2) else n // 2 for n in count()) 2360 >>> it = seekable(source) 2361 >>> pred = lambda x: x > 100 2362 >>> indexes = locate(it, pred=pred) 2363 >>> i = next(indexes) 2364 >>> it.seek(i) 2365 >>> next(it) 2366 106 2367 2368 """ 2369 if window_size is None: 2370 return compress(count(), map(pred, iterable)) 2371 2372 if window_size < 1: 2373 raise ValueError('window size must be at least 1') 2374 2375 it = windowed(iterable, window_size, fillvalue=_marker) 2376 return compress(count(), starmap(pred, it)) 2377 2378 2379def lstrip(iterable, pred): 2380 """Yield the items from *iterable*, but strip any from the beginning 2381 for which *pred* returns ``True``. 2382 2383 For example, to remove a set of items from the start of an iterable: 2384 2385 >>> iterable = (None, False, None, 1, 2, None, 3, False, None) 2386 >>> pred = lambda x: x in {None, False, ''} 2387 >>> list(lstrip(iterable, pred)) 2388 [1, 2, None, 3, False, None] 2389 2390 This function is analogous to to :func:`str.lstrip`, and is essentially 2391 an wrapper for :func:`itertools.dropwhile`. 2392 2393 """ 2394 return dropwhile(pred, iterable) 2395 2396 2397def rstrip(iterable, pred): 2398 """Yield the items from *iterable*, but strip any from the end 2399 for which *pred* returns ``True``. 2400 2401 For example, to remove a set of items from the end of an iterable: 2402 2403 >>> iterable = (None, False, None, 1, 2, None, 3, False, None) 2404 >>> pred = lambda x: x in {None, False, ''} 2405 >>> list(rstrip(iterable, pred)) 2406 [None, False, None, 1, 2, None, 3] 2407 2408 This function is analogous to :func:`str.rstrip`. 2409 2410 """ 2411 cache = [] 2412 cache_append = cache.append 2413 cache_clear = cache.clear 2414 for x in iterable: 2415 if pred(x): 2416 cache_append(x) 2417 else: 2418 yield from cache 2419 cache_clear() 2420 yield x 2421 2422 2423def strip(iterable, pred): 2424 """Yield the items from *iterable*, but strip any from the 2425 beginning and end for which *pred* returns ``True``. 2426 2427 For example, to remove a set of items from both ends of an iterable: 2428 2429 >>> iterable = (None, False, None, 1, 2, None, 3, False, None) 2430 >>> pred = lambda x: x in {None, False, ''} 2431 >>> list(strip(iterable, pred)) 2432 [1, 2, None, 3] 2433 2434 This function is analogous to :func:`str.strip`. 2435 2436 """ 2437 return rstrip(lstrip(iterable, pred), pred) 2438 2439 2440class islice_extended: 2441 """An extension of :func:`itertools.islice` that supports negative values 2442 for *stop*, *start*, and *step*. 2443 2444 >>> iterable = iter('abcdefgh') 2445 >>> list(islice_extended(iterable, -4, -1)) 2446 ['e', 'f', 'g'] 2447 2448 Slices with negative values require some caching of *iterable*, but this 2449 function takes care to minimize the amount of memory required. 2450 2451 For example, you can use a negative step with an infinite iterator: 2452 2453 >>> from itertools import count 2454 >>> list(islice_extended(count(), 110, 99, -2)) 2455 [110, 108, 106, 104, 102, 100] 2456 2457 You can also use slice notation directly: 2458 2459 >>> iterable = map(str, count()) 2460 >>> it = islice_extended(iterable)[10:20:2] 2461 >>> list(it) 2462 ['10', '12', '14', '16', '18'] 2463 2464 """ 2465 2466 def __init__(self, iterable, *args): 2467 it = iter(iterable) 2468 if args: 2469 self._iterable = _islice_helper(it, slice(*args)) 2470 else: 2471 self._iterable = it 2472 2473 def __iter__(self): 2474 return self 2475 2476 def __next__(self): 2477 return next(self._iterable) 2478 2479 def __getitem__(self, key): 2480 if isinstance(key, slice): 2481 return islice_extended(_islice_helper(self._iterable, key)) 2482 2483 raise TypeError('islice_extended.__getitem__ argument must be a slice') 2484 2485 2486def _islice_helper(it, s): 2487 start = s.start 2488 stop = s.stop 2489 if s.step == 0: 2490 raise ValueError('step argument must be a non-zero integer or None.') 2491 step = s.step or 1 2492 2493 if step > 0: 2494 start = 0 if (start is None) else start 2495 2496 if start < 0: 2497 # Consume all but the last -start items 2498 cache = deque(enumerate(it, 1), maxlen=-start) 2499 len_iter = cache[-1][0] if cache else 0 2500 2501 # Adjust start to be positive 2502 i = max(len_iter + start, 0) 2503 2504 # Adjust stop to be positive 2505 if stop is None: 2506 j = len_iter 2507 elif stop >= 0: 2508 j = min(stop, len_iter) 2509 else: 2510 j = max(len_iter + stop, 0) 2511 2512 # Slice the cache 2513 n = j - i 2514 if n <= 0: 2515 return 2516 2517 for index, item in islice(cache, 0, n, step): 2518 yield item 2519 elif (stop is not None) and (stop < 0): 2520 # Advance to the start position 2521 next(islice(it, start, start), None) 2522 2523 # When stop is negative, we have to carry -stop items while 2524 # iterating 2525 cache = deque(islice(it, -stop), maxlen=-stop) 2526 2527 for index, item in enumerate(it): 2528 cached_item = cache.popleft() 2529 if index % step == 0: 2530 yield cached_item 2531 cache.append(item) 2532 else: 2533 # When both start and stop are positive we have the normal case 2534 yield from islice(it, start, stop, step) 2535 else: 2536 start = -1 if (start is None) else start 2537 2538 if (stop is not None) and (stop < 0): 2539 # Consume all but the last items 2540 n = -stop - 1 2541 cache = deque(enumerate(it, 1), maxlen=n) 2542 len_iter = cache[-1][0] if cache else 0 2543 2544 # If start and stop are both negative they are comparable and 2545 # we can just slice. Otherwise we can adjust start to be negative 2546 # and then slice. 2547 if start < 0: 2548 i, j = start, stop 2549 else: 2550 i, j = min(start - len_iter, -1), None 2551 2552 for index, item in list(cache)[i:j:step]: 2553 yield item 2554 else: 2555 # Advance to the stop position 2556 if stop is not None: 2557 m = stop + 1 2558 next(islice(it, m, m), None) 2559 2560 # stop is positive, so if start is negative they are not comparable 2561 # and we need the rest of the items. 2562 if start < 0: 2563 i = start 2564 n = None 2565 # stop is None and start is positive, so we just need items up to 2566 # the start index. 2567 elif stop is None: 2568 i = None 2569 n = start + 1 2570 # Both stop and start are positive, so they are comparable. 2571 else: 2572 i = None 2573 n = start - stop 2574 if n <= 0: 2575 return 2576 2577 cache = list(islice(it, n)) 2578 2579 yield from cache[i::step] 2580 2581 2582def always_reversible(iterable): 2583 """An extension of :func:`reversed` that supports all iterables, not 2584 just those which implement the ``Reversible`` or ``Sequence`` protocols. 2585 2586 >>> print(*always_reversible(x for x in range(3))) 2587 2 1 0 2588 2589 If the iterable is already reversible, this function returns the 2590 result of :func:`reversed()`. If the iterable is not reversible, 2591 this function will cache the remaining items in the iterable and 2592 yield them in reverse order, which may require significant storage. 2593 """ 2594 try: 2595 return reversed(iterable) 2596 except TypeError: 2597 return reversed(list(iterable)) 2598 2599 2600def consecutive_groups(iterable, ordering=lambda x: x): 2601 """Yield groups of consecutive items using :func:`itertools.groupby`. 2602 The *ordering* function determines whether two items are adjacent by 2603 returning their position. 2604 2605 By default, the ordering function is the identity function. This is 2606 suitable for finding runs of numbers: 2607 2608 >>> iterable = [1, 10, 11, 12, 20, 30, 31, 32, 33, 40] 2609 >>> for group in consecutive_groups(iterable): 2610 ... print(list(group)) 2611 [1] 2612 [10, 11, 12] 2613 [20] 2614 [30, 31, 32, 33] 2615 [40] 2616 2617 For finding runs of adjacent letters, try using the :meth:`index` method 2618 of a string of letters: 2619 2620 >>> from string import ascii_lowercase 2621 >>> iterable = 'abcdfgilmnop' 2622 >>> ordering = ascii_lowercase.index 2623 >>> for group in consecutive_groups(iterable, ordering): 2624 ... print(list(group)) 2625 ['a', 'b', 'c', 'd'] 2626 ['f', 'g'] 2627 ['i'] 2628 ['l', 'm', 'n', 'o', 'p'] 2629 2630 Each group of consecutive items is an iterator that shares it source with 2631 *iterable*. When an an output group is advanced, the previous group is 2632 no longer available unless its elements are copied (e.g., into a ``list``). 2633 2634 >>> iterable = [1, 2, 11, 12, 21, 22] 2635 >>> saved_groups = [] 2636 >>> for group in consecutive_groups(iterable): 2637 ... saved_groups.append(list(group)) # Copy group elements 2638 >>> saved_groups 2639 [[1, 2], [11, 12], [21, 22]] 2640 2641 """ 2642 for k, g in groupby( 2643 enumerate(iterable), key=lambda x: x[0] - ordering(x[1]) 2644 ): 2645 yield map(itemgetter(1), g) 2646 2647 2648def difference(iterable, func=sub, *, initial=None): 2649 """This function is the inverse of :func:`itertools.accumulate`. By default 2650 it will compute the first difference of *iterable* using 2651 :func:`operator.sub`: 2652 2653 >>> from itertools import accumulate 2654 >>> iterable = accumulate([0, 1, 2, 3, 4]) # produces 0, 1, 3, 6, 10 2655 >>> list(difference(iterable)) 2656 [0, 1, 2, 3, 4] 2657 2658 *func* defaults to :func:`operator.sub`, but other functions can be 2659 specified. They will be applied as follows:: 2660 2661 A, B, C, D, ... --> A, func(B, A), func(C, B), func(D, C), ... 2662 2663 For example, to do progressive division: 2664 2665 >>> iterable = [1, 2, 6, 24, 120] 2666 >>> func = lambda x, y: x // y 2667 >>> list(difference(iterable, func)) 2668 [1, 2, 3, 4, 5] 2669 2670 If the *initial* keyword is set, the first element will be skipped when 2671 computing successive differences. 2672 2673 >>> it = [10, 11, 13, 16] # from accumulate([1, 2, 3], initial=10) 2674 >>> list(difference(it, initial=10)) 2675 [1, 2, 3] 2676 2677 """ 2678 a, b = tee(iterable) 2679 try: 2680 first = [next(b)] 2681 except StopIteration: 2682 return iter([]) 2683 2684 if initial is not None: 2685 first = [] 2686 2687 return chain(first, starmap(func, zip(b, a))) 2688 2689 2690class SequenceView(Sequence): 2691 """Return a read-only view of the sequence object *target*. 2692 2693 :class:`SequenceView` objects are analogous to Python's built-in 2694 "dictionary view" types. They provide a dynamic view of a sequence's items, 2695 meaning that when the sequence updates, so does the view. 2696 2697 >>> seq = ['0', '1', '2'] 2698 >>> view = SequenceView(seq) 2699 >>> view 2700 SequenceView(['0', '1', '2']) 2701 >>> seq.append('3') 2702 >>> view 2703 SequenceView(['0', '1', '2', '3']) 2704 2705 Sequence views support indexing, slicing, and length queries. They act 2706 like the underlying sequence, except they don't allow assignment: 2707 2708 >>> view[1] 2709 '1' 2710 >>> view[1:-1] 2711 ['1', '2'] 2712 >>> len(view) 2713 4 2714 2715 Sequence views are useful as an alternative to copying, as they don't 2716 require (much) extra storage. 2717 2718 """ 2719 2720 def __init__(self, target): 2721 if not isinstance(target, Sequence): 2722 raise TypeError 2723 self._target = target 2724 2725 def __getitem__(self, index): 2726 return self._target[index] 2727 2728 def __len__(self): 2729 return len(self._target) 2730 2731 def __repr__(self): 2732 return '{}({})'.format(self.__class__.__name__, repr(self._target)) 2733 2734 2735class seekable: 2736 """Wrap an iterator to allow for seeking backward and forward. This 2737 progressively caches the items in the source iterable so they can be 2738 re-visited. 2739 2740 Call :meth:`seek` with an index to seek to that position in the source 2741 iterable. 2742 2743 To "reset" an iterator, seek to ``0``: 2744 2745 >>> from itertools import count 2746 >>> it = seekable((str(n) for n in count())) 2747 >>> next(it), next(it), next(it) 2748 ('0', '1', '2') 2749 >>> it.seek(0) 2750 >>> next(it), next(it), next(it) 2751 ('0', '1', '2') 2752 >>> next(it) 2753 '3' 2754 2755 You can also seek forward: 2756 2757 >>> it = seekable((str(n) for n in range(20))) 2758 >>> it.seek(10) 2759 >>> next(it) 2760 '10' 2761 >>> it.seek(20) # Seeking past the end of the source isn't a problem 2762 >>> list(it) 2763 [] 2764 >>> it.seek(0) # Resetting works even after hitting the end 2765 >>> next(it), next(it), next(it) 2766 ('0', '1', '2') 2767 2768 Call :meth:`peek` to look ahead one item without advancing the iterator: 2769 2770 >>> it = seekable('1234') 2771 >>> it.peek() 2772 '1' 2773 >>> list(it) 2774 ['1', '2', '3', '4'] 2775 >>> it.peek(default='empty') 2776 'empty' 2777 2778 Before the iterator is at its end, calling :func:`bool` on it will return 2779 ``True``. After it will return ``False``: 2780 2781 >>> it = seekable('5678') 2782 >>> bool(it) 2783 True 2784 >>> list(it) 2785 ['5', '6', '7', '8'] 2786 >>> bool(it) 2787 False 2788 2789 You may view the contents of the cache with the :meth:`elements` method. 2790 That returns a :class:`SequenceView`, a view that updates automatically: 2791 2792 >>> it = seekable((str(n) for n in range(10))) 2793 >>> next(it), next(it), next(it) 2794 ('0', '1', '2') 2795 >>> elements = it.elements() 2796 >>> elements 2797 SequenceView(['0', '1', '2']) 2798 >>> next(it) 2799 '3' 2800 >>> elements 2801 SequenceView(['0', '1', '2', '3']) 2802 2803 By default, the cache grows as the source iterable progresses, so beware of 2804 wrapping very large or infinite iterables. Supply *maxlen* to limit the 2805 size of the cache (this of course limits how far back you can seek). 2806 2807 >>> from itertools import count 2808 >>> it = seekable((str(n) for n in count()), maxlen=2) 2809 >>> next(it), next(it), next(it), next(it) 2810 ('0', '1', '2', '3') 2811 >>> list(it.elements()) 2812 ['2', '3'] 2813 >>> it.seek(0) 2814 >>> next(it), next(it), next(it), next(it) 2815 ('2', '3', '4', '5') 2816 >>> next(it) 2817 '6' 2818 2819 """ 2820 2821 def __init__(self, iterable, maxlen=None): 2822 self._source = iter(iterable) 2823 if maxlen is None: 2824 self._cache = [] 2825 else: 2826 self._cache = deque([], maxlen) 2827 self._index = None 2828 2829 def __iter__(self): 2830 return self 2831 2832 def __next__(self): 2833 if self._index is not None: 2834 try: 2835 item = self._cache[self._index] 2836 except IndexError: 2837 self._index = None 2838 else: 2839 self._index += 1 2840 return item 2841 2842 item = next(self._source) 2843 self._cache.append(item) 2844 return item 2845 2846 def __bool__(self): 2847 try: 2848 self.peek() 2849 except StopIteration: 2850 return False 2851 return True 2852 2853 def peek(self, default=_marker): 2854 try: 2855 peeked = next(self) 2856 except StopIteration: 2857 if default is _marker: 2858 raise 2859 return default 2860 if self._index is None: 2861 self._index = len(self._cache) 2862 self._index -= 1 2863 return peeked 2864 2865 def elements(self): 2866 return SequenceView(self._cache) 2867 2868 def seek(self, index): 2869 self._index = index 2870 remainder = index - len(self._cache) 2871 if remainder > 0: 2872 consume(self, remainder) 2873 2874 2875class run_length: 2876 """ 2877 :func:`run_length.encode` compresses an iterable with run-length encoding. 2878 It yields groups of repeated items with the count of how many times they 2879 were repeated: 2880 2881 >>> uncompressed = 'abbcccdddd' 2882 >>> list(run_length.encode(uncompressed)) 2883 [('a', 1), ('b', 2), ('c', 3), ('d', 4)] 2884 2885 :func:`run_length.decode` decompresses an iterable that was previously 2886 compressed with run-length encoding. It yields the items of the 2887 decompressed iterable: 2888 2889 >>> compressed = [('a', 1), ('b', 2), ('c', 3), ('d', 4)] 2890 >>> list(run_length.decode(compressed)) 2891 ['a', 'b', 'b', 'c', 'c', 'c', 'd', 'd', 'd', 'd'] 2892 2893 """ 2894 2895 @staticmethod 2896 def encode(iterable): 2897 return ((k, ilen(g)) for k, g in groupby(iterable)) 2898 2899 @staticmethod 2900 def decode(iterable): 2901 return chain.from_iterable(repeat(k, n) for k, n in iterable) 2902 2903 2904def exactly_n(iterable, n, predicate=bool): 2905 """Return ``True`` if exactly ``n`` items in the iterable are ``True`` 2906 according to the *predicate* function. 2907 2908 >>> exactly_n([True, True, False], 2) 2909 True 2910 >>> exactly_n([True, True, False], 1) 2911 False 2912 >>> exactly_n([0, 1, 2, 3, 4, 5], 3, lambda x: x < 3) 2913 True 2914 2915 The iterable will be advanced until ``n + 1`` truthy items are encountered, 2916 so avoid calling it on infinite iterables. 2917 2918 """ 2919 return len(take(n + 1, filter(predicate, iterable))) == n 2920 2921 2922def circular_shifts(iterable): 2923 """Return a list of circular shifts of *iterable*. 2924 2925 >>> circular_shifts(range(4)) 2926 [(0, 1, 2, 3), (1, 2, 3, 0), (2, 3, 0, 1), (3, 0, 1, 2)] 2927 """ 2928 lst = list(iterable) 2929 return take(len(lst), windowed(cycle(lst), len(lst))) 2930 2931 2932def make_decorator(wrapping_func, result_index=0): 2933 """Return a decorator version of *wrapping_func*, which is a function that 2934 modifies an iterable. *result_index* is the position in that function's 2935 signature where the iterable goes. 2936 2937 This lets you use itertools on the "production end," i.e. at function 2938 definition. This can augment what the function returns without changing the 2939 function's code. 2940 2941 For example, to produce a decorator version of :func:`chunked`: 2942 2943 >>> from more_itertools import chunked 2944 >>> chunker = make_decorator(chunked, result_index=0) 2945 >>> @chunker(3) 2946 ... def iter_range(n): 2947 ... return iter(range(n)) 2948 ... 2949 >>> list(iter_range(9)) 2950 [[0, 1, 2], [3, 4, 5], [6, 7, 8]] 2951 2952 To only allow truthy items to be returned: 2953 2954 >>> truth_serum = make_decorator(filter, result_index=1) 2955 >>> @truth_serum(bool) 2956 ... def boolean_test(): 2957 ... return [0, 1, '', ' ', False, True] 2958 ... 2959 >>> list(boolean_test()) 2960 [1, ' ', True] 2961 2962 The :func:`peekable` and :func:`seekable` wrappers make for practical 2963 decorators: 2964 2965 >>> from more_itertools import peekable 2966 >>> peekable_function = make_decorator(peekable) 2967 >>> @peekable_function() 2968 ... def str_range(*args): 2969 ... return (str(x) for x in range(*args)) 2970 ... 2971 >>> it = str_range(1, 20, 2) 2972 >>> next(it), next(it), next(it) 2973 ('1', '3', '5') 2974 >>> it.peek() 2975 '7' 2976 >>> next(it) 2977 '7' 2978 2979 """ 2980 # See https://sites.google.com/site/bbayles/index/decorator_factory for 2981 # notes on how this works. 2982 def decorator(*wrapping_args, **wrapping_kwargs): 2983 def outer_wrapper(f): 2984 def inner_wrapper(*args, **kwargs): 2985 result = f(*args, **kwargs) 2986 wrapping_args_ = list(wrapping_args) 2987 wrapping_args_.insert(result_index, result) 2988 return wrapping_func(*wrapping_args_, **wrapping_kwargs) 2989 2990 return inner_wrapper 2991 2992 return outer_wrapper 2993 2994 return decorator 2995 2996 2997def map_reduce(iterable, keyfunc, valuefunc=None, reducefunc=None): 2998 """Return a dictionary that maps the items in *iterable* to categories 2999 defined by *keyfunc*, transforms them with *valuefunc*, and 3000 then summarizes them by category with *reducefunc*. 3001 3002 *valuefunc* defaults to the identity function if it is unspecified. 3003 If *reducefunc* is unspecified, no summarization takes place: 3004 3005 >>> keyfunc = lambda x: x.upper() 3006 >>> result = map_reduce('abbccc', keyfunc) 3007 >>> sorted(result.items()) 3008 [('A', ['a']), ('B', ['b', 'b']), ('C', ['c', 'c', 'c'])] 3009 3010 Specifying *valuefunc* transforms the categorized items: 3011 3012 >>> keyfunc = lambda x: x.upper() 3013 >>> valuefunc = lambda x: 1 3014 >>> result = map_reduce('abbccc', keyfunc, valuefunc) 3015 >>> sorted(result.items()) 3016 [('A', [1]), ('B', [1, 1]), ('C', [1, 1, 1])] 3017 3018 Specifying *reducefunc* summarizes the categorized items: 3019 3020 >>> keyfunc = lambda x: x.upper() 3021 >>> valuefunc = lambda x: 1 3022 >>> reducefunc = sum 3023 >>> result = map_reduce('abbccc', keyfunc, valuefunc, reducefunc) 3024 >>> sorted(result.items()) 3025 [('A', 1), ('B', 2), ('C', 3)] 3026 3027 You may want to filter the input iterable before applying the map/reduce 3028 procedure: 3029 3030 >>> all_items = range(30) 3031 >>> items = [x for x in all_items if 10 <= x <= 20] # Filter 3032 >>> keyfunc = lambda x: x % 2 # Evens map to 0; odds to 1 3033 >>> categories = map_reduce(items, keyfunc=keyfunc) 3034 >>> sorted(categories.items()) 3035 [(0, [10, 12, 14, 16, 18, 20]), (1, [11, 13, 15, 17, 19])] 3036 >>> summaries = map_reduce(items, keyfunc=keyfunc, reducefunc=sum) 3037 >>> sorted(summaries.items()) 3038 [(0, 90), (1, 75)] 3039 3040 Note that all items in the iterable are gathered into a list before the 3041 summarization step, which may require significant storage. 3042 3043 The returned object is a :obj:`collections.defaultdict` with the 3044 ``default_factory`` set to ``None``, such that it behaves like a normal 3045 dictionary. 3046 3047 """ 3048 valuefunc = (lambda x: x) if (valuefunc is None) else valuefunc 3049 3050 ret = defaultdict(list) 3051 for item in iterable: 3052 key = keyfunc(item) 3053 value = valuefunc(item) 3054 ret[key].append(value) 3055 3056 if reducefunc is not None: 3057 for key, value_list in ret.items(): 3058 ret[key] = reducefunc(value_list) 3059 3060 ret.default_factory = None 3061 return ret 3062 3063 3064def rlocate(iterable, pred=bool, window_size=None): 3065 """Yield the index of each item in *iterable* for which *pred* returns 3066 ``True``, starting from the right and moving left. 3067 3068 *pred* defaults to :func:`bool`, which will select truthy items: 3069 3070 >>> list(rlocate([0, 1, 1, 0, 1, 0, 0])) # Truthy at 1, 2, and 4 3071 [4, 2, 1] 3072 3073 Set *pred* to a custom function to, e.g., find the indexes for a particular 3074 item: 3075 3076 >>> iterable = iter('abcb') 3077 >>> pred = lambda x: x == 'b' 3078 >>> list(rlocate(iterable, pred)) 3079 [3, 1] 3080 3081 If *window_size* is given, then the *pred* function will be called with 3082 that many items. This enables searching for sub-sequences: 3083 3084 >>> iterable = [0, 1, 2, 3, 0, 1, 2, 3, 0, 1, 2, 3] 3085 >>> pred = lambda *args: args == (1, 2, 3) 3086 >>> list(rlocate(iterable, pred=pred, window_size=3)) 3087 [9, 5, 1] 3088 3089 Beware, this function won't return anything for infinite iterables. 3090 If *iterable* is reversible, ``rlocate`` will reverse it and search from 3091 the right. Otherwise, it will search from the left and return the results 3092 in reverse order. 3093 3094 See :func:`locate` to for other example applications. 3095 3096 """ 3097 if window_size is None: 3098 try: 3099 len_iter = len(iterable) 3100 return (len_iter - i - 1 for i in locate(reversed(iterable), pred)) 3101 except TypeError: 3102 pass 3103 3104 return reversed(list(locate(iterable, pred, window_size))) 3105 3106 3107def replace(iterable, pred, substitutes, count=None, window_size=1): 3108 """Yield the items from *iterable*, replacing the items for which *pred* 3109 returns ``True`` with the items from the iterable *substitutes*. 3110 3111 >>> iterable = [1, 1, 0, 1, 1, 0, 1, 1] 3112 >>> pred = lambda x: x == 0 3113 >>> substitutes = (2, 3) 3114 >>> list(replace(iterable, pred, substitutes)) 3115 [1, 1, 2, 3, 1, 1, 2, 3, 1, 1] 3116 3117 If *count* is given, the number of replacements will be limited: 3118 3119 >>> iterable = [1, 1, 0, 1, 1, 0, 1, 1, 0] 3120 >>> pred = lambda x: x == 0 3121 >>> substitutes = [None] 3122 >>> list(replace(iterable, pred, substitutes, count=2)) 3123 [1, 1, None, 1, 1, None, 1, 1, 0] 3124 3125 Use *window_size* to control the number of items passed as arguments to 3126 *pred*. This allows for locating and replacing subsequences. 3127 3128 >>> iterable = [0, 1, 2, 5, 0, 1, 2, 5] 3129 >>> window_size = 3 3130 >>> pred = lambda *args: args == (0, 1, 2) # 3 items passed to pred 3131 >>> substitutes = [3, 4] # Splice in these items 3132 >>> list(replace(iterable, pred, substitutes, window_size=window_size)) 3133 [3, 4, 5, 3, 4, 5] 3134 3135 """ 3136 if window_size < 1: 3137 raise ValueError('window_size must be at least 1') 3138 3139 # Save the substitutes iterable, since it's used more than once 3140 substitutes = tuple(substitutes) 3141 3142 # Add padding such that the number of windows matches the length of the 3143 # iterable 3144 it = chain(iterable, [_marker] * (window_size - 1)) 3145 windows = windowed(it, window_size) 3146 3147 n = 0 3148 for w in windows: 3149 # If the current window matches our predicate (and we haven't hit 3150 # our maximum number of replacements), splice in the substitutes 3151 # and then consume the following windows that overlap with this one. 3152 # For example, if the iterable is (0, 1, 2, 3, 4...) 3153 # and the window size is 2, we have (0, 1), (1, 2), (2, 3)... 3154 # If the predicate matches on (0, 1), we need to zap (0, 1) and (1, 2) 3155 if pred(*w): 3156 if (count is None) or (n < count): 3157 n += 1 3158 yield from substitutes 3159 consume(windows, window_size - 1) 3160 continue 3161 3162 # If there was no match (or we've reached the replacement limit), 3163 # yield the first item from the window. 3164 if w and (w[0] is not _marker): 3165 yield w[0] 3166 3167 3168def partitions(iterable): 3169 """Yield all possible order-preserving partitions of *iterable*. 3170 3171 >>> iterable = 'abc' 3172 >>> for part in partitions(iterable): 3173 ... print([''.join(p) for p in part]) 3174 ['abc'] 3175 ['a', 'bc'] 3176 ['ab', 'c'] 3177 ['a', 'b', 'c'] 3178 3179 This is unrelated to :func:`partition`. 3180 3181 """ 3182 sequence = list(iterable) 3183 n = len(sequence) 3184 for i in powerset(range(1, n)): 3185 yield [sequence[i:j] for i, j in zip((0,) + i, i + (n,))] 3186 3187 3188def set_partitions(iterable, k=None): 3189 """ 3190 Yield the set partitions of *iterable* into *k* parts. Set partitions are 3191 not order-preserving. 3192 3193 >>> iterable = 'abc' 3194 >>> for part in set_partitions(iterable, 2): 3195 ... print([''.join(p) for p in part]) 3196 ['a', 'bc'] 3197 ['ab', 'c'] 3198 ['b', 'ac'] 3199 3200 3201 If *k* is not given, every set partition is generated. 3202 3203 >>> iterable = 'abc' 3204 >>> for part in set_partitions(iterable): 3205 ... print([''.join(p) for p in part]) 3206 ['abc'] 3207 ['a', 'bc'] 3208 ['ab', 'c'] 3209 ['b', 'ac'] 3210 ['a', 'b', 'c'] 3211 3212 """ 3213 L = list(iterable) 3214 n = len(L) 3215 if k is not None: 3216 if k < 1: 3217 raise ValueError( 3218 "Can't partition in a negative or zero number of groups" 3219 ) 3220 elif k > n: 3221 return 3222 3223 def set_partitions_helper(L, k): 3224 n = len(L) 3225 if k == 1: 3226 yield [L] 3227 elif n == k: 3228 yield [[s] for s in L] 3229 else: 3230 e, *M = L 3231 for p in set_partitions_helper(M, k - 1): 3232 yield [[e], *p] 3233 for p in set_partitions_helper(M, k): 3234 for i in range(len(p)): 3235 yield p[:i] + [[e] + p[i]] + p[i + 1 :] 3236 3237 if k is None: 3238 for k in range(1, n + 1): 3239 yield from set_partitions_helper(L, k) 3240 else: 3241 yield from set_partitions_helper(L, k) 3242 3243 3244class time_limited: 3245 """ 3246 Yield items from *iterable* until *limit_seconds* have passed. 3247 If the time limit expires before all items have been yielded, the 3248 ``timed_out`` parameter will be set to ``True``. 3249 3250 >>> from time import sleep 3251 >>> def generator(): 3252 ... yield 1 3253 ... yield 2 3254 ... sleep(0.2) 3255 ... yield 3 3256 >>> iterable = time_limited(0.1, generator()) 3257 >>> list(iterable) 3258 [1, 2] 3259 >>> iterable.timed_out 3260 True 3261 3262 Note that the time is checked before each item is yielded, and iteration 3263 stops if the time elapsed is greater than *limit_seconds*. If your time 3264 limit is 1 second, but it takes 2 seconds to generate the first item from 3265 the iterable, the function will run for 2 seconds and not yield anything. 3266 3267 """ 3268 3269 def __init__(self, limit_seconds, iterable): 3270 if limit_seconds < 0: 3271 raise ValueError('limit_seconds must be positive') 3272 self.limit_seconds = limit_seconds 3273 self._iterable = iter(iterable) 3274 self._start_time = monotonic() 3275 self.timed_out = False 3276 3277 def __iter__(self): 3278 return self 3279 3280 def __next__(self): 3281 item = next(self._iterable) 3282 if monotonic() - self._start_time > self.limit_seconds: 3283 self.timed_out = True 3284 raise StopIteration 3285 3286 return item 3287 3288 3289def only(iterable, default=None, too_long=None): 3290 """If *iterable* has only one item, return it. 3291 If it has zero items, return *default*. 3292 If it has more than one item, raise the exception given by *too_long*, 3293 which is ``ValueError`` by default. 3294 3295 >>> only([], default='missing') 3296 'missing' 3297 >>> only([1]) 3298 1 3299 >>> only([1, 2]) # doctest: +IGNORE_EXCEPTION_DETAIL 3300 Traceback (most recent call last): 3301 ... 3302 ValueError: Expected exactly one item in iterable, but got 1, 2, 3303 and perhaps more.' 3304 >>> only([1, 2], too_long=TypeError) # doctest: +IGNORE_EXCEPTION_DETAIL 3305 Traceback (most recent call last): 3306 ... 3307 TypeError 3308 3309 Note that :func:`only` attempts to advance *iterable* twice to ensure there 3310 is only one item. See :func:`spy` or :func:`peekable` to check 3311 iterable contents less destructively. 3312 """ 3313 it = iter(iterable) 3314 first_value = next(it, default) 3315 3316 try: 3317 second_value = next(it) 3318 except StopIteration: 3319 pass 3320 else: 3321 msg = ( 3322 'Expected exactly one item in iterable, but got {!r}, {!r}, ' 3323 'and perhaps more.'.format(first_value, second_value) 3324 ) 3325 raise too_long or ValueError(msg) 3326 3327 return first_value 3328 3329 3330def ichunked(iterable, n): 3331 """Break *iterable* into sub-iterables with *n* elements each. 3332 :func:`ichunked` is like :func:`chunked`, but it yields iterables 3333 instead of lists. 3334 3335 If the sub-iterables are read in order, the elements of *iterable* 3336 won't be stored in memory. 3337 If they are read out of order, :func:`itertools.tee` is used to cache 3338 elements as necessary. 3339 3340 >>> from itertools import count 3341 >>> all_chunks = ichunked(count(), 4) 3342 >>> c_1, c_2, c_3 = next(all_chunks), next(all_chunks), next(all_chunks) 3343 >>> list(c_2) # c_1's elements have been cached; c_3's haven't been 3344 [4, 5, 6, 7] 3345 >>> list(c_1) 3346 [0, 1, 2, 3] 3347 >>> list(c_3) 3348 [8, 9, 10, 11] 3349 3350 """ 3351 source = iter(iterable) 3352 3353 while True: 3354 # Check to see whether we're at the end of the source iterable 3355 item = next(source, _marker) 3356 if item is _marker: 3357 return 3358 3359 # Clone the source and yield an n-length slice 3360 source, it = tee(chain([item], source)) 3361 yield islice(it, n) 3362 3363 # Advance the source iterable 3364 consume(source, n) 3365 3366 3367def distinct_combinations(iterable, r): 3368 """Yield the distinct combinations of *r* items taken from *iterable*. 3369 3370 >>> list(distinct_combinations([0, 0, 1], 2)) 3371 [(0, 0), (0, 1)] 3372 3373 Equivalent to ``set(combinations(iterable))``, except duplicates are not 3374 generated and thrown away. For larger input sequences this is much more 3375 efficient. 3376 3377 """ 3378 if r < 0: 3379 raise ValueError('r must be non-negative') 3380 elif r == 0: 3381 yield () 3382 return 3383 pool = tuple(iterable) 3384 generators = [unique_everseen(enumerate(pool), key=itemgetter(1))] 3385 current_combo = [None] * r 3386 level = 0 3387 while generators: 3388 try: 3389 cur_idx, p = next(generators[-1]) 3390 except StopIteration: 3391 generators.pop() 3392 level -= 1 3393 continue 3394 current_combo[level] = p 3395 if level + 1 == r: 3396 yield tuple(current_combo) 3397 else: 3398 generators.append( 3399 unique_everseen( 3400 enumerate(pool[cur_idx + 1 :], cur_idx + 1), 3401 key=itemgetter(1), 3402 ) 3403 ) 3404 level += 1 3405 3406 3407def filter_except(validator, iterable, *exceptions): 3408 """Yield the items from *iterable* for which the *validator* function does 3409 not raise one of the specified *exceptions*. 3410 3411 *validator* is called for each item in *iterable*. 3412 It should be a function that accepts one argument and raises an exception 3413 if that item is not valid. 3414 3415 >>> iterable = ['1', '2', 'three', '4', None] 3416 >>> list(filter_except(int, iterable, ValueError, TypeError)) 3417 ['1', '2', '4'] 3418 3419 If an exception other than one given by *exceptions* is raised by 3420 *validator*, it is raised like normal. 3421 """ 3422 for item in iterable: 3423 try: 3424 validator(item) 3425 except exceptions: 3426 pass 3427 else: 3428 yield item 3429 3430 3431def map_except(function, iterable, *exceptions): 3432 """Transform each item from *iterable* with *function* and yield the 3433 result, unless *function* raises one of the specified *exceptions*. 3434 3435 *function* is called to transform each item in *iterable*. 3436 It should accept one argument. 3437 3438 >>> iterable = ['1', '2', 'three', '4', None] 3439 >>> list(map_except(int, iterable, ValueError, TypeError)) 3440 [1, 2, 4] 3441 3442 If an exception other than one given by *exceptions* is raised by 3443 *function*, it is raised like normal. 3444 """ 3445 for item in iterable: 3446 try: 3447 yield function(item) 3448 except exceptions: 3449 pass 3450 3451 3452def map_if(iterable, pred, func, func_else=lambda x: x): 3453 """Evaluate each item from *iterable* using *pred*. If the result is 3454 equivalent to ``True``, transform the item with *func* and yield it. 3455 Otherwise, transform the item with *func_else* and yield it. 3456 3457 *pred*, *func*, and *func_else* should each be functions that accept 3458 one argument. By default, *func_else* is the identity function. 3459 3460 >>> from math import sqrt 3461 >>> iterable = list(range(-5, 5)) 3462 >>> iterable 3463 [-5, -4, -3, -2, -1, 0, 1, 2, 3, 4] 3464 >>> list(map_if(iterable, lambda x: x > 3, lambda x: 'toobig')) 3465 [-5, -4, -3, -2, -1, 0, 1, 2, 3, 'toobig'] 3466 >>> list(map_if(iterable, lambda x: x >= 0, 3467 ... lambda x: f'{sqrt(x):.2f}', lambda x: None)) 3468 [None, None, None, None, None, '0.00', '1.00', '1.41', '1.73', '2.00'] 3469 """ 3470 for item in iterable: 3471 yield func(item) if pred(item) else func_else(item) 3472 3473 3474def _sample_unweighted(iterable, k): 3475 # Implementation of "Algorithm L" from the 1994 paper by Kim-Hung Li: 3476 # "Reservoir-Sampling Algorithms of Time Complexity O(n(1+log(N/n)))". 3477 3478 # Fill up the reservoir (collection of samples) with the first `k` samples 3479 reservoir = take(k, iterable) 3480 3481 # Generate random number that's the largest in a sample of k U(0,1) numbers 3482 # Largest order statistic: https://en.wikipedia.org/wiki/Order_statistic 3483 W = exp(log(random()) / k) 3484 3485 # The number of elements to skip before changing the reservoir is a random 3486 # number with a geometric distribution. Sample it using random() and logs. 3487 next_index = k + floor(log(random()) / log(1 - W)) 3488 3489 for index, element in enumerate(iterable, k): 3490 3491 if index == next_index: 3492 reservoir[randrange(k)] = element 3493 # The new W is the largest in a sample of k U(0, `old_W`) numbers 3494 W *= exp(log(random()) / k) 3495 next_index += floor(log(random()) / log(1 - W)) + 1 3496 3497 return reservoir 3498 3499 3500def _sample_weighted(iterable, k, weights): 3501 # Implementation of "A-ExpJ" from the 2006 paper by Efraimidis et al. : 3502 # "Weighted random sampling with a reservoir". 3503 3504 # Log-transform for numerical stability for weights that are small/large 3505 weight_keys = (log(random()) / weight for weight in weights) 3506 3507 # Fill up the reservoir (collection of samples) with the first `k` 3508 # weight-keys and elements, then heapify the list. 3509 reservoir = take(k, zip(weight_keys, iterable)) 3510 heapify(reservoir) 3511 3512 # The number of jumps before changing the reservoir is a random variable 3513 # with an exponential distribution. Sample it using random() and logs. 3514 smallest_weight_key, _ = reservoir[0] 3515 weights_to_skip = log(random()) / smallest_weight_key 3516 3517 for weight, element in zip(weights, iterable): 3518 if weight >= weights_to_skip: 3519 # The notation here is consistent with the paper, but we store 3520 # the weight-keys in log-space for better numerical stability. 3521 smallest_weight_key, _ = reservoir[0] 3522 t_w = exp(weight * smallest_weight_key) 3523 r_2 = uniform(t_w, 1) # generate U(t_w, 1) 3524 weight_key = log(r_2) / weight 3525 heapreplace(reservoir, (weight_key, element)) 3526 smallest_weight_key, _ = reservoir[0] 3527 weights_to_skip = log(random()) / smallest_weight_key 3528 else: 3529 weights_to_skip -= weight 3530 3531 # Equivalent to [element for weight_key, element in sorted(reservoir)] 3532 return [heappop(reservoir)[1] for _ in range(k)] 3533 3534 3535def sample(iterable, k, weights=None): 3536 """Return a *k*-length list of elements chosen (without replacement) 3537 from the *iterable*. Like :func:`random.sample`, but works on iterables 3538 of unknown length. 3539 3540 >>> iterable = range(100) 3541 >>> sample(iterable, 5) # doctest: +SKIP 3542 [81, 60, 96, 16, 4] 3543 3544 An iterable with *weights* may also be given: 3545 3546 >>> iterable = range(100) 3547 >>> weights = (i * i + 1 for i in range(100)) 3548 >>> sampled = sample(iterable, 5, weights=weights) # doctest: +SKIP 3549 [79, 67, 74, 66, 78] 3550 3551 The algorithm can also be used to generate weighted random permutations. 3552 The relative weight of each item determines the probability that it 3553 appears late in the permutation. 3554 3555 >>> data = "abcdefgh" 3556 >>> weights = range(1, len(data) + 1) 3557 >>> sample(data, k=len(data), weights=weights) # doctest: +SKIP 3558 ['c', 'a', 'b', 'e', 'g', 'd', 'h', 'f'] 3559 """ 3560 if k == 0: 3561 return [] 3562 3563 iterable = iter(iterable) 3564 if weights is None: 3565 return _sample_unweighted(iterable, k) 3566 else: 3567 weights = iter(weights) 3568 return _sample_weighted(iterable, k, weights) 3569 3570 3571def is_sorted(iterable, key=None, reverse=False, strict=False): 3572 """Returns ``True`` if the items of iterable are in sorted order, and 3573 ``False`` otherwise. *key* and *reverse* have the same meaning that they do 3574 in the built-in :func:`sorted` function. 3575 3576 >>> is_sorted(['1', '2', '3', '4', '5'], key=int) 3577 True 3578 >>> is_sorted([5, 4, 3, 1, 2], reverse=True) 3579 False 3580 3581 If *strict*, tests for strict sorting, that is, returns ``False`` if equal 3582 elements are found: 3583 3584 >>> is_sorted([1, 2, 2]) 3585 True 3586 >>> is_sorted([1, 2, 2], strict=True) 3587 False 3588 3589 The function returns ``False`` after encountering the first out-of-order 3590 item. If there are no out-of-order items, the iterable is exhausted. 3591 """ 3592 3593 compare = (le if reverse else ge) if strict else (lt if reverse else gt) 3594 it = iterable if key is None else map(key, iterable) 3595 return not any(starmap(compare, pairwise(it))) 3596 3597 3598class AbortThread(BaseException): 3599 pass 3600 3601 3602class callback_iter: 3603 """Convert a function that uses callbacks to an iterator. 3604 3605 Let *func* be a function that takes a `callback` keyword argument. 3606 For example: 3607 3608 >>> def func(callback=None): 3609 ... for i, c in [(1, 'a'), (2, 'b'), (3, 'c')]: 3610 ... if callback: 3611 ... callback(i, c) 3612 ... return 4 3613 3614 3615 Use ``with callback_iter(func)`` to get an iterator over the parameters 3616 that are delivered to the callback. 3617 3618 >>> with callback_iter(func) as it: 3619 ... for args, kwargs in it: 3620 ... print(args) 3621 (1, 'a') 3622 (2, 'b') 3623 (3, 'c') 3624 3625 The function will be called in a background thread. The ``done`` property 3626 indicates whether it has completed execution. 3627 3628 >>> it.done 3629 True 3630 3631 If it completes successfully, its return value will be available 3632 in the ``result`` property. 3633 3634 >>> it.result 3635 4 3636 3637 Notes: 3638 3639 * If the function uses some keyword argument besides ``callback``, supply 3640 *callback_kwd*. 3641 * If it finished executing, but raised an exception, accessing the 3642 ``result`` property will raise the same exception. 3643 * If it hasn't finished executing, accessing the ``result`` 3644 property from within the ``with`` block will raise ``RuntimeError``. 3645 * If it hasn't finished executing, accessing the ``result`` property from 3646 outside the ``with`` block will raise a 3647 ``more_itertools.AbortThread`` exception. 3648 * Provide *wait_seconds* to adjust how frequently the it is polled for 3649 output. 3650 3651 """ 3652 3653 def __init__(self, func, callback_kwd='callback', wait_seconds=0.1): 3654 self._func = func 3655 self._callback_kwd = callback_kwd 3656 self._aborted = False 3657 self._future = None 3658 self._wait_seconds = wait_seconds 3659 self._executor = ThreadPoolExecutor(max_workers=1) 3660 self._iterator = self._reader() 3661 3662 def __enter__(self): 3663 return self 3664 3665 def __exit__(self, exc_type, exc_value, traceback): 3666 self._aborted = True 3667 self._executor.shutdown() 3668 3669 def __iter__(self): 3670 return self 3671 3672 def __next__(self): 3673 return next(self._iterator) 3674 3675 @property 3676 def done(self): 3677 if self._future is None: 3678 return False 3679 return self._future.done() 3680 3681 @property 3682 def result(self): 3683 if not self.done: 3684 raise RuntimeError('Function has not yet completed') 3685 3686 return self._future.result() 3687 3688 def _reader(self): 3689 q = Queue() 3690 3691 def callback(*args, **kwargs): 3692 if self._aborted: 3693 raise AbortThread('canceled by user') 3694 3695 q.put((args, kwargs)) 3696 3697 self._future = self._executor.submit( 3698 self._func, **{self._callback_kwd: callback} 3699 ) 3700 3701 while True: 3702 try: 3703 item = q.get(timeout=self._wait_seconds) 3704 except Empty: 3705 pass 3706 else: 3707 q.task_done() 3708 yield item 3709 3710 if self._future.done(): 3711 break 3712 3713 remaining = [] 3714 while True: 3715 try: 3716 item = q.get_nowait() 3717 except Empty: 3718 break 3719 else: 3720 q.task_done() 3721 remaining.append(item) 3722 q.join() 3723 yield from remaining 3724 3725 3726def windowed_complete(iterable, n): 3727 """ 3728 Yield ``(beginning, middle, end)`` tuples, where: 3729 3730 * Each ``middle`` has *n* items from *iterable* 3731 * Each ``beginning`` has the items before the ones in ``middle`` 3732 * Each ``end`` has the items after the ones in ``middle`` 3733 3734 >>> iterable = range(7) 3735 >>> n = 3 3736 >>> for beginning, middle, end in windowed_complete(iterable, n): 3737 ... print(beginning, middle, end) 3738 () (0, 1, 2) (3, 4, 5, 6) 3739 (0,) (1, 2, 3) (4, 5, 6) 3740 (0, 1) (2, 3, 4) (5, 6) 3741 (0, 1, 2) (3, 4, 5) (6,) 3742 (0, 1, 2, 3) (4, 5, 6) () 3743 3744 Note that *n* must be at least 0 and most equal to the length of 3745 *iterable*. 3746 3747 This function will exhaust the iterable and may require significant 3748 storage. 3749 """ 3750 if n < 0: 3751 raise ValueError('n must be >= 0') 3752 3753 seq = tuple(iterable) 3754 size = len(seq) 3755 3756 if n > size: 3757 raise ValueError('n must be <= len(seq)') 3758 3759 for i in range(size - n + 1): 3760 beginning = seq[:i] 3761 middle = seq[i : i + n] 3762 end = seq[i + n :] 3763 yield beginning, middle, end 3764 3765 3766def all_unique(iterable, key=None): 3767 """ 3768 Returns ``True`` if all the elements of *iterable* are unique (no two 3769 elements are equal). 3770 3771 >>> all_unique('ABCB') 3772 False 3773 3774 If a *key* function is specified, it will be used to make comparisons. 3775 3776 >>> all_unique('ABCb') 3777 True 3778 >>> all_unique('ABCb', str.lower) 3779 False 3780 3781 The function returns as soon as the first non-unique element is 3782 encountered. Iterables with a mix of hashable and unhashable items can 3783 be used, but the function will be slower for unhashable items. 3784 """ 3785 seenset = set() 3786 seenset_add = seenset.add 3787 seenlist = [] 3788 seenlist_add = seenlist.append 3789 for element in map(key, iterable) if key else iterable: 3790 try: 3791 if element in seenset: 3792 return False 3793 seenset_add(element) 3794 except TypeError: 3795 if element in seenlist: 3796 return False 3797 seenlist_add(element) 3798 return True 3799 3800 3801def nth_product(index, *args): 3802 """Equivalent to ``list(product(*args))[index]``. 3803 3804 The products of *args* can be ordered lexicographically. 3805 :func:`nth_product` computes the product at sort position *index* without 3806 computing the previous products. 3807 3808 >>> nth_product(8, range(2), range(2), range(2), range(2)) 3809 (1, 0, 0, 0) 3810 3811 ``IndexError`` will be raised if the given *index* is invalid. 3812 """ 3813 pools = list(map(tuple, reversed(args))) 3814 ns = list(map(len, pools)) 3815 3816 c = reduce(mul, ns) 3817 3818 if index < 0: 3819 index += c 3820 3821 if not 0 <= index < c: 3822 raise IndexError 3823 3824 result = [] 3825 for pool, n in zip(pools, ns): 3826 result.append(pool[index % n]) 3827 index //= n 3828 3829 return tuple(reversed(result)) 3830 3831 3832def nth_permutation(iterable, r, index): 3833 """Equivalent to ``list(permutations(iterable, r))[index]``` 3834 3835 The subsequences of *iterable* that are of length *r* where order is 3836 important can be ordered lexicographically. :func:`nth_permutation` 3837 computes the subsequence at sort position *index* directly, without 3838 computing the previous subsequences. 3839 3840 >>> nth_permutation('ghijk', 2, 5) 3841 ('h', 'i') 3842 3843 ``ValueError`` will be raised If *r* is negative or greater than the length 3844 of *iterable*. 3845 ``IndexError`` will be raised if the given *index* is invalid. 3846 """ 3847 pool = list(iterable) 3848 n = len(pool) 3849 3850 if r is None or r == n: 3851 r, c = n, factorial(n) 3852 elif not 0 <= r < n: 3853 raise ValueError 3854 else: 3855 c = factorial(n) // factorial(n - r) 3856 3857 if index < 0: 3858 index += c 3859 3860 if not 0 <= index < c: 3861 raise IndexError 3862 3863 if c == 0: 3864 return tuple() 3865 3866 result = [0] * r 3867 q = index * factorial(n) // c if r < n else index 3868 for d in range(1, n + 1): 3869 q, i = divmod(q, d) 3870 if 0 <= n - d < r: 3871 result[n - d] = i 3872 if q == 0: 3873 break 3874 3875 return tuple(map(pool.pop, result)) 3876 3877 3878def value_chain(*args): 3879 """Yield all arguments passed to the function in the same order in which 3880 they were passed. If an argument itself is iterable then iterate over its 3881 values. 3882 3883 >>> list(value_chain(1, 2, 3, [4, 5, 6])) 3884 [1, 2, 3, 4, 5, 6] 3885 3886 Binary and text strings are not considered iterable and are emitted 3887 as-is: 3888 3889 >>> list(value_chain('12', '34', ['56', '78'])) 3890 ['12', '34', '56', '78'] 3891 3892 3893 Multiple levels of nesting are not flattened. 3894 3895 """ 3896 for value in args: 3897 if isinstance(value, (str, bytes)): 3898 yield value 3899 continue 3900 try: 3901 yield from value 3902 except TypeError: 3903 yield value 3904 3905 3906def product_index(element, *args): 3907 """Equivalent to ``list(product(*args)).index(element)`` 3908 3909 The products of *args* can be ordered lexicographically. 3910 :func:`product_index` computes the first index of *element* without 3911 computing the previous products. 3912 3913 >>> product_index([8, 2], range(10), range(5)) 3914 42 3915 3916 ``ValueError`` will be raised if the given *element* isn't in the product 3917 of *args*. 3918 """ 3919 index = 0 3920 3921 for x, pool in zip_longest(element, args, fillvalue=_marker): 3922 if x is _marker or pool is _marker: 3923 raise ValueError('element is not a product of args') 3924 3925 pool = tuple(pool) 3926 index = index * len(pool) + pool.index(x) 3927 3928 return index 3929 3930 3931def combination_index(element, iterable): 3932 """Equivalent to ``list(combinations(iterable, r)).index(element)`` 3933 3934 The subsequences of *iterable* that are of length *r* can be ordered 3935 lexicographically. :func:`combination_index` computes the index of the 3936 first *element*, without computing the previous combinations. 3937 3938 >>> combination_index('adf', 'abcdefg') 3939 10 3940 3941 ``ValueError`` will be raised if the given *element* isn't one of the 3942 combinations of *iterable*. 3943 """ 3944 element = enumerate(element) 3945 k, y = next(element, (None, None)) 3946 if k is None: 3947 return 0 3948 3949 indexes = [] 3950 pool = enumerate(iterable) 3951 for n, x in pool: 3952 if x == y: 3953 indexes.append(n) 3954 tmp, y = next(element, (None, None)) 3955 if tmp is None: 3956 break 3957 else: 3958 k = tmp 3959 else: 3960 raise ValueError('element is not a combination of iterable') 3961 3962 n, _ = last(pool, default=(n, None)) 3963 3964 # Python versiosn below 3.8 don't have math.comb 3965 index = 1 3966 for i, j in enumerate(reversed(indexes), start=1): 3967 j = n - j 3968 if i <= j: 3969 index += factorial(j) // (factorial(i) * factorial(j - i)) 3970 3971 return factorial(n + 1) // (factorial(k + 1) * factorial(n - k)) - index 3972 3973 3974def permutation_index(element, iterable): 3975 """Equivalent to ``list(permutations(iterable, r)).index(element)``` 3976 3977 The subsequences of *iterable* that are of length *r* where order is 3978 important can be ordered lexicographically. :func:`permutation_index` 3979 computes the index of the first *element* directly, without computing 3980 the previous permutations. 3981 3982 >>> permutation_index([1, 3, 2], range(5)) 3983 19 3984 3985 ``ValueError`` will be raised if the given *element* isn't one of the 3986 permutations of *iterable*. 3987 """ 3988 index = 0 3989 pool = list(iterable) 3990 for i, x in zip(range(len(pool), -1, -1), element): 3991 r = pool.index(x) 3992 index = index * i + r 3993 del pool[r] 3994 3995 return index 3996 3997 3998class countable: 3999 """Wrap *iterable* and keep a count of how many items have been consumed. 4000 4001 The ``items_seen`` attribute starts at ``0`` and increments as the iterable 4002 is consumed: 4003 4004 >>> iterable = map(str, range(10)) 4005 >>> it = countable(iterable) 4006 >>> it.items_seen 4007 0 4008 >>> next(it), next(it) 4009 ('0', '1') 4010 >>> list(it) 4011 ['2', '3', '4', '5', '6', '7', '8', '9'] 4012 >>> it.items_seen 4013 10 4014 """ 4015 4016 def __init__(self, iterable): 4017 self._it = iter(iterable) 4018 self.items_seen = 0 4019 4020 def __iter__(self): 4021 return self 4022 4023 def __next__(self): 4024 item = next(self._it) 4025 self.items_seen += 1 4026 4027 return item 4028 4029 4030def chunked_even(iterable, n): 4031 """Break *iterable* into lists of approximately length *n*. 4032 Items are distributed such the lengths of the lists differ by at most 4033 1 item. 4034 4035 >>> iterable = [1, 2, 3, 4, 5, 6, 7] 4036 >>> n = 3 4037 >>> list(chunked_even(iterable, n)) # List lengths: 3, 2, 2 4038 [[1, 2, 3], [4, 5], [6, 7]] 4039 >>> list(chunked(iterable, n)) # List lengths: 3, 3, 1 4040 [[1, 2, 3], [4, 5, 6], [7]] 4041 4042 """ 4043 4044 len_method = getattr(iterable, '__len__', None) 4045 4046 if len_method is None: 4047 return _chunked_even_online(iterable, n) 4048 else: 4049 return _chunked_even_finite(iterable, len_method(), n) 4050 4051 4052def _chunked_even_online(iterable, n): 4053 buffer = [] 4054 maxbuf = n + (n - 2) * (n - 1) 4055 for x in iterable: 4056 buffer.append(x) 4057 if len(buffer) == maxbuf: 4058 yield buffer[:n] 4059 buffer = buffer[n:] 4060 yield from _chunked_even_finite(buffer, len(buffer), n) 4061 4062 4063def _chunked_even_finite(iterable, N, n): 4064 if N < 1: 4065 return 4066 4067 # Lists are either size `full_size <= n` or `partial_size = full_size - 1` 4068 q, r = divmod(N, n) 4069 num_lists = q + (1 if r > 0 else 0) 4070 q, r = divmod(N, num_lists) 4071 full_size = q + (1 if r > 0 else 0) 4072 partial_size = full_size - 1 4073 num_full = N - partial_size * num_lists 4074 num_partial = num_lists - num_full 4075 4076 buffer = [] 4077 iterator = iter(iterable) 4078 4079 # Yield num_full lists of full_size 4080 for x in iterator: 4081 buffer.append(x) 4082 if len(buffer) == full_size: 4083 yield buffer 4084 buffer = [] 4085 num_full -= 1 4086 if num_full <= 0: 4087 break 4088 4089 # Yield num_partial lists of partial_size 4090 for x in iterator: 4091 buffer.append(x) 4092 if len(buffer) == partial_size: 4093 yield buffer 4094 buffer = [] 4095 num_partial -= 1 4096 4097 4098def zip_broadcast(*objects, scalar_types=(str, bytes), strict=False): 4099 """A version of :func:`zip` that "broadcasts" any scalar 4100 (i.e., non-iterable) items into output tuples. 4101 4102 >>> iterable_1 = [1, 2, 3] 4103 >>> iterable_2 = ['a', 'b', 'c'] 4104 >>> scalar = '_' 4105 >>> list(zip_broadcast(iterable_1, iterable_2, scalar)) 4106 [(1, 'a', '_'), (2, 'b', '_'), (3, 'c', '_')] 4107 4108 The *scalar_types* keyword argument determines what types are considered 4109 scalar. It is set to ``(str, bytes)`` by default. Set it to ``None`` to 4110 treat strings and byte strings as iterable: 4111 4112 >>> list(zip_broadcast('abc', 0, 'xyz', scalar_types=None)) 4113 [('a', 0, 'x'), ('b', 0, 'y'), ('c', 0, 'z')] 4114 4115 If the *strict* keyword argument is ``True``, then 4116 ``UnequalIterablesError`` will be raised if any of the iterables have 4117 different lengthss. 4118 """ 4119 4120 def is_scalar(obj): 4121 if scalar_types and isinstance(obj, scalar_types): 4122 return True 4123 try: 4124 iter(obj) 4125 except TypeError: 4126 return True 4127 else: 4128 return False 4129 4130 size = len(objects) 4131 if not size: 4132 return 4133 4134 iterables, iterable_positions = [], [] 4135 scalars, scalar_positions = [], [] 4136 for i, obj in enumerate(objects): 4137 if is_scalar(obj): 4138 scalars.append(obj) 4139 scalar_positions.append(i) 4140 else: 4141 iterables.append(iter(obj)) 4142 iterable_positions.append(i) 4143 4144 if len(scalars) == size: 4145 yield tuple(objects) 4146 return 4147 4148 zipper = _zip_equal if strict else zip 4149 for item in zipper(*iterables): 4150 new_item = [None] * size 4151 4152 for i, elem in zip(iterable_positions, item): 4153 new_item[i] = elem 4154 4155 for i, elem in zip(scalar_positions, scalars): 4156 new_item[i] = elem 4157 4158 yield tuple(new_item) 4159 4160 4161def unique_in_window(iterable, n, key=None): 4162 """Yield the items from *iterable* that haven't been seen recently. 4163 *n* is the size of the lookback window. 4164 4165 >>> iterable = [0, 1, 0, 2, 3, 0] 4166 >>> n = 3 4167 >>> list(unique_in_window(iterable, n)) 4168 [0, 1, 2, 3, 0] 4169 4170 The *key* function, if provided, will be used to determine uniqueness: 4171 4172 >>> list(unique_in_window('abAcda', 3, key=lambda x: x.lower())) 4173 ['a', 'b', 'c', 'd', 'a'] 4174 4175 The items in *iterable* must be hashable. 4176 4177 """ 4178 if n <= 0: 4179 raise ValueError('n must be greater than 0') 4180 4181 window = deque(maxlen=n) 4182 uniques = set() 4183 use_key = key is not None 4184 4185 for item in iterable: 4186 k = key(item) if use_key else item 4187 if k in uniques: 4188 continue 4189 4190 if len(uniques) == n: 4191 uniques.discard(window[0]) 4192 4193 uniques.add(k) 4194 window.append(k) 4195 4196 yield item 4197 4198 4199def duplicates_everseen(iterable, key=None): 4200 """Yield duplicate elements after their first appearance. 4201 4202 >>> list(duplicates_everseen('mississippi')) 4203 ['s', 'i', 's', 's', 'i', 'p', 'i'] 4204 >>> list(duplicates_everseen('AaaBbbCccAaa', str.lower)) 4205 ['a', 'a', 'b', 'b', 'c', 'c', 'A', 'a', 'a'] 4206 4207 This function is analagous to :func:`unique_everseen` and is subject to 4208 the same performance considerations. 4209 4210 """ 4211 seen_set = set() 4212 seen_list = [] 4213 use_key = key is not None 4214 4215 for element in iterable: 4216 k = key(element) if use_key else element 4217 try: 4218 if k not in seen_set: 4219 seen_set.add(k) 4220 else: 4221 yield element 4222 except TypeError: 4223 if k not in seen_list: 4224 seen_list.append(k) 4225 else: 4226 yield element 4227 4228 4229def duplicates_justseen(iterable, key=None): 4230 """Yields serially-duplicate elements after their first appearance. 4231 4232 >>> list(duplicates_justseen('mississippi')) 4233 ['s', 's', 'p'] 4234 >>> list(duplicates_justseen('AaaBbbCccAaa', str.lower)) 4235 ['a', 'a', 'b', 'b', 'c', 'c', 'a', 'a'] 4236 4237 This function is analagous to :func:`unique_justseen`. 4238 4239 """ 4240 return flatten( 4241 map( 4242 lambda group_tuple: islice_extended(group_tuple[1])[1:], 4243 groupby(iterable, key), 4244 ) 4245 ) 4246 4247 4248def minmax(iterable_or_value, *others, key=None, default=_marker): 4249 """Returns both the smallest and largest items in an iterable 4250 or the largest of two or more arguments. 4251 4252 >>> minmax([3, 1, 5]) 4253 (1, 5) 4254 4255 >>> minmax(4, 2, 6) 4256 (2, 6) 4257 4258 If a *key* function is provided, it will be used to transform the input 4259 items for comparison. 4260 4261 >>> minmax([5, 30], key=str) # '30' sorts before '5' 4262 (30, 5) 4263 4264 If a *default* value is provided, it will be returned if there are no 4265 input items. 4266 4267 >>> minmax([], default=(0, 0)) 4268 (0, 0) 4269 4270 Otherwise ``ValueError`` is raised. 4271 4272 This function is based on the 4273 `recipe <http://code.activestate.com/recipes/577916/>`__ by 4274 Raymond Hettinger and takes care to minimize the number of comparisons 4275 performed. 4276 """ 4277 iterable = (iterable_or_value, *others) if others else iterable_or_value 4278 4279 it = iter(iterable) 4280 4281 try: 4282 lo = hi = next(it) 4283 except StopIteration as e: 4284 if default is _marker: 4285 raise ValueError( 4286 '`minmax()` argument is an empty iterable. ' 4287 'Provide a `default` value to suppress this error.' 4288 ) from e 4289 return default 4290 4291 # Different branches depending on the presence of key. This saves a lot 4292 # of unimportant copies which would slow the "key=None" branch 4293 # significantly down. 4294 if key is None: 4295 for x, y in zip_longest(it, it, fillvalue=lo): 4296 if y < x: 4297 x, y = y, x 4298 if x < lo: 4299 lo = x 4300 if hi < y: 4301 hi = y 4302 4303 else: 4304 lo_key = hi_key = key(lo) 4305 4306 for x, y in zip_longest(it, it, fillvalue=lo): 4307 4308 x_key, y_key = key(x), key(y) 4309 4310 if y_key < x_key: 4311 x, y, x_key, y_key = y, x, y_key, x_key 4312 if x_key < lo_key: 4313 lo, lo_key = x, x_key 4314 if hi_key < y_key: 4315 hi, hi_key = y, y_key 4316 4317 return lo, hi 4318