From d1d5e24692ec8dae357f26b186c373ca9c9d9868 Mon Sep 17 00:00:00 2001 From: upstream Date: Wed, 7 Oct 2026 20:10:16 +0200 Subject: [PATCH] Upstream more-itertools/more-itertools state of the touched files before 71e070b35d --- docs/versions.rst | 915 +++++ more_itertools/more.py | 5644 ++++++++++++++++++++++++++++++ tests/test_more.py | 7444 ++++++++++++++++++++++++++++++++++++++++ 3 files changed, 14003 insertions(+) create mode 100644 docs/versions.rst create mode 100644 more_itertools/more.py create mode 100644 tests/test_more.py diff --git a/docs/versions.rst b/docs/versions.rst new file mode 100644 index 0000000..ecc3811 --- /dev/null +++ b/docs/versions.rst @@ -0,0 +1,915 @@ +=============== +Version History +=============== + +.. automodule:: more_itertools + :noindex: + +Unreleased +---------- + +* Changes to existing functions: + * :func:`bucket` now honours a falsey *validator* callable (such as an empty ``set`` subclass that defines ``__call__``) instead of silently replacing it with the default that accepts every key. + * :func:`ichunked` now matches :func:`chunked` for ``n = 0`` and negative *n*: it yields nothing for a zero chunk size (without consuming the source) and raises ``ValueError: n must be at least 0`` instead of leaking :func:`itertools.islice`'s error. + * :func:`running_min`, :func:`running_max`, :func:`running_mean`, and :func:`running_statistics` now consistently validate *maxlen* as an integer and support the index protocol. + * :func:`numeric_range` now preserves its original values and length when reversed, including floating-point ranges and ranges near datetime limits. + * :func:`constrained_batches` now raises ``ValueError`` for a nonpositive *max_count*. + * :func:`running_min`, :func:`running_max`, :func:`running_mean`, and :func:`running_statistics` now consistently validate *maxlen* as an integer and support the index protocol. + * :func:`iter_index` was fixed to accept negative *start* and *stop* with general iterables (thanks to gaoflow) + * :func:`split_before`, :func:`split_after`, and :func:`split_when` no longer yield an empty list for an empty *iterable* when *maxsplit* is ``0`` + * :func:`one` and :func:`only` were fixed to raise a user-supplied *too_short*/*too_long* exception even when the exception object is falsy, and to no longer build the default error message (which reprs iterable items) when a custom exception is given (thanks to itzzdev09) + +11.1.0 +------ + +* Changes to existing functions: + * :func:`numeric_range` was updated to fix its handling of empty ranges (thanks to rhettinger) + * :func:`peekable` was updated to fix typing issues (thanks to DORI2001, powellnorma, Pandede, m9810223, and rhettinger) + * :func:`islice_extended` was optimized for memory usage and speed (thanks to ben42code, rhettinger, and pochmann) + * :func:`serialize` now supports the generator methods ``throw``, ``send``, and ``close`` (thanks to rhettinger) + * :func:`seekable` now supports implements ``__getitem__`` for cached elements (thanks to SAY-5, jenstroeger, and JamesParrott) + +11.0.2 +------ + +* Updates + * The changes to the type hints for :func:`always_iterable` in 11.0.0 have been + reverted due to reported problems (thanks to ngoldbaum and rhettinger) + * :func:`zip_equal` was removed in 11.0.0. It had been deprecated and raising ``DeprecationWarning`` + since 2021, but it's removal should have been documented in 11.0.0's release notes. + We regret the error. + * :func:`running_statistics` has been moved from ``more_itertools.more`` to + ``more_itertools.recipes``. Its docstring was also improved (thanks to mastash3ff + and rhettinger) + +11.0.1 +------ + +* Updates + * This release restores :func:`pairwise`, which was mistakenly removed in 11.0.0 + instead of being deprecated. It is now marked as deprecated and will be removed + in a future major release. Use :func:`itertools.pairwise` as a replacement. + +11.0.0 +------ + +* Potentially breaking changes + * Python 3.9 support was dropped, since it went EOL on 2025-10-31 + * :func:`zip_equal` was removed. Applications should use the standard library + function :func:`zip` with ``strict=True`` as a replacement. + * :func:`callback_iter` is deprecated. It will be removed in a future major release. + * :func:`iequals` no longer returns ``True`` when called with ``([], [ANY])`` (thanks to rhettinger and pochmann3) + * The ``pred`` argument for :func:`locate` and :func:`replace` must now be able to + handle a variable number of arguments. See their docstrings for details. (thanks to veeceey, james-wasson, and pochmann3) + * :func:`unique_everseen` now raises ``TypeError`` when input elements are not hashable. + The silent fallback mechanism for unhashable inputs has been removed. (thanks to rhettinger and karlicoss) + * :func:`windowed` now raises ``ValueError`` when given a window size of ``0`` (thanks to rhettinger and brevzin) + * Parameter names for :func:`flatten`, :func:`iter_except`, :func:`repeatfunc`, :func:`prepend`, :func:`transpose`, + and :func:`sum_of_squares` have been changed. Code that uses these functions will need to be updated if it + uses keywords for positional arguments. (thanks to rhettinger) + +* New functions + * :func:`concurrent_tee`: similar to :func:`tee`, but with guaranteed threading + semantics (thanks to rhettinger) + * :func:`random_derangement`: returns a random derangement of elements from an + iterable (thanks to rhettinger) + * :func:`running_mean`: yields the mean of values in an iterable, or in a sliding + window (thanks to rhettinger) + * :func:`running_statistics`: provides statistics for the values in an iterable, or + in a sliding window (thanks to rhettinger) + * :func:`serialize`: wraps a non-concurrent iterator with a lock to enforce + sequential access (thanks to rhettinger) + * :func:`synchronized`: wraps an iterator-returning callable to make its iterators + thread-safe (thanks to rhettinger) + * :func:`sized_iterator`: wraps an iterable with a known length and implements ``__len__`` (thanks to assaf127) + +* Changes to existing functions + * The type hints for :func:`always_iterable` were improved (thanks to rhettinger and maltevesper) + * A potential bug in :func:`callback_iter` was fixed + * A bug in :func:`exactly_n`'s handling of negative arguments was fixed (thanks to rhettinger) + * :func:`extract` now accepts a `monotonic` argument for improved performance (thanks to rhettinger) + * A bug in :func:`numeric_range`'s handling of negative steps was fixed (thanks to bysiber) + * :func:`grouper` implementation was updated to match the ``itertools`` docs (thanks to rhettinger) + * :func:`nth_product`, :func:`product_index`, :func:`gray_product`, + :func:`partial_product`, :func:`nth_product`, :func:`product_index`, and + :func:`random_product` now accept a ``repeat`` keyword argument (thanks to rhettinger) + * :func:`powerset_of_sets` can now construct ``frozenset`` instances (thanks to rhettinger) + * The type hints for :func:`stagger` were improved (thanks to nathanjmcdougall) + * Memory efficiency was improved for :func:`partition` was improved (thanks to rhettinger and yueyinqiu) + * Performance for the functions :func:`count_cycle`, :func:`difference`, + :func:`random_product`, was improved :func:`nth_combination`, :func:`nth_combination_with_replacement`, + :func:`nth_permuwas improvedtation`, :func:`nth_product`, :func:`product_index`, + :func:`random_permutation`, :func:`substrings`, and :func:`value_chain` was + improved (thanks to rhettinger) + * Performance for :func:`ichunked` was improved (thanks to pochmann3 and rhettinger) + +* Other changes + * The docstrings for several functions were improved (thanks to nathanjmcdougall, olliemath, r266-tech, and rhettinger) + + +10.8.0 +------ + +* New functions: + * :func:`derangements` was added (thanks to debruijn) + * :func:`argmin` and :func:`argmax` were added (thanks to rhettinger) + * :func:`running_median` was added (thanks to rhettinger) + * :func:`extract` was added (thanks to rhettinger) + * :func:`interleave_randomly` was added (thanks to ktbarrett) + +* Changes to existing functions: + * The type hints and docstring for :func:`batched` were improved (thanks to qobilidop and inventshah) + * The memory usage of :func:`islice_extended` was reduced (thanks to ben42code) + * The performance of :func:`sample` and :func:`consecutive_groups`, :func:`dft`, :func:`idft`, :func:`map_if`, :func:`count_cycle`, and :func:`tail` were improved (thanks to rhettinger) + * The performance of :func:`before_and_after`, :func:`mark_ends`, and :func:`interleave_longest` were improved (thanks to pochmann3) + * :func:`nth_prime` now accepts an ``approximate`` keyword. When set to ``True``, a faster but less accurate method is used to return a result. (thanks to rhettinger) + * :func:`last` now works when its input has ``__reversed__`` set to ``None`` (thanks to inventshah) + * The :func:`unzip` function was simplified (thanks to pochmann3) + * The :func:`reshape` function now accepts ``shape`` values that represent multidimensional matrices (thanks to rhettinger) + +* Other changes: + * An issue with dark themes and documentation display was fixed (thanks to pochmann3, moreati, and pradyunsg) + * Variable names in several functions were improved (thanks to rhettinger) + * The docstrings for :func:`dft`, :func:`idft`, :func:`minmax`, :func:`sample`, and :func:`multinomial` were improved (thanks to rhettinger) + * Packaging and package index metadata were improved (thanks to cdce8p) + * Several aspects of the documentation were improved (thanks to rhettinger, saadmanrafat) + * The Makefile now refers to `python` instead of `python3` (thanks to ktbarrett) + * Test coverage was improved (thanks to rhettinger) + * Python 3.14 is now tested by GitHub Actions + +10.7.0 +------ + +* New functions: + * :func:`multinomial` was added (thanks to rhettinger) + +* Changes to existing functions: + * :func:`ichunked`, :func:`iterate`, :func:`one`, :func:`only`, :func:`powerset_of_sets`, and :func:`strictly_n` were optimized (thanks to rhettinger) + * :func:`exactly_n` now uses less memory (thanks to rhettinger) + * :func:`dft` and :func:`idft` were optimized for Python versions below 3.12 (thanks to rhettinger) + * :func:`is_prime` no longer shares state with the users random number generator (thanks to rhettinger) + +* Other changes: + * Some docstring issues were fixed (thanks to lpulley and ricbit) + * The type hints for :func:`groupby_transform` were improved (thanks to rhettinger) + * The ``furo`` theme is now used for docs (thanks to AA-turner) + + +10.6.0 +------ + +* New functions: + * :func:`is_prime` and :func:`nth_prime` were added (thanks to JamesParrott and rhettinger) + * :func:`loops` was added (thanks to rhettinger) + +* Changes to existing functions: + * :func:`factor` was optimized to handle larger inputs and use less memory (thanks to rhettinger) + * :func:`spy` was optimized to enable nested calls (thanks to rhettinger) + * :func:`polynomial_from_roots` was made non-recursive and able to handle larger numbers of roots (thanks to pochmann3 and rhettinger) + * :func:`is_sorted` now only relies on less than comparisons (thanks to rhettinger) + * The docstring for :func:`outer_product` was improved (thanks to rhettinger) + * The type annotations for :func:`sample` were improved (thanks to rhettinger) + +* Other changes: + * Python 3.13 is officially supported. Python 3.8 is no longer officially supported. (thanks to hugovk, JamesParrott, and stankudrow) + * `mypy` checks were fixed (thanks to JamesParrott) + +10.5.0 +------ + +* Bug fixes + * A missing symbol in ``more.pyi`` was fixed (thanks to eberts-google and nathanielmanistaatgoogle) + +* Other changes + * :func:`all_equal` was optimized (thanks to pochmann3 and rhettinger) + +10.4.0 +------ + +* Changes to existing functions + * :func:`circular_shifts` now accepts a ``steps`` parameter (thanks to rhettinger) + * :func:`distinct_permutations` now accepts iterables with non-comparable items (thanks to hgustafsson, JamesParrott, and pochmann3) + * :class:`run_length`, :func:`totient`, :func:`sliding_window`, and :func:`triplewise` were optimized (thanks to rhettinger) + * :class:`ilen` was optimized (thanks to pochmann3 and rhettinger) + * :func:`sample` was improved, and now accepts ``counts`` and ``strict`` parameters (thanks to rhettinger) + * :func:`set_partitions` now accepts ``min_size`` and ``max_size`` parameters (thanks to Pandede) + * :func:`seekable`'s ``relative_seek`` method remembers previous calls (thanks to dkrikun) + * :func:`sort_together` now accepts a ``strict`` parameter (thanks to rhettinger and Pandede) + +* Other changes + * The docs for :func:`is_sorted` and :func:`strictly_n` were improved (thanks to pochmann3 and fakuivan) + * The typing information for :func:`windowed_complete`, :func:`zip_broadcast`, and :func:`zip_equal` were improved (thanks to m472, eyalho, akisatoon1, jbosboom, and Pandede) + +10.3.0 +------ + +* New functions + * :func:`powerset_of_sets`, :func:`dft`, and :func:`idft` (thanks to rhettinger) + * :func:`join_mappings` (thanks to NeilGirdhar and rhettinger) + * :func:`doublestarmap` (thanks to Skeen, monk-time, DamianB-BitFlipper, and ergoithz) + * :func:`unique` (thanks to rhettinger) + +* Changes to existing functions + * :func:`collapse`, :func:`chunked_even`, :func:`ichunked`, :func:`padded`, and :func:`windowed` were optimized and improved (thanks to james-wasson) + * :func:`totient` was optimized (thanks to rhettinger) + * :func:`roundrobin` was updated and improved (thanks to rhettinger) + * :func:`all_equal` now accepts a *key* parameter. + * The docs for :func:`value_chain` were improved (thanks to bjrtx) + * The type annotations for :class:`countable` were improved (thanks to aidanholm) + +* Other changes + * Unit tests were improved (thanks to haukex) + * Some documentation issues were fixed (thanks to bjrtx and DimitriPapadopoulos) + +10.2.0 +------ + +* New functions + * :func:`iter_suppress` (thanks to jaraco, pochmann, and rhettinger) + * :func:`filter_map` (thanks to struktured) + * :func:`classify_unique` (thanks to haukex) + * :func:`totient` (from the itertools docs) + * :func:`reshape` (from the itertools docs) + +* Changes to existing functions + * :func:`factor`, :func:`iter_index`, :func:`sieve`, and :func:`unique_justseen` were updated to match the itertools docs + * :func:`first` was optimized (thanks to pochmann) + * :func:`takewhile_inclusive` was refactored (thanks to eltoder) + * :func:`combination_with_replacement_index` was optimized (thanks to elliotwutingfeng and rhettinger) + * :func:`nth_permutation`, :func:`nth_combination_with_replacement`, :func:`combination_index`, and :func:`combination_with_replacement_index` were optimized (thanks to rhettinger) + * :func:`batched` now accepts a `strict` argument (adapted from itertools docs) + * :func:`time_limited` was improved for Windows (thanks to haukex) + +* Other changes + * Several typing updates were made (thanks to obaltian and ilai-deutel) + * Some documentation issues were fixed (thanks to F-park, DimitriPapadopoulos, peterbygrave, shuuji3, eltoder, and homeworkprod) + +10.1.0 +------ + +* New functions + * :func:`takewhile_inclusive` (thanks to OlegAlexander) + * :func:`outer_product` (thanks to OlegAlexander) + +* Changes to existing functions + * :func:`zip_broadcast` was improved (thanks to kalekundert and pochmann) + * :func:`consume` had its type annotation fixed (thanks to obaltian) + +* Other changes + * Some documentation and testing issues were fixed (thanks to OlegAlexander) + +10.0.0 +------ + +* Potentially breaking changes + * Python 3.7 support was dropped, since it went EOL on 2023-06-27 + * :func:`batched` no longer issues a ``DeprecationWarning``; it is now an alias for ``itertools.batched`` for Python 3.12+ + * :func:`batched` and :func:`matmul` now yield tuples instead of lists + +* New functions + * :func:`combination_with_replacement_index` (thanks to Schoyen) + * :func:`nth_combination_with_replacement` (thanks to Schoyen) + * :func:`polynomial_eval` (from the Python itertools docs) + * :func:`polynomial_derivative` (from the Python itertools docs) + * :func:`sum_of_squares` (from the Python itertools docs) + +* Changes to existing functions + * :func:`seekable` now has ``relative_seek`` method (thanks to karlb) + * :func:`chunked_even` was optimized (thanks to elliotwutingfeng) + * :func:`numeric_range` was optimized (thanks to eltoder) + * :func:`duplicates_justseen`, :func:`pairwise`, :func:`partial_product`, and :func:`partition` were updated and optimized (thanks to pochmann) + * :func:`unique_in_window` had its implementation updated (thanks to elliotwutingfeng) + * :func:`iterate` now breaks when its ``func`` argument raises ``StopIteration`` (thanks to jrebiffe) + +* Other changes + * Some documentation and testing issues were fixed (thanks to lonnen and XuehaiPan) + +9.1.0 +----- + +* New functions + * :func:`iter_index` (from the Python itertools docs) + * :func:`transpose` (from the Python itertools docs) + * :func:`matmul` (from the Python itertools docs) + * :func:`factor` (from the Python itertools docs) + * :func:`gray_product` (thanks to haukex) + * :func:`partial_product` (thanks to lonnen) + +* Changes to existing functions + * :func:`sieve` was updated to match the Python itertools docs + * :func:`maxsplit` was updated to fix a bug (thanks to abingham) + * :func:`sliced` had its `type hint `__ updated (thanks to ad-chaos) + + +* Other changes + * The ``batched`` function is marked as deprecated and will be removed in a future major release. For Python 3.12 and above, use ``itertools.batched`` instead. (thanks to neutrinoceros) + * The type hints now used postponed evaluation of annotations from PEP 563 (thanks to Isira-Seneviratne) + * Some documentation issues were fixed (thanks to Voskov and jdkandersson) + +9.0.0 +----- + +* Potentially breaking changes + * :func:`grouper` no longer accepts an integer as its first argument. Previously this raised a ``DeprecationWarning``. + * :func:`collate` has been removed. Use the built-in :func:`heapq.merge` instead. + * :func:`windowed` now yields nothing when its iterable is empty. + * This library now advertises support for Python 3.7+. + +* New functions + * :func:`constrained_batches` + * :func:`batched` (from the Python itertools docs) + * :func:`polynomial_from_roots` (from the Python itertools docs) + * :func:`sieve` (from the Python itertools docs) + +* Other changes + * Some documentation issues were fixed (thanks to nanouasyn) + +8.14.0 +------ + +* New functions + * :func:`longest_common_prefix` (thanks to nanouasyn) + * :func:`iequals` (thanks to nanouasyn) + +* Changes to existing functions + * `concurrent.futures.ThreadPoolExecutor` is now imported lazily in :func:`callback_iter`. + * :func:`tail` is now optimized for iterables with a fixed length. + +* Other changes + * Some documentation issues were fixed (thanks to pochmann and timgates42) + * This library is now marked for Python 3.10 compatibility in PyPI (thanks to chayim) + +8.13.0 +------ + +* New functions + * The :func:`subslices` recipe from the `itertools` docs was added (thanks to rhettinger) + +* Changes to existing functions + * The :func:`ichunked` function is now more efficient (thanks to hjtran0 and seanmacavaney) + * The :func:`difference` function is now more efficient (thanks to Masynchin) + * The :func:`grouper` recipe now has more features, mirroring the one in the `itertools` docs (thanks to rhettinger) + +* Other changes + * Some documentation issues were fixed (thanks to medvied and Freed-Wu) + * The `more_itertools` package is now built with `flit` (thanks to mgorny) + +8.12.0 +------ + +* Bug fixes + * Some documentation issues were fixed (thanks to Masynchin, spookylukey, astrojuanlu, and stephengmatthews) + * Python 3.5 support was temporarily restored (thanks to mattbonnell) + +8.11.0 +------ + +* New functions + * The :func:`before_and_after`, :func:`sliding_window`, and :func:`triplewise` recipes from the Python 3.10 docs were added + * :func:`duplicates_everseen` and :func:`duplicates_justseen` (thanks to OrBin and DavidPratt512) + * :func:`minmax` (thanks to Ricocotam, MSeifert04, and ruancomelli) + * :func:`strictly_n` (thanks to hwalinga and NotWearingPants) + * :func:`unique_in_window` + +* Changes to existing functions + * :func:`groupby_transform` had its type stub improved (thanks to mjk4 and ruancomelli) + * :func:`is_sorted` now accepts a ``strict`` parameter (thanks to Dutcho and ruancomelli) + * :func:`zip_broadcast` was updated to fix a bug (thanks to kalekundert) + +8.10.0 +------ + +* Changes to existing functions + * The type stub for :func:`iter_except` was improved (thanks to MarcinKonowalczyk) + +* Other changes: + * Type stubs now ship with the source release (thanks to saaketp) + * The Sphinx docs were improved (thanks to MarcinKonowalczyk) + +8.9.0 +----- + +* New functions + * :func:`interleave_evenly` (thanks to mbugert) + * :func:`repeat_each` (thanks to FinalSh4re) + * :func:`chunked_even` (thanks to valtron) + * :func:`map_if` (thanks to sassbalint) + * :func:`zip_broadcast` (thanks to kalekundert) + +* Changes to existing functions + * The type stub for :func:`chunked` was improved (thanks to PhilMacKay) + * The type stubs for :func:`zip_equal` and `zip_offset` were improved (thanks to maffoo) + * Building Sphinx docs locally was improved (thanks to MarcinKonowalczyk) + +8.8.0 +----- + +* New functions + * :func:`countable` (thanks to krzysieq) + +* Changes to existing functions + * :func:`split_before` was updated to handle empty collections (thanks to TiunovNN) + * :func:`unique_everseen` got a performance boost (thanks to Numerlor) + * The type hint for :func:`value_chain` was corrected (thanks to vr2262) + +8.7.0 +----- + +* New functions + * :func:`convolve` (from the Python itertools docs) + * :func:`product_index`, :func:`combination_index`, and :func:`permutation_index` (thanks to N8Brooks) + * :func:`value_chain` (thanks to jenstroeger) + +* Changes to existing functions + * :func:`distinct_combinations` now uses a non-recursive algorithm (thanks to knutdrand) + * :func:`pad_none` is now the preferred name for :func:`padnone`, though the latter remains available. + * :func:`pairwise` will now use the Python standard library implementation on Python 3.10+ + * :func:`sort_together` now accepts a ``key`` argument (thanks to brianmaissy) + * :func:`seekable` now has a ``peek`` method, and can indicate whether the iterator it's wrapping is exhausted (thanks to gsakkis) + * :func:`time_limited` can now indicate whether its iterator has expired (thanks to roysmith) + * The implementation of :func:`unique_everseen` was improved (thanks to plammens) + +* Other changes: + * Various documentation updates (thanks to cthoyt, Evantm, and cyphase) + +8.6.0 +----- + +* New itertools + * :func:`all_unique` (thanks to brianmaissy) + * :func:`nth_product` and :func:`nth_permutation` (thanks to N8Brooks) + +* Changes to existing itertools + * :func:`chunked` and :func:`sliced` now accept a ``strict`` parameter (thanks to shlomif and jtwool) + +* Other changes + * Python 3.5 has reached its end of life and is no longer supported. + * Python 3.9 is officially supported. + * Various documentation fixes (thanks to timgates42) + +8.5.0 +----- + +* New itertools + * :func:`windowed_complete` (thanks to MarcinKonowalczyk) + +* Changes to existing itertools: + * The :func:`is_sorted` implementation was improved (thanks to cool-RR) + * The :func:`groupby_transform` now accepts a ``reducefunc`` parameter. + * The :func:`last` implementation was improved (thanks to brianmaissy) + +* Other changes + * Various documentation fixes (thanks to craigrosie, samuelstjean, PiCT0) + * The tests for :func:`distinct_combinations` were improved (thanks to Minabsapi) + * Automated tests now run on GitHub Actions. All commits now check: + * That unit tests pass + * That the examples in docstrings work + * That test coverage remains high (using `coverage`) + * For linting errors (using `flake8`) + * For consistent style (using `black`) + * That the type stubs work (using `mypy`) + * That the docs build correctly (using `sphinx`) + * That packages build correctly (using `twine`) + +8.4.0 +----- + +* New itertools + * :func:`mark_ends` (thanks to kalekundert) + * :func:`is_sorted` + +* Changes to existing itertools: + * :func:`islice_extended` can now be used with real slices (thanks to cool-RR) + * The implementations for :func:`filter_except` and :func:`map_except` were improved (thanks to SergBobrovsky) + +* Other changes + * Automated tests now enforce code style (using `black `__) + * The various signatures of :func:`islice_extended` and :func:`numeric_range` now appear in the docs (thanks to dsfulf) + * The test configuration for mypy was updated (thanks to blueyed) + + +8.3.0 +----- + +* New itertools + * :func:`zip_equal` (thanks to frankier and alexmojaki) + +* Changes to existing itertools: + * :func:`split_at`, :func:`split_before`, :func:`split_after`, and :func:`split_when` all got a ``maxsplit`` parameter (thanks to jferard and ilai-deutel) + * :func:`split_at` now accepts a ``keep_separator`` parameter (thanks to jferard) + * :func:`distinct_permutations` can now generate ``r``-length permutations (thanks to SergBobrovsky and ilai-deutel) + * The :func:`windowed` implementation was improved (thanks to SergBobrovsky) + * The :func:`spy` implementation was improved (thanks to has2k1) + +* Other changes + * Type stubs are now tested with ``stubtest`` (thanks to ilai-deutel) + * Tests now run with ``python -m unittest`` instead of ``python setup.py test`` (thanks to jdufresne) + +8.2.0 +----- + +* Bug fixes + * The .pyi files for typing were updated. (thanks to blueyed and ilai-deutel) + +* Changes to existing itertools: + * :func:`numeric_range` now behaves more like the built-in :func:`range`. (thanks to jferard) + * :func:`bucket` now allows for enumerating keys. (thanks to alexchandel) + * :func:`sliced` should work for numpy arrays. (thanks to sswingle) + * :func:`seekable` now has a ``maxlen`` parameter. + +8.1.0 +----- + +* Bug fixes + * :func:`partition` works with ``pred=None`` again. (thanks to MSeifert04) + +* New itertools + * :func:`sample` (thanks to tommyod) + * :func:`nth_or_last` (thanks to d-ryzhikov) + +* Changes to existing itertools: + * The implementation for :func:`divide` was improved. (thanks to jferard) + +8.0.2 +----- + +* Bug fixes + * The type stub files are now part of the wheel distribution (thanks to keisheiled) + +8.0.1 +----- + +* Bug fixes + * The type stub files now work for functions imported from the + root package (thanks to keisheiled) + +8.0.0 +----- + +* New itertools and other additions + * This library now ships type hints for use with mypy. + (thanks to ilai-deutel for the implementation, and to gabbard and fmagin for assistance) + * :func:`split_when` (thanks to jferard) + * :func:`repeat_last` (thanks to d-ryzhikov) + +* Changes to existing itertools: + * The implementation for :func:`set_partitions` was improved. (thanks to jferard) + * :func:`partition` was optimized for expensive predicates. (thanks to stevecj) + * :func:`unique_everseen` and :func:`groupby_transform` were re-factored. (thanks to SergBobrovsky) + * The implementation for :func:`difference` was improved. (thanks to Jabbey92) + +* Other changes + * Python 3.4 has reached its end of life and is no longer supported. + * Python 3.8 is officially supported. (thanks to jdufresne) + * The ``collate`` function has been deprecated. + It raises a ``DeprecationWarning`` if used, and will be removed in a future release. + * :func:`one` and :func:`only` now provide more informative error messages. (thanks to gabbard) + * Unit tests were moved outside of the main package (thanks to jdufresne) + * Various documentation fixes (thanks to kriomant, gabbard, jdufresne) + + +7.2.0 +----- + +* New itertools + * :func:`distinct_combinations` + * :func:`set_partitions` (thanks to kbarrett) + * :func:`filter_except` + * :func:`map_except` + +7.1.0 +----- + +* New itertools + * :func:`ichunked` (thanks davebelais and youtux) + * :func:`only` (thanks jaraco) + +* Changes to existing itertools: + * :func:`numeric_range` now supports ranges specified by + ``datetime.datetime`` and ``datetime.timedelta`` objects (thanks to MSeifert04 for tests). + * :func:`difference` now supports an *initial* keyword argument. + + +* Other changes + * Various documentation fixes (thanks raimon49, pylang) + +7.0.0 +----- + +* New itertools: + * :func:`time_limited` + * :func:`partitions` (thanks to rominf and Saluev) + * :func:`substrings_indexes` (thanks to rominf) + +* Changes to existing itertools: + * :func:`collapse` now treats ``bytes`` objects the same as ``str`` objects. (thanks to Sweenpet) + +The major version update is due to the change in the default behavior of +:func:`collapse`. It now treats ``bytes`` objects the same as ``str`` objects. +This aligns its behavior with :func:`always_iterable`. + +.. code-block:: python + + >>> from more_itertools import collapse + >>> iterable = [[1, 2], b'345', [6]] + >>> print(list(collapse(iterable))) + [1, 2, b'345', 6] + +6.0.0 +----- + +* Major changes: + * Python 2.7 is no longer supported. The 5.0.0 release will be the last + version targeting Python 2.7. + * All future releases will target the active versions of Python 3. + As of 2019, those are Python 3.4 and above. + * The ``six`` library is no longer a dependency. + * The :func:`accumulate` function is no longer part of this library. You + may import a better version from the standard ``itertools`` module. + +* Changes to existing itertools: + * The order of the parameters in :func:`grouper` have changed to match + the latest recipe in the itertools documentation. Use of the old order + will be supported in this release, but emit a ``DeprecationWarning``. + The legacy behavior will be dropped in a future release. (thanks to jaraco) + * :func:`distinct_permutations` was improved (thanks to jferard - see also `permutations with unique values `_ at StackOverflow.) + * An unused parameter was removed from :func:`substrings`. (thanks to pylang) + +* Other changes: + * The docs for :func:`unique_everseen` were improved. (thanks to jferard and MSeifert04) + * Several Python 2-isms were removed. (thanks to jaraco, MSeifert04, and hugovk) + +5.0.0 +----- + +* New itertools: + * :func:`split_into` (thanks to rovyko) + * :func:`unzip` (thanks to bmintz) + * :func:`substrings` (thanks to pylang) + +* Changes to existing itertools: + * :func:`ilen` was optimized a bit (thanks to MSeifert04, achampion, and bmintz) + * :func:`first_true` now returns ``None`` by default. This is the reason for the major version bump - see below. (thanks to sk and OJFord) + +* Other changes: + * Some code for old Python versions was removed (thanks to hugovk) + * Some documentation mistakes were corrected (thanks to belm0 and hugovk) + * Tests now run properly on 32-bit versions of Python (thanks to Millak) + * Newer versions of CPython and PyPy are now tested against + +The major version update is due to the change in the default return value of +:func:`first_true`. It's now ``None``. + +.. code-block:: python + + >>> from more_itertools import first_true + >>> iterable = [0, '', False, [], ()] # All these are False + >>> answer = first_true(iterable) + >>> print(answer) + None + +4.3.0 +----- + +* New itertools: + * :func:`last` (thanks to tmshn) + * :func:`replace` (thanks to pylang) + * :func:`rlocate` (thanks to jferard and pylang) + +* Improvements to existing itertools: + * :func:`locate` can now search for multiple items + +* Other changes: + * The docs now include a nice table of tools (thanks MSeifert04) + +4.2.0 +----- + +* New itertools: + * :func:`map_reduce` (thanks to pylang) + * :func:`prepend` (from the `Python 3.7 docs `_) + +* Improvements to existing itertools: + * :func:`bucket` now complies with PEP 479 (thanks to irmen) + +* Other changes: + * Python 3.7 is now supported (thanks to irmen) + * Python 3.3 is no longer supported + * The test suite no longer requires third-party modules to run + * The API docs now include links to source code + +4.1.0 +----- + +* New itertools: + * :func:`split_at` (thanks to michael-celani) + * :func:`circular_shifts` (thanks to hiqua) + * :func:`make_decorator` - see the blog post `Yo, I heard you like decorators `_ + for a tour (thanks to pylang) + * :func:`always_reversible` (thanks to michael-celani) + * :func:`nth_combination` (from the `Python 3.7 docs `_) + +* Improvements to existing itertools: + * :func:`seekable` now has an ``elements`` method to return cached items. + * The performance tradeoffs between :func:`roundrobin` and + :func:`interleave_longest` are now documented (thanks michael-celani, + pylang, and MSeifert04) + +4.0.1 +----- + +* No code changes - this release fixes how the docs display on PyPI. + +4.0.0 +----- + +* New itertools: + * :func:`consecutive_groups` (Based on the example in the `Python 2.4 docs `_) + * :func:`seekable` (If you're looking for how to "reset" an iterator, + you're in luck!) + * :func:`exactly_n` (thanks to michael-celani) + * :func:`run_length.encode` and :func:`run_length.decode` + * :func:`difference` + +* Improvements to existing itertools: + * The number of items between filler elements in :func:`intersperse` can + now be specified (thanks to pylang) + * :func:`distinct_permutations` and :func:`peekable` got some minor + adjustments (thanks to MSeifert04) + * :func:`always_iterable` now returns an iterator object. It also now + allows different types to be considered iterable (thanks to jaraco) + * :func:`bucket` can now limit the keys it stores in memory + * :func:`one` now allows for custom exceptions (thanks to kalekundert) + +* Other changes: + * A few typos were fixed (thanks to EdwardBetts) + * All tests can now be run with ``python setup.py test`` + +The major version update is due to the change in the return value of :func:`always_iterable`. +It now always returns iterator objects: + +.. code-block:: python + + >>> from more_itertools import always_iterable + # Non-iterable objects are wrapped with iter(tuple(obj)) + >>> always_iterable(12345) + + >>> list(always_iterable(12345)) + [12345] + # Iterable objects are wrapped with iter() + >>> always_iterable([1, 2, 3, 4, 5]) + + +3.2.0 +----- + +* New itertools: + * :func:`lstrip`, :func:`rstrip`, and :func:`strip` + (thanks to MSeifert04 and pylang) + * :func:`islice_extended` +* Improvements to existing itertools: + * Some bugs with slicing :func:`peekable`-wrapped iterables were fixed + +3.1.0 +----- + +* New itertools: + * :func:`numeric_range` (Thanks to BebeSparkelSparkel and MSeifert04) + * :func:`count_cycle` (Thanks to BebeSparkelSparkel) + * :func:`locate` (Thanks to pylang and MSeifert04) +* Improvements to existing itertools: + * A few itertools are now slightly faster due to some function + optimizations. (Thanks to MSeifert04) +* The docs have been substantially revised with installation notes, + categories for library functions, links, and more. (Thanks to pylang) + + +3.0.0 +----- + +* Removed itertools: + * ``context`` has been removed due to a design flaw - see below for + replacement options. (thanks to NeilGirdhar) +* Improvements to existing itertools: + * ``side_effect`` now supports ``before`` and ``after`` keyword + arguments. (Thanks to yardsale8) +* PyPy and PyPy3 are now supported. + +The major version change is due to the removal of the ``context`` function. +Replace it with standard ``with`` statement context management: + +.. code-block:: python + + # Don't use context() anymore + file_obj = StringIO() + consume(print(x, file=f) for f in context(file_obj) for x in u'123') + + # Use a with statement instead + file_obj = StringIO() + with file_obj as f: + consume(print(x, file=f) for x in u'123') + +2.6.0 +----- + +* New itertools: + * ``adjacent`` and ``groupby_transform`` (Thanks to diazona) + * ``always_iterable`` (Thanks to jaraco) + * (Removed in 3.0.0) ``context`` (Thanks to yardsale8) + * ``divide`` (Thanks to mozbhearsum) +* Improvements to existing itertools: + * ``ilen`` is now slightly faster. (Thanks to wbolster) + * ``peekable`` can now prepend items to an iterable. (Thanks to diazona) + +2.5.0 +----- + +* New itertools: + * ``distribute`` (Thanks to mozbhearsum and coady) + * ``sort_together`` (Thanks to clintval) + * ``stagger`` and ``zip_offset`` (Thanks to joshbode) + * ``padded`` +* Improvements to existing itertools: + * ``peekable`` now handles negative indexes and slices with negative + components properly. + * ``intersperse`` is now slightly faster. (Thanks to pylang) + * ``windowed`` now accepts a ``step`` keyword argument. + (Thanks to pylang) +* Python 3.6 is now supported. + +2.4.1 +----- + +* Move docs 100% to readthedocs.io. + +2.4 +----- + +* New itertools: + * ``accumulate``, ``all_equal``, ``first_true``, ``partition``, and + ``tail`` from the itertools documentation. + * ``bucket`` (Thanks to Rosuav and cvrebert) + * ``collapse`` (Thanks to abarnet) + * ``interleave`` and ``interleave_longest`` (Thanks to abarnet) + * ``side_effect`` (Thanks to nvie) + * ``sliced`` (Thanks to j4mie and coady) + * ``split_before`` and ``split_after`` (Thanks to astronouth7303) + * ``spy`` (Thanks to themiurgo and mathieulongtin) +* Improvements to existing itertools: + * ``chunked`` is now simpler and more friendly to garbage collection. + (Contributed by coady, with thanks to piskvorky) + * ``collate`` now delegates to ``heapq.merge`` when possible. + (Thanks to kmike and julianpistorius) + * ``peekable``-wrapped iterables are now indexable and sliceable. + Iterating through ``peekable``-wrapped iterables is also faster. + * ``one`` and ``unique_to_each`` have been simplified. + (Thanks to coady) + + +2.3 +----- + +* Added ``one`` from ``jaraco.util.itertools``. (Thanks, jaraco!) +* Added ``distinct_permutations`` and ``unique_to_each``. (Contributed by + bbayles) +* Added ``windowed``. (Contributed by bbayles, with thanks to buchanae, + jaraco, and abarnert) +* Simplified the implementation of ``chunked``. (Thanks, nvie!) +* Python 3.5 is now supported. Python 2.6 is no longer supported. +* Python 3 is now supported directly; there is no 2to3 step. + +2.2 +----- + +* Added ``iterate`` and ``with_iter``. (Thanks, abarnert!) + +2.1 +----- + +* Added (tested!) implementations of the recipes from the itertools + documentation. (Thanks, Chris Lonnen!) +* Added ``ilen``. (Thanks for the inspiration, Matt Basta!) + +2.0 +----- + +* ``chunked`` now returns lists rather than tuples. After all, they're + homogeneous. This slightly backward-incompatible change is the reason for + the major version bump. +* Added ``@consumer``. +* Improved test machinery. + +1.1 +----- + +* Added ``first`` function. +* Added Python 3 support. +* Added a default arg to ``peekable.peek()``. +* Noted how to easily test whether a peekable iterator is exhausted. +* Rewrote documentation. + +1.0 +----- + +* Initial release, with ``collate``, ``peekable``, and ``chunked``. Could + really use better docs. diff --git a/more_itertools/more.py b/more_itertools/more.py new file mode 100644 index 0000000..9c56ab3 --- /dev/null +++ b/more_itertools/more.py @@ -0,0 +1,5644 @@ +__lazy_modules__ = frozenset({'queue', 'threading'}) + +import math +import types + +from collections import Counter, defaultdict, deque +from collections.abc import Sequence +from contextlib import suppress +from functools import cached_property, partial, wraps +from heapq import heapify, heapreplace +from itertools import ( + chain, + combinations, + compress, + count, + cycle, + dropwhile, + groupby, + islice, + permutations, + repeat, + starmap, + takewhile, + tee, + zip_longest, + product, +) +from math import comb, e, exp, floor, fsum, log, log1p, perm, tau +from math import ceil, prod +from queue import Empty, Queue +from random import random, randrange, shuffle, uniform +from operator import ( + attrgetter, + getitem, + is_not, + itemgetter, + lt, + neg, + sub, + gt, +) +from sys import maxsize +from time import monotonic +from threading import Lock + +from .recipes import ( + _marker, + consume, + first_true, + flatten, + is_prime, + nth, + powerset, + sieve, + take, + unique_everseen, + all_equal, + batched, +) + +__all__ = [ + 'AbortThread', + 'SequenceView', + 'adjacent', + 'all_unique', + 'always_iterable', + 'always_reversible', + 'argmax', + 'argmin', + 'bucket', + 'callback_iter', + 'chunked', + 'chunked_even', + 'circular_shifts', + 'classify_unique', + 'collapse', + 'combination_index', + 'combination_with_replacement_index', + 'concurrent_tee', + 'consecutive_groups', + 'constrained_batches', + 'consumer', + 'count_cycle', + 'countable', + 'derangements', + 'dft', + 'difference', + 'distinct_combinations', + 'distinct_permutations', + 'distribute', + 'divide', + 'doublestarmap', + 'duplicates_everseen', + 'duplicates_justseen', + 'exactly_n', + 'extract', + 'filter_except', + 'filter_map', + 'first', + 'gray_product', + 'groupby_transform', + 'ichunked', + 'idft', + 'iequals', + 'ilen', + 'interleave', + 'interleave_evenly', + 'interleave_longest', + 'interleave_randomly', + 'intersperse', + 'is_sorted', + 'islice_extended', + 'iter_suppress', + 'iterate', + 'join_mappings', + 'last', + 'locate', + 'longest_common_prefix', + 'lstrip', + 'make_decorator', + 'map_except', + 'map_if', + 'map_reduce', + 'mark_ends', + 'minmax', + 'nth_combination_with_replacement', + 'nth_or_last', + 'nth_permutation', + 'nth_prime', + 'nth_product', + 'numeric_range', + 'one', + 'only', + 'outer_product', + 'padded', + 'partial_product', + 'partitions', + 'peekable', + 'permutation_index', + 'powerset_of_sets', + 'product_index', + 'raise_', + 'random_ordered_range', + 'repeat_each', + 'repeat_last', + 'replace', + 'rlocate', + 'rstrip', + 'run_length', + 'sample', + 'seekable', + 'serialize', + 'set_partitions', + 'side_effect', + 'sized_iterator', + 'sliced', + 'sort_together', + 'split_after', + 'split_at', + 'split_before', + 'split_into', + 'split_when', + 'spy', + 'stagger', + 'strictly_n', + 'strip', + 'subfactorial', + 'substrings', + 'substrings_indexes', + 'synchronized', + 'takewhile_inclusive', + 'time_limited', + 'unique_in_window', + 'unique_to_each', + 'unzip', + 'value_chain', + 'windowed', + 'windowed_complete', + 'with_iter', + 'zip_broadcast', + 'zip_offset', +] + +# math.sumprod is available for Python 3.12+ +try: + from math import sumprod as _fsumprod + +except ImportError: # pragma: no cover + # Extended precision algorithms from T. J. Dekker, + # "A Floating-Point Technique for Extending the Available Precision" + # https://csclub.uwaterloo.ca/~pbarfuss/dekker1971.pdf + # Formulas: (5.5) (5.6) and (5.8). Code: mul12() + + def dl_split(x: float): + "Split a float into two half-precision components." + t = x * 134217729.0 # Veltkamp constant = 2.0 ** 27 + 1 + hi = t - (t - x) + lo = x - hi + return hi, lo + + def dl_mul(x, y): + "Lossless multiplication." + xx_hi, xx_lo = dl_split(x) + yy_hi, yy_lo = dl_split(y) + p = xx_hi * yy_hi + q = xx_hi * yy_lo + xx_lo * yy_hi + z = p + q + zz = p - z + q + xx_lo * yy_lo + return z, zz + + def _fsumprod(p, q): + return fsum(chain.from_iterable(map(dl_mul, p, q))) + + +def chunked(iterable, n, strict=False): + """Break *iterable* into lists of length *n*: + + >>> list(chunked([1, 2, 3, 4, 5, 6], 3)) + [[1, 2, 3], [4, 5, 6]] + + By the default, the last yielded list will have fewer than *n* elements + if the length of *iterable* is not divisible by *n*: + + >>> list(chunked([1, 2, 3, 4, 5, 6, 7, 8], 3)) + [[1, 2, 3], [4, 5, 6], [7, 8]] + + To use a fill-in value instead, see the :func:`grouper` recipe. + + If the length of *iterable* is not divisible by *n* and *strict* is + ``True``, then ``ValueError`` will be raised before the last + list is yielded. + + """ + if n is not None and n < 0: + raise ValueError('n must be at least 0') + + iterator = iter(partial(take, n, iter(iterable)), []) + if strict: + if n is None: + raise ValueError('n must not be None when using strict mode.') + + def ret(): + for chunk in iterator: + if len(chunk) != n: + raise ValueError('iterable is not divisible by n.') + yield chunk + + return ret() + else: + return iterator + + +def first(iterable, default=_marker): + """Return the first item of *iterable*, or *default* if *iterable* is + empty. + + >>> first([0, 1, 2, 3]) + 0 + >>> first([], 'some default') + 'some default' + + If *default* is not provided and there are no items in the iterable, + raise ``ValueError``. + + :func:`first` is useful when you have a generator of expensive-to-retrieve + values and want any arbitrary one. It is marginally shorter than + ``next(iter(iterable), default)``. + + """ + for item in iterable: + return item + if default is _marker: + raise ValueError( + 'first() was called on an empty iterable, ' + 'and no default value was provided.' + ) + return default + + +def last(iterable, default=_marker): + """Return the last item of *iterable*, or *default* if *iterable* is + empty. + + >>> last([0, 1, 2, 3]) + 3 + >>> last([], 'some default') + 'some default' + + If *default* is not provided and there are no items in the iterable, + raise ``ValueError``. + """ + try: + if getattr(iterable, '__reversed__', None): + return next(reversed(iterable)) + return deque(iterable, maxlen=1)[-1] + except (IndexError, StopIteration): + if default is _marker: + raise ValueError( + 'last() was called on an empty iterable, ' + 'and no default value was provided.' + ) + return default + + +def nth_or_last(iterable, n, default=_marker): + """Return the nth or the last item of *iterable*, + or *default* if *iterable* is empty. + + >>> nth_or_last([0, 1, 2, 3], 2) + 2 + >>> nth_or_last([0, 1], 2) + 1 + >>> nth_or_last([], 0, 'some default') + 'some default' + + If *default* is not provided and there are no items in the iterable, + raise ``ValueError``. + """ + return last(islice(iterable, n + 1), default=default) + + +class peekable: + """Wrap an iterator to allow lookahead and prepending elements. + + Call :meth:`peek` on the result to get the value that will be returned + by :func:`next`. This won't advance the iterator: + + >>> p = peekable(['a', 'b']) + >>> p.peek() + 'a' + >>> next(p) + 'a' + + Pass :meth:`peek` a default value to return that instead of raising + ``StopIteration`` when the iterator is exhausted. + + >>> p = peekable([]) + >>> p.peek('hi') + 'hi' + + peekables also offer a :meth:`prepend` method, which "inserts" items + at the head of the iterable: + + >>> p = peekable([1, 2, 3]) + >>> p.prepend(10, 11, 12) + >>> next(p) + 10 + >>> p.peek() + 11 + >>> list(p) + [11, 12, 1, 2, 3] + + peekables can be indexed. Index 0 is the item that will be returned by + :func:`next`, index 1 is the item after that, and so on: + The values up to the given index will be cached. + + >>> p = peekable(['a', 'b', 'c', 'd']) + >>> p[0] + 'a' + >>> p[1] + 'b' + >>> next(p) + 'a' + + Negative indexes are supported, but be aware that they will cache the + remaining items in the source iterator, which may require significant + storage. + + To check whether a peekable is exhausted, check its truth value: + + >>> p = peekable(['a', 'b']) + >>> if p: # peekable has items + ... list(p) + ['a', 'b'] + >>> if not p: # peekable is exhausted + ... list(p) + [] + + """ + + def __init__(self, iterable): + self._it = iter(iterable) + self._cache = deque() + + def __iter__(self): + return self + + def __bool__(self): + try: + self.peek() + except StopIteration: + return False + return True + + def peek(self, default=_marker): + """Return the item that will be next returned from ``next()``. + + Return ``default`` if there are no items left. If ``default`` is not + provided, raise ``StopIteration``. + + """ + if not self._cache: + try: + self._cache.append(next(self._it)) + except StopIteration: + if default is _marker: + raise + return default + return self._cache[0] + + def prepend(self, *items): + """Stack up items to be the next ones returned from ``next()`` or + ``self.peek()``. The items will be returned in + first in, first out order:: + + >>> p = peekable([1, 2, 3]) + >>> p.prepend(10, 11, 12) + >>> next(p) + 10 + >>> list(p) + [11, 12, 1, 2, 3] + + It is possible, by prepending items, to "resurrect" a peekable that + previously raised ``StopIteration``. + + >>> p = peekable([]) + >>> next(p) + Traceback (most recent call last): + ... + StopIteration + >>> p.prepend(1) + >>> next(p) + 1 + >>> next(p) + Traceback (most recent call last): + ... + StopIteration + + """ + self._cache.extendleft(reversed(items)) + + __class_getitem__ = classmethod(types.GenericAlias) + + def __next__(self): + if self._cache: + return self._cache.popleft() + + return next(self._it) + + def _get_slice(self, index): + # Normalize the slice's arguments + step = 1 if (index.step is None) else index.step + if step > 0: + start = 0 if (index.start is None) else index.start + stop = maxsize if (index.stop is None) else index.stop + elif step < 0: + start = -1 if (index.start is None) else index.start + stop = (-maxsize - 1) if (index.stop is None) else index.stop + else: + raise ValueError('slice step cannot be zero') + + # If either the start or stop index is negative, we'll need to cache + # the rest of the iterable in order to slice from the right side. + if (start < 0) or (stop < 0): + self._cache.extend(self._it) + # Otherwise we'll need to find the rightmost index and cache to that + # point. + else: + n = min(max(start, stop) + 1, maxsize) + cache_len = len(self._cache) + if n >= cache_len: + self._cache.extend(islice(self._it, n - cache_len)) + + return list(self._cache)[index] + + def __getitem__(self, index): + if isinstance(index, slice): + return self._get_slice(index) + + cache_len = len(self._cache) + if index < 0: + self._cache.extend(self._it) + elif index >= cache_len: + self._cache.extend(islice(self._it, index + 1 - cache_len)) + + return self._cache[index] + + +def consumer(func): + """Decorator that automatically advances a PEP-342-style "reverse iterator" + to its first yield point so you don't have to call ``next()`` on it + manually. + + >>> @consumer + ... def tally(): + ... i = 0 + ... while True: + ... print('Thing number %s is %s.' % (i, (yield))) + ... i += 1 + ... + >>> t = tally() + >>> t.send('red') + Thing number 0 is red. + >>> t.send('fish') + Thing number 1 is fish. + + Without the decorator, you would have to call ``next(t)`` before + ``t.send()`` could be used. + + """ + + @wraps(func) + def wrapper(*args, **kwargs): + gen = func(*args, **kwargs) + next(gen) + return gen + + return wrapper + + +def ilen(iterable): + """Return the number of items in *iterable*. + + For example, there are 168 prime numbers below 1,000: + + >>> ilen(sieve(1000)) + 168 + + Equivalent to, but faster than:: + + def ilen(iterable): + count = 0 + for _ in iterable: + count += 1 + return count + + This fully consumes the iterable, so handle with care. + + """ + # This is the "most beautiful of the fast variants" of this function. + # If you think you can improve on it, please ensure that your version + # is both 10x faster and 10x more beautiful. + return sum(compress(repeat(1), zip(iterable))) + + +def iterate(func, start): + """Return ``start``, ``func(start)``, ``func(func(start))``, ... + + Produces an infinite iterator. To add a stopping condition, + use :func:`take`, ``takewhile``, or :func:`takewhile_inclusive`:. + + >>> take(10, iterate(lambda x: 2*x, 1)) + [1, 2, 4, 8, 16, 32, 64, 128, 256, 512] + + >>> collatz = lambda x: 3*x + 1 if x%2==1 else x // 2 + >>> list(takewhile_inclusive(lambda x: x!=1, iterate(collatz, 10))) + [10, 5, 16, 8, 4, 2, 1] + + """ + with suppress(StopIteration): + while True: + yield start + start = func(start) + + +def with_iter(context_manager): + """Wrap an iterable in a ``with`` statement, so it closes once exhausted. + + For example, this will close the file when the iterator is exhausted:: + + upper_lines = (line.upper() for line in with_iter(open('foo'))) + + Note that you have to actually exhaust the iterator for opened files to be closed. + + Any context manager which returns an iterable is a candidate for + ``with_iter``. + + """ + with context_manager as iterable: + yield from iterable + + +class sized_iterator: + """Wrapper for *iterable* that implements ``__len__``. + + >>> it = map(str, range(5)) + >>> sized_it = sized_iterator(it, 5) + >>> len(sized_it) + 5 + >>> list(sized_it) + ['0', '1', '2', '3', '4'] + + This is useful for tools that use :func:`len`, like + `tqdm `__ . + + The wrapper doesn't validate the provided *length*, so be sure to choose + a value that reflects reality. + """ + + def __init__(self, iterable, length): + self._iterator = iter(iterable) + self._length = length + + def __next__(self): + return next(self._iterator) + + def __iter__(self): + return self + + def __len__(self): + return self._length + + +def one(iterable, too_short=None, too_long=None): + """Return the first item from *iterable*, which is expected to contain only + that item. Raise an exception if *iterable* is empty or has more than one + item. + + :func:`one` is useful for ensuring that an iterable contains only one item. + For example, it can be used to retrieve the result of a database query + that is expected to return a single row. + + If *iterable* is empty, ``ValueError`` will be raised. You may specify a + different exception with the *too_short* keyword: + + >>> it = [] + >>> one(it) # doctest: +IGNORE_EXCEPTION_DETAIL + Traceback (most recent call last): + ... + ValueError: too few items in iterable (expected 1)' + >>> too_short = IndexError('too few items') + >>> one(it, too_short=too_short) # doctest: +IGNORE_EXCEPTION_DETAIL + Traceback (most recent call last): + ... + IndexError: too few items + + Similarly, if *iterable* contains more than one item, ``ValueError`` will + be raised. You may specify a different exception with the *too_long* + keyword: + + >>> it = ['too', 'many'] + >>> one(it) # doctest: +IGNORE_EXCEPTION_DETAIL + Traceback (most recent call last): + ... + ValueError: Expected exactly one item in iterable, but got 'too', + 'many', and perhaps more. + >>> too_long = RuntimeError + >>> one(it, too_long=too_long) # doctest: +IGNORE_EXCEPTION_DETAIL + Traceback (most recent call last): + ... + RuntimeError + + Note that :func:`one` attempts to advance *iterable* twice to ensure there + is only one item. See :func:`spy` or :func:`peekable` to check iterable + contents less destructively. + + """ + iterator = iter(iterable) + for first in iterator: + for second in iterator: + if too_long is not None: + raise too_long + raise ValueError( + f'Expected exactly one item in iterable, but got {first!r}, ' + f'{second!r}, and perhaps more.' + ) + return first + if too_short is not None: + raise too_short + raise ValueError('too few items in iterable (expected 1)') + + +def raise_(exception, *args): + raise exception(*args) + + +def strictly_n(iterable, n, too_short=None, too_long=None): + """Validate that *iterable* has exactly *n* items and return them if + it does. If it has fewer than *n* items, call function *too_short* + with the actual number of items. If it has more than *n* items, call function + *too_long* with the number ``n + 1``. + + >>> iterable = ['a', 'b', 'c', 'd'] + >>> n = 4 + >>> list(strictly_n(iterable, n)) + ['a', 'b', 'c', 'd'] + + Note that the returned iterable must be consumed in order for the check to + be made. + + By default, *too_short* and *too_long* are functions that raise + ``ValueError``. + + >>> list(strictly_n('ab', 3)) # doctest: +IGNORE_EXCEPTION_DETAIL + Traceback (most recent call last): + ... + ValueError: too few items in iterable (got 2) + + >>> list(strictly_n('abc', 2)) # doctest: +IGNORE_EXCEPTION_DETAIL + Traceback (most recent call last): + ... + ValueError: too many items in iterable (got at least 3) + + You can instead supply functions that do something else. + *too_short* will be called with the number of items in *iterable*. + *too_long* will be called with `n + 1`. + + >>> def too_short(item_count): + ... raise RuntimeError + >>> it = strictly_n('abcd', 6, too_short=too_short) + >>> list(it) # doctest: +IGNORE_EXCEPTION_DETAIL + Traceback (most recent call last): + ... + RuntimeError + + >>> def too_long(item_count): + ... print('The boss is going to hear about this') + >>> it = strictly_n('abcdef', 4, too_long=too_long) + >>> list(it) + The boss is going to hear about this + ['a', 'b', 'c', 'd'] + + """ + if too_short is None: + too_short = lambda item_count: raise_( + ValueError, + f'Too few items in iterable (got {item_count})', + ) + + if too_long is None: + too_long = lambda item_count: raise_( + ValueError, + f'Too many items in iterable (got at least {item_count})', + ) + + it = iter(iterable) + + sent = 0 + for item in islice(it, n): + yield item + sent += 1 + + if sent < n: + too_short(sent) + return + + for item in it: + too_long(n + 1) + return + + +def distinct_permutations(iterable, r=None): + """Yield successive distinct permutations of the elements in *iterable*. + + >>> sorted(distinct_permutations([1, 0, 1])) + [(0, 1, 1), (1, 0, 1), (1, 1, 0)] + + Equivalent to yielding from ``set(permutations(iterable))``, except + duplicates are not generated and thrown away. For larger input sequences + this is much more efficient. + + If the elements of the input iterable are sortable, the output tuples are + produced in sorted order. + + Duplicate permutations arise when there are duplicated elements in the + input iterable. The number of items returned is + `n! / (x_1! * x_2! * ... * x_n!)`, where `n` is the total number of + items input, and each `x_i` is the count of a distinct item in the input + sequence. The function :func:`multinomial` computes this directly. + + If *r* is given, only the *r*-length permutations are yielded. + + >>> sorted(distinct_permutations([1, 0, 1], r=2)) + [(0, 1), (1, 0), (1, 1)] + >>> sorted(distinct_permutations(range(3), r=2)) + [(0, 1), (0, 2), (1, 0), (1, 2), (2, 0), (2, 1)] + + *iterable* need not be sortable, but note that using equal (``x == y``) + but non-identical (``id(x) != id(y)``) elements may produce surprising + behavior. For example, ``1`` and ``True`` are equal but non-identical: + + >>> list(distinct_permutations([1, True, '3'])) # doctest: +SKIP + [ + (1, True, '3'), + (1, '3', True), + ('3', 1, True) + ] + >>> list(distinct_permutations([1, 2, '3'])) # doctest: +SKIP + [ + (1, 2, '3'), + (1, '3', 2), + (2, 1, '3'), + (2, '3', 1), + ('3', 1, 2), + ('3', 2, 1) + ] + """ + + # Algorithm: https://w.wiki/Qai + def _full(A): + while True: + # Yield the permutation we have + yield tuple(A) + + # Find the largest index i such that A[i] < A[i + 1] + for i in range(size - 2, -1, -1): + if A[i] < A[i + 1]: + break + # If no such index exists, this permutation is the last one + else: + return + + # Find the largest index j greater than j such that A[i] < A[j] + for j in range(size - 1, i, -1): + if A[i] < A[j]: + break + + # Swap the value of A[i] with that of A[j], then reverse the + # sequence from A[i + 1] to form the new permutation + A[i], A[j] = A[j], A[i] + A[i + 1 :] = A[: i - size : -1] # A[i + 1:][::-1] + + # Algorithm: modified from the above + def _partial(A, r): + # Split A into the first r items and the last r items + head, tail = A[:r], A[r:] + right_head_indexes = range(r - 1, -1, -1) + left_tail_indexes = range(len(tail)) + + while True: + # Yield the permutation we have + yield tuple(head) + + # Starting from the right, find the first index of the head with + # value smaller than the maximum value of the tail - call it i. + pivot = tail[-1] + for i in right_head_indexes: + if head[i] < pivot: + break + pivot = head[i] + else: + return + + # Starting from the left, find the first value of the tail + # with a value greater than head[i] and swap. + for j in left_tail_indexes: + if tail[j] > head[i]: + head[i], tail[j] = tail[j], head[i] + break + # If we didn't find one, start from the right and find the first + # index of the head with a value greater than head[i] and swap. + else: + for j in right_head_indexes: + if head[j] > head[i]: + head[i], head[j] = head[j], head[i] + break + + # Reverse head[i + 1:] and swap it with tail[:r - (i + 1)] + tail += head[: i - r : -1] # head[i + 1:][::-1] + i += 1 + head[i:], tail[:] = tail[: r - i], tail[r - i :] + + items = list(iterable) + + try: + items.sort() + sortable = True + except TypeError: + sortable = False + + indices_dict = defaultdict(list) + + for item in items: + indices_dict[items.index(item)].append(item) + + indices = [items.index(item) for item in items] + indices.sort() + + equivalent_items = {k: cycle(v) for k, v in indices_dict.items()} + + def permuted_items(permuted_indices): + return tuple( + next(equivalent_items[index]) for index in permuted_indices + ) + + size = len(items) + if r is None: + r = size + + # functools.partial(_partial, ... ) + algorithm = _full if (r == size) else partial(_partial, r=r) + + if 0 < r <= size: + if sortable: + return algorithm(items) + else: + return ( + permuted_items(permuted_indices) + for permuted_indices in algorithm(indices) + ) + + return iter(() if r else ((),)) + + +def derangements(iterable, r=None): + """Yield successive derangements of the elements in *iterable*. + + A derangement is a permutation in which no element appears at its original + index. In other words, a derangement is a permutation that has no fixed points. + + Suppose Alice, Bob, Carol, and Dave are playing Secret Santa. + The code below outputs all of the different ways to assign gift recipients + such that nobody is assigned to himself or herself: + + >>> for d in derangements(['Alice', 'Bob', 'Carol', 'Dave']): + ... print(', '.join(d)) + Bob, Alice, Dave, Carol + Bob, Carol, Dave, Alice + Bob, Dave, Alice, Carol + Carol, Alice, Dave, Bob + Carol, Dave, Alice, Bob + Carol, Dave, Bob, Alice + Dave, Alice, Bob, Carol + Dave, Carol, Alice, Bob + Dave, Carol, Bob, Alice + + If *r* is given, only the *r*-length derangements are yielded. + + >>> sorted(derangements(range(3), 2)) + [(1, 0), (1, 2), (2, 0)] + >>> sorted(derangements([0, 2, 3], 2)) + [(2, 0), (2, 3), (3, 0)] + + Elements are treated as unique based on their position, not on their value. + + Consider the Secret Santa example with two *different* people who have + the *same* name. Then there are two valid gift assignments even though + it might appear that a person is assigned to themselves: + + >>> names = ['Alice', 'Bob', 'Bob'] + >>> list(derangements(names)) + [('Bob', 'Bob', 'Alice'), ('Bob', 'Alice', 'Bob')] + + To avoid confusion, make the inputs distinct: + + >>> deduped = [f'{name}{index}' for index, name in enumerate(names)] + >>> list(derangements(deduped)) + [('Bob1', 'Bob2', 'Alice0'), ('Bob2', 'Alice0', 'Bob1')] + + The number of derangements of a set of size *n* is known as the + "subfactorial of n". For n > 0, the subfactorial is: + ``round(math.factorial(n) / math.e)``. The more-itertools function + :func:`subfactorial` computes this directly. + + References: + + * Article: https://www.numberanalytics.com/blog/ultimate-guide-to-derangements-in-combinatorics + * Sizes: https://oeis.org/A000166 + """ + xs = tuple(iterable) + ys = tuple(range(len(xs))) + return compress( + permutations(xs, r=r), + map(all, map(map, repeat(is_not), repeat(ys), permutations(ys, r=r))), + ) + + +def intersperse(e, iterable, n=1): + """Intersperse filler element *e* among the items in *iterable*, leaving + *n* items between each filler element. + + >>> list(intersperse('!', [1, 2, 3, 4, 5])) + [1, '!', 2, '!', 3, '!', 4, '!', 5] + + >>> list(intersperse(None, [1, 2, 3, 4, 5], n=2)) + [1, 2, None, 3, 4, None, 5] + + """ + if n == 0: + raise ValueError('n must be > 0') + elif n == 1: + # interleave(repeat(e), iterable) -> e, x_0, e, x_1, e, x_2... + # islice(..., 1, None) -> x_0, e, x_1, e, x_2... + return islice(interleave(repeat(e), iterable), 1, None) + else: + # interleave(filler, chunks) -> [e], [x_0, x_1], [e], [x_2, x_3]... + # islice(..., 1, None) -> [x_0, x_1], [e], [x_2, x_3]... + # flatten(...) -> x_0, x_1, e, x_2, x_3... + filler = repeat([e]) + chunks = chunked(iterable, n) + return flatten(islice(interleave(filler, chunks), 1, None)) + + +def unique_to_each(*iterables): + """Return the elements from each of the input iterables that aren't in the + other input iterables. + + For example, suppose you have a set of packages, each with a set of + dependencies:: + + {'pkg_1': {'A', 'B'}, 'pkg_2': {'B', 'C'}, 'pkg_3': {'B', 'D'}} + + If you remove one package, which dependencies can also be removed? + + If ``pkg_1`` is removed, then ``A`` is no longer necessary - it is not + associated with ``pkg_2`` or ``pkg_3``. Similarly, ``C`` is only needed for + ``pkg_2``, and ``D`` is only needed for ``pkg_3``:: + + >>> unique_to_each({'A', 'B'}, {'B', 'C'}, {'B', 'D'}) + [['A'], ['C'], ['D']] + + If there are duplicates in one input iterable that aren't in the others + they will be duplicated in the output. Input order is preserved:: + + >>> unique_to_each("mississippi", "missouri") + [['p', 'p'], ['o', 'u', 'r']] + + It is assumed that the elements of each iterable are hashable. + + """ + pool = [list(it) for it in iterables] + counts = Counter(chain.from_iterable(map(set, pool))) + uniques = {element for element in counts if counts[element] == 1} + return [list(filter(uniques.__contains__, it)) for it in pool] + + +def windowed(seq, n, fillvalue=None, step=1): + """Return a sliding window of width *n* over the given iterable. + + >>> all_windows = windowed([1, 2, 3, 4, 5], 3) + >>> list(all_windows) + [(1, 2, 3), (2, 3, 4), (3, 4, 5)] + + When the window is larger than the iterable, *fillvalue* is used in place + of missing values: + + >>> list(windowed([1, 2, 3], 4)) + [(1, 2, 3, None)] + + Each window will advance in increments of *step*: + + >>> list(windowed([1, 2, 3, 4, 5, 6], 3, fillvalue='!', step=2)) + [(1, 2, 3), (3, 4, 5), (5, 6, '!')] + + To slide into the iterable's items, use :func:`chain` to add filler items + to the left: + + >>> iterable = [1, 2, 3, 4] + >>> n = 3 + >>> padding = [None] * (n - 1) + >>> list(windowed(chain(padding, iterable), 3)) + [(None, None, 1), (None, 1, 2), (1, 2, 3), (2, 3, 4)] + """ + if n <= 0: + raise ValueError('n must be > 0') + if step < 1: + raise ValueError('step must be >= 1') + + iterator = iter(seq) + + # Generate first window + window = deque(islice(iterator, n), maxlen=n) + + # Deal with the first window not being full + if not window: + return + if len(window) < n: + yield tuple(window) + ((fillvalue,) * (n - len(window))) + return + yield tuple(window) + + # Create the filler for the next windows. The padding ensures + # we have just enough elements to fill the last window. + padding = (fillvalue,) * (n - 1 if step >= n else step - 1) + filler = map(window.append, chain(iterator, padding)) + + # Generate the rest of the windows + for _ in islice(filler, step - 1, None, step): + yield tuple(window) + + +def substrings(iterable): + """Yield all of the substrings of *iterable*. + + >>> [''.join(s) for s in substrings('more')] + ['m', 'o', 'r', 'e', 'mo', 'or', 're', 'mor', 'ore', 'more'] + + Note that non-string iterables can also be subdivided. + + >>> list(substrings([0, 1, 2])) + [(0,), (1,), (2,), (0, 1), (1, 2), (0, 1, 2)] + + Like subslices() but returns tuples instead of lists + and returns the shortest substrings first. + + """ + seq = tuple(iterable) + item_count = len(seq) + for n in range(1, item_count + 1): + slices = map(slice, range(item_count), range(n, item_count + 1)) + yield from map(getitem, repeat(seq), slices) + + +def substrings_indexes(seq, reverse=False): + """Yield all substrings and their positions in *seq* + + The items yielded will be a tuple of the form ``(substr, i, j)``, where + ``substr == seq[i:j]``. + + This function only works for iterables that support slicing, such as + ``str`` objects. + + >>> for item in substrings_indexes('more'): + ... print(item) + ('m', 0, 1) + ('o', 1, 2) + ('r', 2, 3) + ('e', 3, 4) + ('mo', 0, 2) + ('or', 1, 3) + ('re', 2, 4) + ('mor', 0, 3) + ('ore', 1, 4) + ('more', 0, 4) + + Set *reverse* to ``True`` to yield the same items in the opposite order. + + + """ + r = range(1, len(seq) + 1) + if reverse: + r = reversed(r) + return ( + (seq[i : i + L], i, i + L) for L in r for i in range(len(seq) - L + 1) + ) + + +class bucket: + """Wrap *iterable* and return an object that buckets the iterable into + child iterables based on a *key* function. + + >>> iterable = ['a1', 'b1', 'c1', 'a2', 'b2', 'c2', 'b3'] + >>> s = bucket(iterable, key=lambda x: x[0]) # Bucket by 1st character + >>> sorted(list(s)) # Get the keys + ['a', 'b', 'c'] + >>> a_iterable = s['a'] + >>> next(a_iterable) + 'a1' + >>> next(a_iterable) + 'a2' + >>> list(s['b']) + ['b1', 'b2', 'b3'] + + The original iterable will be advanced and its items will be cached until + they are used by the child iterables. This may require significant storage. + + By default, attempting to select a bucket to which no items belong will + exhaust the iterable and cache all values. + If you specify a *validator* function, selected buckets will instead be + checked against it. + + >>> from itertools import count + >>> it = count(1, 2) # Infinite sequence of odd numbers + >>> key = lambda x: x % 10 # Bucket by last digit + >>> validator = lambda x: x in {1, 3, 5, 7, 9} # Odd digits only + >>> s = bucket(it, key=key, validator=validator) + >>> 2 in s + False + >>> list(s[2]) + [] + + .. seealso:: :func:`map_reduce`, :func:`groupby_transform` + + If storage is not a concern, :func:`map_reduce` returns a Python + dictionary, which is generally easier to work with. If the elements + with the same key are already adjacent, :func:`groupby_transform` + or :func:`itertools.groupby` can be used without any caching overhead. + + """ + + def __init__(self, iterable, key, validator=None): + self._it = iter(iterable) + self._key = key + self._cache = defaultdict(deque) + # Use ``is None`` rather than truthiness: a valid validator may be a + # callable object that is falsey (e.g. an empty ``set`` subclass with + # ``__call__``), and ``validator or ...`` would silently replace it. + self._validator = (lambda x: True) if validator is None else validator + + def __contains__(self, value): + if not self._validator(value): + return False + + try: + item = next(self[value]) + except StopIteration: + return False + else: + self._cache[value].appendleft(item) + + return True + + def _get_values(self, value): + """ + Helper to yield items from the parent iterator that match *value*. + Items that don't match are stored in the local cache as they + are encountered. + """ + while True: + # If we've cached some items that match the target value, emit + # the first one and evict it from the cache. + if self._cache.get(value): + yield self._cache[value].popleft() + # Otherwise we need to advance the parent iterator to search for + # a matching item, caching the rest. + else: + while True: + try: + item = next(self._it) + except StopIteration: + return + item_value = self._key(item) + if item_value == value: + if value not in self._cache: + self._cache[value] = deque() + yield item + break + elif self._validator(item_value): + self._cache[item_value].append(item) + + def __iter__(self): + for item in self._it: + item_value = self._key(item) + if self._validator(item_value): + self._cache[item_value].append(item) + + return iter(self._cache) + + def __getitem__(self, value): + if not self._validator(value): + return iter(()) + + return self._get_values(value) + + +def spy(iterable, n=1): + """Return a 2-tuple with a list containing the first *n* elements of + *iterable*, and an iterator with the same items as *iterable*. + This allows you to "look ahead" at the items in the iterable without + advancing it. + + There is one item in the list by default: + + >>> iterable = 'abcdefg' + >>> head, iterable = spy(iterable) + >>> head + ['a'] + >>> list(iterable) + ['a', 'b', 'c', 'd', 'e', 'f', 'g'] + + You may use unpacking to retrieve items instead of lists: + + >>> (head,), iterable = spy('abcdefg') + >>> head + 'a' + >>> (first, second), iterable = spy('abcdefg', 2) + >>> first + 'a' + >>> second + 'b' + + The number of items requested can be larger than the number of items in + the iterable: + + >>> iterable = [1, 2, 3, 4, 5] + >>> head, iterable = spy(iterable, 10) + >>> head + [1, 2, 3, 4, 5] + >>> list(iterable) + [1, 2, 3, 4, 5] + + """ + p, q = tee(iterable) + return take(n, q), p + + +def interleave(*iterables): + """Return a new iterable yielding from each iterable in turn, + until the shortest is exhausted. + + >>> list(interleave([1, 2, 3], [4, 5], [6, 7, 8])) + [1, 4, 6, 2, 5, 7] + + For a version that doesn't terminate after the shortest iterable is + exhausted, see :func:`interleave_longest`. + + """ + return chain.from_iterable(zip(*iterables)) + + +def interleave_longest(*iterables): + """Return a new iterable yielding from each iterable in turn, + skipping any that are exhausted. + + >>> list(interleave_longest([1, 2, 3], [4, 5], [6, 7, 8])) + [1, 4, 6, 2, 5, 7, 3, 8] + + This function produces the same output as :func:`roundrobin`, but may + perform better for some inputs (in particular when the number of iterables + is large). + + """ + for xs in zip_longest(*iterables, fillvalue=_marker): + for x in xs: + if x is not _marker: + yield x + + +def interleave_evenly(iterables, lengths=None): + """ + Interleave multiple iterables so that their elements are evenly distributed + throughout the output sequence. + + >>> iterables = [1, 2, 3, 4, 5], ['a', 'b'] + >>> list(interleave_evenly(iterables)) + [1, 2, 'a', 3, 4, 'b', 5] + + >>> iterables = [[1, 2, 3], [4, 5], [6, 7, 8]] + >>> list(interleave_evenly(iterables)) + [1, 6, 4, 2, 7, 3, 8, 5] + + This function requires iterables of known length. Iterables without + ``__len__()`` can be used by manually specifying lengths with *lengths*: + + >>> from itertools import combinations, repeat + >>> iterables = [combinations(range(4), 2), ['a', 'b', 'c']] + >>> lengths = [4 * (4 - 1) // 2, 3] + >>> list(interleave_evenly(iterables, lengths=lengths)) + [(0, 1), (0, 2), 'a', (0, 3), (1, 2), 'b', (1, 3), (2, 3), 'c'] + + Based on Bresenham's algorithm. + """ + if lengths is None: + try: + lengths = [len(it) for it in iterables] + except TypeError: + raise ValueError( + 'Iterable lengths could not be determined automatically. ' + 'Specify them with the lengths keyword.' + ) + elif len(iterables) != len(lengths): + raise ValueError('Mismatching number of iterables and lengths.') + + dims = len(lengths) + + if not dims: + return + + # sort iterables by length, descending + lengths_permute = sorted( + range(dims), key=lambda i: lengths[i], reverse=True + ) + lengths_desc = [lengths[i] for i in lengths_permute] + iters_desc = [iter(iterables[i]) for i in lengths_permute] + + # the longest iterable is the primary one (Bresenham: the longest + # distance along an axis) + delta_primary, deltas_secondary = lengths_desc[0], lengths_desc[1:] + iter_primary, iters_secondary = iters_desc[0], iters_desc[1:] + errors = [delta_primary // dims] * len(deltas_secondary) + + to_yield = sum(lengths) + while to_yield: + yield next(iter_primary) + to_yield -= 1 + # update errors for each secondary iterable + errors = [e - delta for e, delta in zip(errors, deltas_secondary)] + + # those iterables for which the error is negative are yielded + # ("diagonal step" in Bresenham) + for i, e_ in enumerate(errors): + if e_ < 0: + yield next(iters_secondary[i]) + to_yield -= 1 + errors[i] += delta_primary + + +def interleave_randomly(*iterables): + """Repeatedly select one of the input *iterables* at random and yield the next + item from it. + + >>> iterables = [1, 2, 3], 'abc', (True, False, None) + >>> list(interleave_randomly(*iterables)) # doctest: +SKIP + ['a', 'b', 1, 'c', True, False, None, 2, 3] + + The relative order of the items in each input iterable will preserved. Note the + sequences of items with this property are not equally likely to be generated. + + """ + iterators = [iter(e) for e in iterables] + while iterators: + idx = randrange(len(iterators)) + try: + yield next(iterators[idx]) + except StopIteration: + # equivalent to `list.pop` but slightly faster + iterators[idx] = iterators[-1] + del iterators[-1] + + +def collapse(iterable, base_type=None, levels=None): + """Flatten an iterable with multiple levels of nesting (e.g., a list of + lists of tuples) into non-iterable types. + + >>> iterable = [(1, 2), ([3, 4], [[5], [6]])] + >>> list(collapse(iterable)) + [1, 2, 3, 4, 5, 6] + + Binary and text strings are not considered iterable and + will not be collapsed. + + To avoid collapsing other types, specify *base_type*: + + >>> iterable = ['ab', ('cd', 'ef'), ['gh', 'ij']] + >>> list(collapse(iterable, base_type=tuple)) + ['ab', ('cd', 'ef'), 'gh', 'ij'] + + Specify *levels* to stop flattening after a certain level: + + >>> iterable = [('a', ['b']), ('c', ['d'])] + >>> list(collapse(iterable)) # Fully flattened + ['a', 'b', 'c', 'd'] + >>> list(collapse(iterable, levels=1)) # Only one level flattened + ['a', ['b'], 'c', ['d']] + + """ + stack = deque() + # Add our first node group, treat the iterable as a single node + stack.appendleft((0, repeat(iterable, 1))) + + while stack: + node_group = stack.popleft() + level, nodes = node_group + + # Check if beyond max level + if levels is not None and level > levels: + yield from nodes + continue + + for node in nodes: + # Check if done iterating + if isinstance(node, (str, bytes)) or ( + (base_type is not None) and isinstance(node, base_type) + ): + yield node + # Otherwise try to create child nodes + else: + try: + tree = iter(node) + except TypeError: + yield node + else: + # Save our current location + stack.appendleft(node_group) + # Append the new child node + stack.appendleft((level + 1, tree)) + # Break to process child node + break + + +def side_effect(func, iterable, chunk_size=None, before=None, after=None): + """Invoke *func* on each item in *iterable* (or on each *chunk_size* group + of items) before yielding the item. + + `func` must be a function that takes a single argument. Its return value + will be discarded. + + *before* and *after* are optional functions that take no arguments. They + will be executed before iteration starts and after it ends, respectively. + + `side_effect` can be used for logging, updating progress bars, or anything + that is not functionally "pure." + + Emitting a status message: + + >>> from more_itertools import consume + >>> func = lambda item: print('Received {}'.format(item)) + >>> consume(side_effect(func, range(2))) + Received 0 + Received 1 + + Operating on chunks of items: + + >>> pair_sums = [] + >>> func = lambda chunk: pair_sums.append(sum(chunk)) + >>> list(side_effect(func, [0, 1, 2, 3, 4, 5], 2)) + [0, 1, 2, 3, 4, 5] + >>> list(pair_sums) + [1, 5, 9] + + Writing to a file-like object: + + >>> from io import StringIO + >>> from more_itertools import consume + >>> f = StringIO() + >>> func = lambda x: print(x, file=f) + >>> before = lambda: print('HEADER', file=f) + >>> after = f.close + >>> it = ['a', 'b', 'c'] + >>> consume(side_effect(func, it, before=before, after=after)) + >>> f.closed + True + + """ + try: + if before is not None: + before() + + if chunk_size is None: + for item in iterable: + func(item) + yield item + else: + for chunk in chunked(iterable, chunk_size): + func(chunk) + yield from chunk + finally: + if after is not None: + after() + + +def sliced(seq, n, strict=False): + """Yield slices of length *n* from the sequence *seq*. + + >>> list(sliced((1, 2, 3, 4, 5, 6), 3)) + [(1, 2, 3), (4, 5, 6)] + + By the default, the last yielded slice will have fewer than *n* elements + if the length of *seq* is not divisible by *n*: + + >>> list(sliced((1, 2, 3, 4, 5, 6, 7, 8), 3)) + [(1, 2, 3), (4, 5, 6), (7, 8)] + + If the length of *seq* is not divisible by *n* and *strict* is + ``True``, then ``ValueError`` will be raised before the last + slice is yielded. + + This function will only work for iterables that support slicing. + For non-sliceable iterables, see :func:`chunked`. + + """ + if n < 0: + raise ValueError('n must be at least 0') + + iterator = takewhile(len, (seq[i : i + n] for i in count(0, n))) + if strict: + + def ret(): + for _slice in iterator: + if len(_slice) != n: + raise ValueError("seq is not divisible by n.") + yield _slice + + return ret() + else: + return iterator + + +def split_at(iterable, pred, maxsplit=-1, keep_separator=False): + """Yield lists of items from *iterable*, where each list is delimited by + an item where callable *pred* returns ``True``. + + >>> list(split_at('abcdcba', lambda x: x == 'b')) + [['a'], ['c', 'd', 'c'], ['a']] + + >>> list(split_at(range(10), lambda n: n % 2 == 1)) + [[0], [2], [4], [6], [8], []] + + At most *maxsplit* splits are done. If *maxsplit* is not specified or -1, + then there is no limit on the number of splits: + + >>> list(split_at(range(10), lambda n: n % 2 == 1, maxsplit=2)) + [[0], [2], [4, 5, 6, 7, 8, 9]] + + By default, the delimiting items are not included in the output. + To include them, set *keep_separator* to ``True``. + + >>> list(split_at('abcdcba', lambda x: x == 'b', keep_separator=True)) + [['a'], ['b'], ['c', 'd', 'c'], ['b'], ['a']] + + """ + if maxsplit == 0: + yield list(iterable) + return + + buf = [] + it = iter(iterable) + for item in it: + if pred(item): + yield buf + if keep_separator: + yield [item] + if maxsplit == 1: + yield list(it) + return + buf = [] + maxsplit -= 1 + else: + buf.append(item) + yield buf + + +def split_before(iterable, pred, maxsplit=-1): + """Yield lists of items from *iterable*, where each list ends just before + an item for which callable *pred* returns ``True``: + + >>> list(split_before('OneTwo', lambda s: s.isupper())) + [['O', 'n', 'e'], ['T', 'w', 'o']] + + >>> list(split_before(range(10), lambda n: n % 3 == 0)) + [[0, 1, 2], [3, 4, 5], [6, 7, 8], [9]] + + At most *maxsplit* splits are done. If *maxsplit* is not specified or -1, + then there is no limit on the number of splits: + + >>> list(split_before(range(10), lambda n: n % 3 == 0, maxsplit=2)) + [[0, 1, 2], [3, 4, 5], [6, 7, 8, 9]] + """ + if maxsplit == 0: + buf = list(iterable) + if buf: + yield buf + return + + buf = [] + it = iter(iterable) + for item in it: + if pred(item) and buf: + yield buf + if maxsplit == 1: + yield [item, *it] + return + buf = [] + maxsplit -= 1 + buf.append(item) + if buf: + yield buf + + +def split_after(iterable, pred, maxsplit=-1): + """Yield lists of items from *iterable*, where each list ends with an + item where callable *pred* returns ``True``: + + >>> list(split_after('one1two2', lambda s: s.isdigit())) + [['o', 'n', 'e', '1'], ['t', 'w', 'o', '2']] + + >>> list(split_after(range(10), lambda n: n % 3 == 0)) + [[0], [1, 2, 3], [4, 5, 6], [7, 8, 9]] + + At most *maxsplit* splits are done. If *maxsplit* is not specified or -1, + then there is no limit on the number of splits: + + >>> list(split_after(range(10), lambda n: n % 3 == 0, maxsplit=2)) + [[0], [1, 2, 3], [4, 5, 6, 7, 8, 9]] + + """ + if maxsplit == 0: + buf = list(iterable) + if buf: + yield buf + return + + buf = [] + it = iter(iterable) + for item in it: + buf.append(item) + if pred(item) and buf: + yield buf + if maxsplit == 1: + buf = list(it) + if buf: + yield buf + return + buf = [] + maxsplit -= 1 + if buf: + yield buf + + +def split_when(iterable, pred, maxsplit=-1): + """Split *iterable* into pieces based on the output of *pred*. + *pred* should be a function that takes successive pairs of items and + returns ``True`` if the iterable should be split in between them. + + For example, to find runs of increasing numbers, split the iterable when + element ``i`` is larger than element ``i + 1``: + + >>> list(split_when([1, 2, 3, 3, 2, 5, 2, 4, 2], lambda x, y: x > y)) + [[1, 2, 3, 3], [2, 5], [2, 4], [2]] + + At most *maxsplit* splits are done. If *maxsplit* is not specified or -1, + then there is no limit on the number of splits: + + >>> list(split_when([1, 2, 3, 3, 2, 5, 2, 4, 2], + ... lambda x, y: x > y, maxsplit=2)) + [[1, 2, 3, 3], [2, 5], [2, 4, 2]] + + """ + if maxsplit == 0: + buf = list(iterable) + if buf: + yield buf + return + + it = iter(iterable) + try: + cur_item = next(it) + except StopIteration: + return + + buf = [cur_item] + for next_item in it: + if pred(cur_item, next_item): + yield buf + if maxsplit == 1: + yield [next_item, *it] + return + buf = [] + maxsplit -= 1 + + buf.append(next_item) + cur_item = next_item + + yield buf + + +def split_into(iterable, sizes): + """Yield a list of sequential items from *iterable* of length 'n' for each + integer 'n' in *sizes*. + + >>> list(split_into([1,2,3,4,5,6], [1,2,3])) + [[1], [2, 3], [4, 5, 6]] + + If the sum of *sizes* is smaller than the length of *iterable*, then the + remaining items of *iterable* will not be returned. + + >>> list(split_into([1,2,3,4,5,6], [2,3])) + [[1, 2], [3, 4, 5]] + + If the sum of *sizes* is larger than the length of *iterable*, fewer items + will be returned in the iteration that overruns the *iterable* and further + lists will be empty: + + >>> list(split_into([1,2,3,4], [1,2,3,4])) + [[1], [2, 3], [4], []] + + When a ``None`` object is encountered in *sizes*, the returned list will + contain items up to the end of *iterable* the same way that + :func:`itertools.slice` does: + + >>> list(split_into([1,2,3,4,5,6,7,8,9,0], [2,3,None])) + [[1, 2], [3, 4, 5], [6, 7, 8, 9, 0]] + + :func:`split_into` can be useful for grouping a series of items where the + sizes of the groups are not uniform. An example would be where in a row + from a table, multiple columns represent elements of the same feature + (e.g. a point represented by x,y,z) but, the format is not the same for + all columns. + """ + # convert the iterable argument into an iterator so its contents can + # be consumed by islice in case it is a generator + it = iter(iterable) + + for size in sizes: + if size is None: + yield list(it) + return + else: + yield list(islice(it, size)) + + +def padded(iterable, fillvalue=None, n=None, next_multiple=False): + """Yield the elements from *iterable*, followed by *fillvalue*, such that + at least *n* items are emitted. + + >>> list(padded([1, 2, 3], '?', 5)) + [1, 2, 3, '?', '?'] + + If *next_multiple* is ``True``, *fillvalue* will be emitted until the + number of items emitted is a multiple of *n*: + + >>> list(padded([1, 2, 3, 4], n=3, next_multiple=True)) + [1, 2, 3, 4, None, None] + + If *n* is ``None``, *fillvalue* will be emitted indefinitely. + + To create an *iterable* of exactly size *n*, you can truncate with + :func:`islice`. + + >>> list(islice(padded([1, 2, 3], '?'), 5)) + [1, 2, 3, '?', '?'] + >>> list(islice(padded([1, 2, 3, 4, 5, 6, 7, 8], '?'), 5)) + [1, 2, 3, 4, 5] + + """ + iterator = iter(iterable) + iterator_with_repeat = chain(iterator, repeat(fillvalue)) + + if n is None: + return iterator_with_repeat + elif n < 1: + raise ValueError('n must be at least 1') + elif next_multiple: + + def slice_generator(): + for first in iterator: + yield (first,) + yield islice(iterator_with_repeat, n - 1) + + # While elements exist produce slices of size n + return chain.from_iterable(slice_generator()) + else: + # Ensure the first batch is at least size n then iterate + return chain(islice(iterator_with_repeat, n), iterator) + + +def repeat_each(iterable, n=2): + """Repeat each element in *iterable* *n* times. + + >>> list(repeat_each('ABC', 3)) + ['A', 'A', 'A', 'B', 'B', 'B', 'C', 'C', 'C'] + """ + return chain.from_iterable(map(repeat, iterable, repeat(n))) + + +def repeat_last(iterable, default=None): + """After the *iterable* is exhausted, keep yielding its last element. + + >>> list(islice(repeat_last(range(3)), 5)) + [0, 1, 2, 2, 2] + + If the iterable is empty, yield *default* forever:: + + >>> list(islice(repeat_last(range(0), 42), 5)) + [42, 42, 42, 42, 42] + + """ + item = _marker + for item in iterable: + yield item + final = default if item is _marker else item + yield from repeat(final) + + +def distribute(n, iterable): + """Distribute the items from *iterable* among *n* smaller iterables. + + >>> group_1, group_2 = distribute(2, [1, 2, 3, 4, 5, 6]) + >>> list(group_1) + [1, 3, 5] + >>> list(group_2) + [2, 4, 6] + + If the length of *iterable* is not evenly divisible by *n*, then the + length of the returned iterables will not be identical: + + >>> children = distribute(3, [1, 2, 3, 4, 5, 6, 7]) + >>> [list(c) for c in children] + [[1, 4, 7], [2, 5], [3, 6]] + + If the length of *iterable* is smaller than *n*, then the last returned + iterables will be empty: + + >>> children = distribute(5, [1, 2, 3]) + >>> [list(c) for c in children] + [[1], [2], [3], [], []] + + This function uses :func:`itertools.tee` and may require significant + storage. + + If you need the order items in the smaller iterables to match the + original iterable, see :func:`divide`. + + """ + if n < 1: + raise ValueError('n must be at least 1') + + children = tee(iterable, n) + return [islice(it, index, None, n) for index, it in enumerate(children)] + + +def stagger(iterable, offsets=(-1, 0, 1), longest=False, fillvalue=None): + """Yield tuples whose elements are offset from *iterable*. + The amount by which the `i`-th item in each tuple is offset is given by + the `i`-th item in *offsets*. + + >>> list(stagger([0, 1, 2, 3])) + [(None, 0, 1), (0, 1, 2), (1, 2, 3)] + >>> list(stagger(range(8), offsets=(0, 2, 4))) + [(0, 2, 4), (1, 3, 5), (2, 4, 6), (3, 5, 7)] + + By default, the sequence will end when the final element of a tuple is the + last item in the iterable. To continue until the first element of a tuple + is the last item in the iterable, set *longest* to ``True``:: + + >>> list(stagger([0, 1, 2, 3], longest=True)) + [(None, 0, 1), (0, 1, 2), (1, 2, 3), (2, 3, None), (3, None, None)] + + By default, ``None`` will be used to replace offsets beyond the end of the + sequence. Specify *fillvalue* to use some other value. + + """ + children = tee(iterable, len(offsets)) + + return zip_offset( + *children, offsets=offsets, longest=longest, fillvalue=fillvalue + ) + + +def zip_offset(*iterables, offsets, longest=False, fillvalue=None): + """``zip`` the input *iterables* together, but offset the `i`-th iterable + by the `i`-th item in *offsets*. + + >>> list(zip_offset('0123', 'abcdef', offsets=(0, 1))) + [('0', 'b'), ('1', 'c'), ('2', 'd'), ('3', 'e')] + + This can be used as a lightweight alternative to SciPy or pandas to analyze + data sets in which some series have a lead or lag relationship. + + By default, the sequence will end when the shortest iterable is exhausted. + To continue until the longest iterable is exhausted, set *longest* to + ``True``. + + >>> list(zip_offset('0123', 'abcdef', offsets=(0, 1), longest=True)) + [('0', 'b'), ('1', 'c'), ('2', 'd'), ('3', 'e'), (None, 'f')] + + By default, ``None`` will be used to replace offsets beyond the end of the + sequence. Specify *fillvalue* to use some other value. + + """ + if len(iterables) != len(offsets): + raise ValueError("Number of iterables and offsets didn't match") + + staggered = [] + for it, n in zip(iterables, offsets): + if n < 0: + staggered.append(chain(repeat(fillvalue, -n), it)) + elif n > 0: + staggered.append(islice(it, n, None)) + else: + staggered.append(it) + + if longest: + return zip_longest(*staggered, fillvalue=fillvalue) + + return zip(*staggered) + + +def sort_together( + iterables, key_list=(0,), key=None, reverse=False, strict=False +): + """Return the input iterables sorted together, with *key_list* as the + priority for sorting. All iterables are trimmed to the length of the + shortest one. + + This can be used like the sorting function in a spreadsheet. If each + iterable represents a column of data, the key list determines which + columns are used for sorting. + + By default, all iterables are sorted using the ``0``-th iterable:: + + >>> iterables = [(4, 3, 2, 1), ('a', 'b', 'c', 'd')] + >>> sort_together(iterables) + [(1, 2, 3, 4), ('d', 'c', 'b', 'a')] + + Set a different key list to sort according to another iterable. + Specifying multiple keys dictates how ties are broken:: + + >>> iterables = [(3, 1, 2), (0, 1, 0), ('c', 'b', 'a')] + >>> sort_together(iterables, key_list=(1, 2)) + [(2, 3, 1), (0, 0, 1), ('a', 'c', 'b')] + + To sort by a function of the elements of the iterable, pass a *key* + function. Its arguments are the elements of the iterables corresponding to + the key list:: + + >>> names = ('a', 'b', 'c') + >>> lengths = (1, 2, 3) + >>> widths = (5, 2, 1) + >>> def area(length, width): + ... return length * width + >>> sort_together([names, lengths, widths], key_list=(1, 2), key=area) + [('c', 'b', 'a'), (3, 2, 1), (1, 2, 5)] + + Set *reverse* to ``True`` to sort in descending order. + + >>> sort_together([(1, 2, 3), ('c', 'b', 'a')], reverse=True) + [(3, 2, 1), ('a', 'b', 'c')] + + If the *strict* keyword argument is ``True``, then + ``ValueError`` will be raised if any of the iterables have + different lengths. + + """ + if key is None: + # if there is no key function, the key argument to sorted is an + # itemgetter + key_argument = itemgetter(*key_list) + else: + # if there is a key function, call it with the items at the offsets + # specified by the key function as arguments + key_list = list(key_list) + if len(key_list) == 1: + # if key_list contains a single item, pass the item at that offset + # as the only argument to the key function + key_offset = key_list[0] + key_argument = lambda zipped_items: key(zipped_items[key_offset]) + else: + # if key_list contains multiple items, use itemgetter to return a + # tuple of items, which we pass as *args to the key function + get_key_items = itemgetter(*key_list) + key_argument = lambda zipped_items: key( + *get_key_items(zipped_items) + ) + + transposed = zip(*iterables, strict=strict) + reordered = sorted(transposed, key=key_argument, reverse=reverse) + untransposed = zip(*reordered, strict=strict) + return list(untransposed) + + +def unzip(iterable): + """The inverse of :func:`zip`, this function disaggregates the elements + of the zipped *iterable*. + + The ``i``-th iterable contains the ``i``-th element from each element + of the zipped iterable. The first element is used to determine the + length of the remaining elements. + + >>> iterable = [('a', 1), ('b', 2), ('c', 3), ('d', 4)] + >>> letters, numbers = unzip(iterable) + >>> list(letters) + ['a', 'b', 'c', 'd'] + >>> list(numbers) + [1, 2, 3, 4] + + This is similar to using ``zip(*iterable)``, but it avoids reading + *iterable* into memory. Note, however, that this function uses + :func:`itertools.tee` and thus may require significant storage. + + """ + head, iterable = spy(iterable) + if not head: + # empty iterable, e.g. zip([], [], []) + return () + # spy returns a one-length iterable as head + head = head[0] + iterables = tee(iterable, len(head)) + + # If we have an iterable like iter([(1, 2, 3), (4, 5), (6,)]), + # the second unzipped iterable fails at the third tuple since + # it tries to access (6,)[1]. + # Same with the third unzipped iterable and the second tuple. + # To support these "improperly zipped" iterables, we suppress + # the IndexError, which just stops the unzipped iterables at + # first length mismatch. + return tuple( + iter_suppress(map(itemgetter(i), it), IndexError) + for i, it in enumerate(iterables) + ) + + +def divide(n, iterable): + """Divide the elements from *iterable* into *n* parts, maintaining + order. + + >>> group_1, group_2 = divide(2, [1, 2, 3, 4, 5, 6]) + >>> list(group_1) + [1, 2, 3] + >>> list(group_2) + [4, 5, 6] + + If the length of *iterable* is not evenly divisible by *n*, then the + length of the returned iterables will not be identical: + + >>> children = divide(3, [1, 2, 3, 4, 5, 6, 7]) + >>> [list(c) for c in children] + [[1, 2, 3], [4, 5], [6, 7]] + + If the length of the iterable is smaller than n, then the last returned + iterables will be empty: + + >>> children = divide(5, [1, 2, 3]) + >>> [list(c) for c in children] + [[1], [2], [3], [], []] + + This function will exhaust the iterable before returning. + If order is not important, see :func:`distribute`, which does not first + pull the iterable into memory. + + """ + if n < 1: + raise ValueError('n must be at least 1') + + try: + iterable[:0] + except TypeError: + seq = tuple(iterable) + else: + seq = iterable + + q, r = divmod(len(seq), n) + + ret = [] + stop = 0 + for i in range(1, n + 1): + start = stop + stop += q + 1 if i <= r else q + ret.append(iter(seq[start:stop])) + + return ret + + +def always_iterable(obj, base_type=(str, bytes)): + """If *obj* is iterable, return an iterator over its items:: + + >>> obj = (1, 2, 3) + >>> list(always_iterable(obj)) + [1, 2, 3] + + If *obj* is not iterable, return a one-item iterable containing *obj*:: + + >>> obj = 1 + >>> list(always_iterable(obj)) + [1] + + If *obj* is ``None``, return an empty iterable: + + >>> obj = None + >>> list(always_iterable(None)) + [] + + By default, binary and text strings are not considered iterable:: + + >>> obj = 'foo' + >>> list(always_iterable(obj)) + ['foo'] + + If *base_type* is set, objects for which ``isinstance(obj, base_type)`` + returns ``True`` won't be considered iterable. + + >>> obj = {'a': 1} + >>> list(always_iterable(obj)) # Iterate over the dict's keys + ['a'] + >>> list(always_iterable(obj, base_type=dict)) # Treat dicts as a unit + [{'a': 1}] + + Set *base_type* to ``None`` to avoid any special handling and treat objects + Python considers iterable as iterable: + + >>> obj = 'foo' + >>> list(always_iterable(obj, base_type=None)) + ['f', 'o', 'o'] + """ + if obj is None: + return iter(()) + + if (base_type is not None) and isinstance(obj, base_type): + return iter((obj,)) + + try: + return iter(obj) + except TypeError: + return iter((obj,)) + + +def adjacent(predicate, iterable, distance=1): + """Return an iterable over `(bool, item)` tuples where the `item` is + drawn from *iterable* and the `bool` indicates whether + that item satisfies the *predicate* or is adjacent to an item that does. + + For example, to find whether items are adjacent to a ``3``:: + + >>> list(adjacent(lambda x: x == 3, range(6))) + [(False, 0), (False, 1), (True, 2), (True, 3), (True, 4), (False, 5)] + + Set *distance* to change what counts as adjacent. For example, to find + whether items are two places away from a ``3``: + + >>> list(adjacent(lambda x: x == 3, range(6), distance=2)) + [(False, 0), (True, 1), (True, 2), (True, 3), (True, 4), (True, 5)] + + This is useful for contextualizing the results of a search function. + For example, a code comparison tool might want to identify lines that + have changed, but also surrounding lines to give the viewer of the diff + context. + + The predicate function will only be called once for each item in the + iterable. + + See also :func:`groupby_transform`, which can be used with this function + to group ranges of items with the same `bool` value. + + """ + # Allow distance=0 mainly for testing that it reproduces results with map() + if distance < 0: + raise ValueError('distance must be at least 0') + + i1, i2 = tee(iterable) + padding = [False] * distance + selected = chain(padding, map(predicate, i1), padding) + adjacent_to_selected = map(any, windowed(selected, 2 * distance + 1)) + return zip(adjacent_to_selected, i2) + + +def groupby_transform(iterable, keyfunc=None, valuefunc=None, reducefunc=None): + """An extension of :func:`itertools.groupby` that can apply transformations + to the grouped data. + + * *keyfunc* is a function computing a key value for each item in *iterable* + * *valuefunc* is a function that transforms the individual items from + *iterable* after grouping + * *reducefunc* is a function that transforms each group of items + + >>> iterable = 'aAAbBBcCC' + >>> keyfunc = lambda k: k.upper() + >>> valuefunc = lambda v: v.lower() + >>> reducefunc = lambda g: ''.join(g) + >>> list(groupby_transform(iterable, keyfunc, valuefunc, reducefunc)) + [('A', 'aaa'), ('B', 'bbb'), ('C', 'ccc')] + + Each optional argument defaults to an identity function if not specified. + + :func:`groupby_transform` is useful when grouping elements of an iterable + using a separate iterable as the key. To do this, :func:`zip` the iterables + and pass a *keyfunc* that extracts the first element and a *valuefunc* + that extracts the second element:: + + >>> from operator import itemgetter + >>> keys = [0, 0, 1, 1, 1, 2, 2, 2, 3] + >>> values = 'abcdefghi' + >>> iterable = zip(keys, values) + >>> grouper = groupby_transform(iterable, itemgetter(0), itemgetter(1)) + >>> [(k, ''.join(g)) for k, g in grouper] + [(0, 'ab'), (1, 'cde'), (2, 'fgh'), (3, 'i')] + + Note that the order of items in the iterable is significant. + Only adjacent items are grouped together, so if you don't want any + duplicate groups, you should sort the iterable by the key function + or consider :func:`bucket` or :func:`map_reduce`. :func:`map_reduce` + consumes the iterable immediately and returns a dictionary, while + :func:`bucket` does not. + + .. seealso:: :func:`bucket`, :func:`map_reduce` + + """ + ret = groupby(iterable, keyfunc) + if valuefunc: + ret = ((k, map(valuefunc, g)) for k, g in ret) + if reducefunc: + ret = ((k, reducefunc(g)) for k, g in ret) + + return ret + + +class numeric_range(Sequence): + """An extension of the built-in ``range()`` function whose arguments can + be any orderable numeric type. + + With only *stop* specified, *start* defaults to ``0`` and *step* + defaults to ``1``. The output items will match the type of *stop*: + + >>> list(numeric_range(3.5)) + [0.0, 1.0, 2.0, 3.0] + + With only *start* and *stop* specified, *step* defaults to ``1``, typed + to match *stop* - *start*. The output items will match the type of + ``start + step``: + + >>> from decimal import Decimal + >>> start = Decimal('2.1') + >>> stop = Decimal('5.1') + >>> list(numeric_range(start, stop)) + [Decimal('2.1'), Decimal('3.1'), Decimal('4.1')] + + With *start*, *stop*, and *step* specified the output items will match + the type of ``start + step``: + + >>> from fractions import Fraction + >>> start = Fraction(1, 2) # Start at 1/2 + >>> stop = Fraction(5, 2) # End at 5/2 + >>> step = Fraction(1, 2) # Count by 1/2 + >>> list(numeric_range(start, stop, step)) + [Fraction(1, 2), Fraction(1, 1), Fraction(3, 2), Fraction(2, 1)] + + If *step* is zero, ``ValueError`` is raised. Negative steps are supported: + + >>> list(numeric_range(3, -1, -1.0)) + [3.0, 2.0, 1.0, 0.0] + + Be aware of the limitations of floating-point numbers; the representation + of the yielded numbers may be surprising. + + ``datetime.datetime`` objects can be used for *start* and *stop*, if *step* + is a ``datetime.timedelta`` object: + + >>> import datetime + >>> start = datetime.datetime(2019, 1, 1) + >>> stop = datetime.datetime(2019, 1, 3) + >>> step = datetime.timedelta(days=1) + >>> items = iter(numeric_range(start, stop, step)) + >>> next(items) + datetime.datetime(2019, 1, 1, 0, 0) + >>> next(items) + datetime.datetime(2019, 1, 2, 0, 0) + + """ + + _EMPTY_HASH = hash(range(0, 0)) + + def __init__(self, *args): + argc = len(args) + if argc == 1: + (self._stop,) = args + self._start = type(self._stop)(0) + self._step = type(self._stop - self._start)(1) + elif argc == 2: + self._start, self._stop = args + self._step = type(self._stop - self._start)(1) + elif argc == 3: + self._start, self._stop, self._step = args + elif argc == 0: + raise TypeError( + f'numeric_range expected at least 1 argument, got {argc}' + ) + else: + raise TypeError( + f'numeric_range expected at most 3 arguments, got {argc}' + ) + + self._zero = type(self._step)(0) + if self._step == self._zero: + raise ValueError('numeric_range() arg 3 must not be zero') + self._growing = self._step > self._zero + + def __bool__(self): + if self._growing: + return self._start < self._stop + else: + return self._start > self._stop + + def __contains__(self, elem): + try: + self.index(elem) + except ValueError: + return False + return True + + def __eq__(self, other): + # numeric_range object equality is intended to mirror the built-in range + # object's equality. + # https://github.com/python/cpython/blob/f5c4880151b609e0a0a0b05c292d36b18038c061/Objects/rangeobject.c#L499 + if not isinstance(other, numeric_range): + return False + + if self is other: + return True + + len_self = len(self) + if len_self != len(other): + return False + + if not len_self: + return True + + if self._start != other._start: + return False + + if len_self == 1: + return True + + return self._step == other._step + + def __getitem__(self, key): + if isinstance(key, int): + return self._get_by_index(key) + elif isinstance(key, slice): + start_idx, stop_idx, step_idx = key.indices(self._len) + return numeric_range( + self._start + start_idx * self._step, + self._start + stop_idx * self._step, + self._step * step_idx, + ) + else: + raise TypeError( + 'numeric range indices must be ' + f'integers or slices, not {type(key).__name__}' + ) + + def __hash__(self): + # numeric_range hashing is intended to mirror the built-in range object's + # hashing. + # https://github.com/python/cpython/blob/f5c4880151b609e0a0a0b05c292d36b18038c061/Objects/rangeobject.c#L570 + len_self = len(self) + if not len_self: + return hash((len_self, None, None)) + if len_self == 1: + return hash((len_self, self._start, None)) + return hash((len_self, self._start, self._step)) + + def __iter__(self): + values = (self._start + (n * self._step) for n in count()) + if self._growing: + return takewhile(partial(gt, self._stop), values) + else: + return takewhile(partial(lt, self._stop), values) + + def __len__(self): + return self._len + + @cached_property + def _len(self): + if self._growing: + start = self._start + stop = self._stop + step = self._step + else: + start = self._stop + stop = self._start + step = -self._step + distance = stop - start + if distance <= self._zero: + return 0 + else: # distance > 0 and step > 0: regular euclidean division + q, r = divmod(distance, step) + n = int(q) + int(r != self._zero) + # The division above measures the distance to cover, but the items + # are produced by repeated multiplication (see `__iter__`). With + # inexact arithmetic the two disagree at the boundary, so settle + # the count against the items themselves. + while n and not self._before_stop(n - 1): + n -= 1 + while self._before_stop(n): + n += 1 + return n + + def _before_stop(self, i): + value = self._start + i * self._step + return value < self._stop if self._growing else value > self._stop + + def __reduce__(self): + return numeric_range, (self._start, self._stop, self._step) + + def __repr__(self): + if self._step == 1: + return f"numeric_range({self._start!r}, {self._stop!r})" + return ( + f"numeric_range({self._start!r}, {self._stop!r}, {self._step!r})" + ) + + def __reversed__(self): + start = self._start + step = self._step + for i in reversed(range(self._len)): + yield start + i * step + + def count(self, value): + return int(value in self) + + def index(self, value): + if self._growing: + if self._start <= value < self._stop: + q, _ = divmod(value - self._start, self._step) + return self._index_near(int(q), value) + else: + if self._start >= value > self._stop: + q, _ = divmod(self._start - value, -self._step) + return self._index_near(int(q), value) + + raise ValueError(f"{value} is not in numeric range") + + def _index_near(self, i, value): + # `i` is the quotient of the division of `value` by the step, which + # locates the value on the grid of items this range produces. For + # exact types that quotient is the index. For inexact ones (floats) + # it can land one short, and the remainder of the division is not a + # reliable membership test either: `numeric_range(0.0, 1.0, 0.1)` + # yields 0.30000000000000004, which leaves a non-zero remainder. + # So compare against the item the range actually produces there. + for candidate in (i, i + 1): + if 0 <= candidate < self._len: + if self._start + candidate * self._step == value: + return candidate + + raise ValueError(f"{value} is not in numeric range") + + def _get_by_index(self, i): + if i < 0: + i += self._len + if i < 0 or i >= self._len: + raise IndexError("numeric range object index out of range") + return self._start + i * self._step + + +def count_cycle(iterable, n=None): + """Cycle through the items from *iterable* up to *n* times, yielding + the number of completed cycles along with each item. If *n* is omitted the + process repeats indefinitely. + + >>> list(count_cycle('AB', 3)) + [(0, 'A'), (0, 'B'), (1, 'A'), (1, 'B'), (2, 'A'), (2, 'B')] + + """ + if n is not None: + return product(range(n), iterable) + seq = tuple(iterable) + if not seq: + return iter(()) + return zip(repeat_each(count(), len(seq)), cycle(seq)) + + +def mark_ends(iterable): + """Yield 3-tuples of the form ``(is_first, is_last, item)``. + + >>> list(mark_ends('ABC')) + [(True, False, 'A'), (False, False, 'B'), (False, True, 'C')] + + Use this when looping over an iterable to take special action on its first + and/or last items: + + >>> iterable = ['Header', 100, 200, 'Footer'] + >>> total = 0 + >>> for is_first, is_last, item in mark_ends(iterable): + ... if is_first: + ... continue # Skip the header + ... if is_last: + ... continue # Skip the footer + ... total += item + >>> print(total) + 300 + """ + it = iter(iterable) + for a in it: + first = True + for b in it: + yield first, False, a + a = b + first = False + yield first, True, a + + +def locate(iterable, pred=bool, window_size=None): + """Yield the index of each item in *iterable* for which *pred* returns + ``True``. + + *pred* defaults to :func:`bool`, which will select truthy items: + + >>> list(locate([0, 1, 1, 0, 1, 0, 0])) + [1, 2, 4] + + Set *pred* to a custom function to, e.g., find the indexes for a particular + item. + + >>> list(locate(['a', 'b', 'c', 'b'], lambda x: x == 'b')) + [1, 3] + + If *window_size* is given, then the *pred* function will be called with + the values in each window. This enables searching for sub-sequences. + Note that *pred* may receive fewer than *window_size* arguments at the end of + the iterable. + + >>> iterable = [0, 1, 2, 3, 0, 1, 2, 3, 0, 1, 2, 3] + >>> pred = lambda *args: args == (1, 2, 3) + >>> list(locate(iterable, pred=pred, window_size=3)) + [1, 5, 9] + + Use with :func:`seekable` to find indexes and then retrieve the associated + items: + + >>> from itertools import count + >>> from more_itertools import seekable + >>> source = (3 * n + 1 if (n % 2) else n // 2 for n in count()) + >>> it = seekable(source) + >>> pred = lambda x: x > 100 + >>> indexes = locate(it, pred=pred) + >>> i = next(indexes) + >>> it.seek(i) + >>> next(it) + 106 + + """ + if window_size is None: + return compress(count(), map(pred, iterable)) + + if window_size < 1: + raise ValueError('window size must be at least 1') + + it = windowed(iterable, window_size, fillvalue=_marker) + return compress( + count(), + (pred(*(x for x in w if x is not _marker)) for w in it), + ) + + +def longest_common_prefix(iterables): + """Yield elements of the longest common prefix among given *iterables*. + + >>> ''.join(longest_common_prefix(['abcd', 'abc', 'abf'])) + 'ab' + + """ + return (c[0] for c in takewhile(all_equal, zip(*iterables))) + + +def lstrip(iterable, pred): + """Yield the items from *iterable*, but strip any from the beginning + for which *pred* returns ``True``. + + For example, to remove a set of items from the start of an iterable: + + >>> iterable = (None, False, None, 1, 2, None, 3, False, None) + >>> pred = lambda x: x in {None, False, ''} + >>> list(lstrip(iterable, pred)) + [1, 2, None, 3, False, None] + + This function is analogous to :func:`str.lstrip`, and is essentially + a wrapper for :func:`itertools.dropwhile`. + + """ + return dropwhile(pred, iterable) + + +def rstrip(iterable, pred): + """Yield the items from *iterable*, but strip any from the end + for which *pred* returns ``True``. + + For example, to remove a set of items from the end of an iterable: + + >>> iterable = (None, False, None, 1, 2, None, 3, False, None) + >>> pred = lambda x: x in {None, False, ''} + >>> list(rstrip(iterable, pred)) + [None, False, None, 1, 2, None, 3] + + This function is analogous to :func:`str.rstrip`. + + """ + cache = [] + cache_append = cache.append + cache_clear = cache.clear + for x in iterable: + if pred(x): + cache_append(x) + else: + yield from cache + cache_clear() + yield x + + +def strip(iterable, pred): + """Yield the items from *iterable*, but strip any from the + beginning and end for which *pred* returns ``True``. + + For example, to remove a set of items from both ends of an iterable: + + >>> iterable = (None, False, None, 1, 2, None, 3, False, None) + >>> pred = lambda x: x in {None, False, ''} + >>> list(strip(iterable, pred)) + [1, 2, None, 3] + + This function is analogous to :func:`str.strip`. + + """ + return rstrip(lstrip(iterable, pred), pred) + + +class islice_extended: + """An extension of :func:`itertools.islice` that supports negative values + for *stop*, *start*, and *step*. + + >>> iterator = iter('abcdefgh') + >>> list(islice_extended(iterator, -4, -1)) + ['e', 'f', 'g'] + + Slices with negative values require some caching of *iterable*, but this + function takes care to minimize the amount of memory required. + + For example, you can use a negative step with an infinite iterator: + + >>> from itertools import count + >>> list(islice_extended(count(), 110, 99, -2)) + [110, 108, 106, 104, 102, 100] + + You can also use slice notation directly: + + >>> iterator = map(str, count()) + >>> it = islice_extended(iterator)[10:20:2] + >>> list(it) + ['10', '12', '14', '16', '18'] + + """ + + def __init__(self, iterable, *args): + it = iter(iterable) + if args: + self._iterator = _islice_helper(it, slice(*args)) + else: + self._iterator = it + + def __iter__(self): + return self + + def __next__(self): + return next(self._iterator) + + def __getitem__(self, key): + if isinstance(key, slice): + return islice_extended(_islice_helper(self._iterator, key)) + + raise TypeError('islice_extended.__getitem__ argument must be a slice') + + +def _islice_helper(it, s): + start = s.start + stop = s.stop + if s.step == 0: + raise ValueError('step argument must be a non-zero integer or None.') + step = s.step or 1 + + if step > 0: + start = 0 if (start is None) else start + + if start < 0: + # Consume all but the last -start items + counter = count(1) + wrapper = compress(it, counter) + cache = deque(wrapper, maxlen=-start) + len_iter = next(counter) - 1 + + # Adjust start to be positive + i = max(len_iter + start, 0) + + # Adjust stop to be positive + if stop is None: + j = len_iter + elif stop >= 0: + j = min(stop, len_iter) + else: + j = max(len_iter + stop, 0) + + # Slice the cache + n = j - i + if n <= 0: + return + + for index in range(n): + if index % step == 0: + # pop and yield the item. + # We don't want to use an intermediate variable + # it would extend the lifetime of the current item + yield cache.popleft() + else: + # just pop and discard the item + cache.popleft() + elif (stop is not None) and (stop < 0): + # Advance to the start position + next(islice(it, start, start), None) + + # When stop is negative, we have to carry -stop items while + # iterating + cache = deque(islice(it, -stop), maxlen=-stop) + + for index, item in enumerate(it): + if index % step == 0: + # pop and yield the item. + # We don't want to use an intermediate variable + # it would extend the lifetime of the current item + yield cache.popleft() + else: + # just pop and discard the item + cache.popleft() + cache.append(item) + else: + # When both start and stop are positive we have the normal case + yield from islice(it, start, stop, step) + else: + start = -1 if (start is None) else start + + if (stop is not None) and (stop < 0): + # Consume all but the last items + n = -stop - 1 + counter = count(1) + wrapper = compress(it, counter) + cache = deque(wrapper, maxlen=n) + len_iter = next(counter) - 1 + + # If start and stop are both negative they are comparable and + # we can just slice. Otherwise we can adjust start to be negative + # and then slice. + if start < 0: + i, j = start, stop + else: + i, j = min(start - len_iter, -1), None + + yield from list(cache)[i:j:step] + else: + # Advance to the stop position + if stop is not None: + m = stop + 1 + next(islice(it, m, m), None) + + # stop is positive, so if start is negative they are not comparable + # and we need the rest of the items. + if start < 0: + i = start + n = None + # stop is None and start is positive, so we just need items up to + # the start index. + elif stop is None: + i = None + n = start + 1 + # Both stop and start are positive, so they are comparable. + else: + i = None + n = start - stop + if n <= 0: + return + + cache = list(islice(it, n)) + + yield from cache[i::step] + + +def always_reversible(iterable): + """An extension of :func:`reversed` that supports all iterables, not + just those which implement the ``Reversible`` or ``Sequence`` protocols. + + >>> print(*always_reversible(x for x in range(3))) + 2 1 0 + + If the iterable is already reversible, this function returns the + result of :func:`reversed()`. If the iterable is not reversible, + this function will cache the remaining items in the iterable and + yield them in reverse order, which may require significant storage. + """ + try: + return reversed(iterable) + except TypeError: + return reversed(list(iterable)) + + +def consecutive_groups(iterable, ordering=None): + """Yield groups of consecutive items using :func:`itertools.groupby`. + The *ordering* function determines whether two items are adjacent by + returning their position. + + By default, the ordering function is the identity function. This is + suitable for finding runs of numbers: + + >>> iterable = [1, 10, 11, 12, 20, 30, 31, 32, 33, 40] + >>> for group in consecutive_groups(iterable): + ... print(list(group)) + [1] + [10, 11, 12] + [20] + [30, 31, 32, 33] + [40] + + To find runs of adjacent letters, apply :func:`ord` function + to convert letters to ordinals. + + >>> iterable = 'abcdfgilmnop' + >>> ordering = ord + >>> for group in consecutive_groups(iterable, ordering): + ... print(list(group)) + ['a', 'b', 'c', 'd'] + ['f', 'g'] + ['i'] + ['l', 'm', 'n', 'o', 'p'] + + Each group of consecutive items is an iterator that shares its source with + *iterable*. When an output group is advanced, the previous group is + no longer available unless its elements are copied (e.g., into a ``list``). + + >>> iterable = [1, 2, 11, 12, 21, 22] + >>> saved_groups = [] + >>> for group in consecutive_groups(iterable): + ... saved_groups.append(list(group)) # Copy group elements + >>> saved_groups + [[1, 2], [11, 12], [21, 22]] + + """ + if ordering is None: + key = lambda x: x[0] - x[1] + else: + key = lambda x: x[0] - ordering(x[1]) + + for k, g in groupby(enumerate(iterable), key=key): + yield map(itemgetter(1), g) + + +def difference(iterable, func=sub, *, initial=None): + """This function is the inverse of :func:`itertools.accumulate`. By default + it will compute the first difference of *iterable* using + :func:`operator.sub`: + + >>> from itertools import accumulate + >>> iterable = accumulate([0, 1, 2, 3, 4]) # produces 0, 1, 3, 6, 10 + >>> list(difference(iterable)) + [0, 1, 2, 3, 4] + + *func* defaults to :func:`operator.sub`, but other functions can be + specified. They will be applied as follows:: + + A, B, C, D, ... --> A, func(B, A), func(C, B), func(D, C), ... + + For example, to do progressive division: + + >>> iterable = [1, 2, 6, 24, 120] + >>> func = lambda x, y: x // y + >>> list(difference(iterable, func)) + [1, 2, 3, 4, 5] + + If the *initial* keyword is set, the first element will be skipped when + computing successive differences. + + >>> it = [10, 11, 13, 16] # from accumulate([1, 2, 3], initial=10) + >>> list(difference(it, initial=10)) + [1, 2, 3] + + """ + a, b = tee(iterable) + try: + first = [next(b)] + except StopIteration: + return iter([]) + + if initial is not None: + return map(func, b, a) + + return chain(first, map(func, b, a)) + + +class SequenceView(Sequence): + """Return a read-only view of the sequence object *target*. + + :class:`SequenceView` objects are analogous to Python's built-in + "dictionary view" types. They provide a dynamic view of a sequence's items, + meaning that when the sequence updates, so does the view. + + >>> seq = ['0', '1', '2'] + >>> view = SequenceView(seq) + >>> view + SequenceView(['0', '1', '2']) + >>> seq.append('3') + >>> view + SequenceView(['0', '1', '2', '3']) + + Sequence views support indexing, slicing, and length queries. They act + like the underlying sequence, except they don't allow assignment: + + >>> view[1] + '1' + >>> view[1:-1] + ['1', '2'] + >>> len(view) + 4 + + Sequence views are useful as an alternative to copying, as they don't + require (much) extra storage. + + """ + + def __init__(self, target): + if not isinstance(target, Sequence): + raise TypeError + self._target = target + + def __getitem__(self, index): + return self._target[index] + + def __len__(self): + return len(self._target) + + def __repr__(self): + return f'{self.__class__.__name__}({self._target!r})' + + +class seekable: + """Wrap an iterator to allow for seeking backward and forward. This + progressively caches the items in the source iterable so they can be + re-visited. + + Call :meth:`seek` with an index to seek to that position in the source + iterable. + + To "reset" an iterator, seek to ``0``: + + >>> from itertools import count + >>> it = seekable((str(n) for n in count())) + >>> next(it), next(it), next(it) + ('0', '1', '2') + >>> it.seek(0) + >>> next(it), next(it), next(it) + ('0', '1', '2') + + You can also seek forward: + + >>> it = seekable((str(n) for n in range(20))) + >>> it.seek(10) + >>> next(it) + '10' + >>> it.seek(20) # Seeking past the end of the source isn't a problem + >>> list(it) + [] + >>> it.seek(0) # Resetting works even after hitting the end + >>> next(it) + '0' + + Call :meth:`relative_seek` to seek relative to the source iterator's + current position. + + >>> it = seekable((str(n) for n in range(20))) + >>> next(it), next(it), next(it) + ('0', '1', '2') + >>> it.relative_seek(2) + >>> next(it) + '5' + >>> it.relative_seek(-3) # Source is at '6', we move back to '3' + >>> next(it) + '3' + >>> it.relative_seek(-3) # Source is at '4', we move back to '1' + >>> next(it) + '1' + + + Call :meth:`peek` to look ahead one item without advancing the iterator: + + >>> it = seekable('1234') + >>> it.peek() + '1' + >>> list(it) + ['1', '2', '3', '4'] + >>> it.peek(default='empty') + 'empty' + + Before the iterator is at its end, calling :func:`bool` on it will return + ``True``. After it will return ``False``: + + >>> it = seekable('5678') + >>> bool(it) + True + >>> list(it) + ['5', '6', '7', '8'] + >>> bool(it) + False + + You may view the contents of the cache with the :meth:`elements` method. + That returns a :class:`SequenceView`, a view that updates automatically: + + >>> it = seekable((str(n) for n in range(10))) + >>> next(it), next(it), next(it) + ('0', '1', '2') + >>> elements = it.elements() + >>> elements + SequenceView(['0', '1', '2']) + >>> next(it) + '3' + >>> elements + SequenceView(['0', '1', '2', '3']) + + Indexing the :class:`seekable` directly returns items from the cache: + + >>> it = seekable((str(n) for n in range(10))) + >>> next(it), next(it), next(it) + ('0', '1', '2') + >>> it[-1] + '2' + >>> it[0] + '0' + + By default, the cache grows as the source iterable progresses, so beware of + wrapping very large or infinite iterables. Supply *maxlen* to limit the + size of the cache (this of course limits how far back you can seek). + + >>> from itertools import count + >>> it = seekable((str(n) for n in count()), maxlen=2) + >>> next(it), next(it), next(it), next(it) + ('0', '1', '2', '3') + >>> list(it.elements()) + ['2', '3'] + >>> it.seek(0) + >>> next(it), next(it), next(it), next(it) + ('2', '3', '4', '5') + >>> next(it) + '6' + + """ + + def __init__(self, iterable, maxlen=None): + self._maxlen_zero = maxlen == 0 + self._source = ( + peekable(iterable) if self._maxlen_zero else iter(iterable) + ) + if maxlen is None: + self._cache = [] + else: + self._cache = deque([], maxlen) + self._index = None + + def __iter__(self): + return self + + def __next__(self): + if self._index is not None: + try: + item = self._cache[self._index] + except IndexError: + self._index = None + else: + self._index += 1 + return item + + item = next(self._source) + self._cache.append(item) + return item + + def __bool__(self): + try: + self.peek() + except StopIteration: + return False + return True + + def peek(self, default=_marker): + if self._maxlen_zero: + return self._source.peek(default) + try: + peeked = next(self) + except StopIteration: + if default is _marker: + raise + return default + if self._index is None: + self._index = len(self._cache) + self._index -= 1 + return peeked + + def elements(self): + return SequenceView(self._cache) + + def seek(self, index): + self._index = index + remainder = index - len(self._cache) + if remainder > 0: + consume(self, remainder) + + def relative_seek(self, count): + if self._index is None: + self._index = len(self._cache) + + self.seek(max(self._index + count, 0)) + + def __getitem__(self, index): + return self._cache[index] + + +class run_length: + """ + :func:`run_length.encode` compresses an iterable with run-length encoding. + It yields groups of repeated items with the count of how many times they + were repeated: + + >>> uncompressed = 'abbcccdddd' + >>> list(run_length.encode(uncompressed)) + [('a', 1), ('b', 2), ('c', 3), ('d', 4)] + + :func:`run_length.decode` decompresses an iterable that was previously + compressed with run-length encoding. It yields the items of the + decompressed iterable: + + >>> compressed = [('a', 1), ('b', 2), ('c', 3), ('d', 4)] + >>> list(run_length.decode(compressed)) + ['a', 'b', 'b', 'c', 'c', 'c', 'd', 'd', 'd', 'd'] + + """ + + @staticmethod + def encode(iterable): + return ((k, ilen(g)) for k, g in groupby(iterable)) + + @staticmethod + def decode(iterable): + return chain.from_iterable(starmap(repeat, iterable)) + + +def exactly_n(iterable, n, predicate=bool): + """Return ``True`` if exactly ``n`` items in the iterable are ``True`` + according to the *predicate* function. + + >>> exactly_n([True, True, False], 2) + True + >>> exactly_n([True, True, False], 1) + False + >>> exactly_n([0, 1, 2, 3, 4, 5], 3, lambda x: x < 3) + True + + The iterable will be advanced until ``n + 1`` truthy items are encountered, + so avoid calling it on infinite iterables. + + """ + iterator = filter(predicate, iterable) + if n <= 0: + if n < 0: + return False + for _ in iterator: + return False + return True + + iterator = islice(iterator, n - 1, None) + for _ in iterator: + for _ in iterator: + return False + return True + return False + + +def circular_shifts(iterable, steps=1): + """Yield the circular shifts of *iterable*. + + >>> list(circular_shifts(range(4))) + [(0, 1, 2, 3), (1, 2, 3, 0), (2, 3, 0, 1), (3, 0, 1, 2)] + + Set *steps* to the number of places to rotate to the left + (or to the right if negative). Defaults to 1. + + >>> list(circular_shifts(range(4), 2)) + [(0, 1, 2, 3), (2, 3, 0, 1)] + + >>> list(circular_shifts(range(4), -1)) + [(0, 1, 2, 3), (3, 0, 1, 2), (2, 3, 0, 1), (1, 2, 3, 0)] + + """ + buffer = deque(iterable) + if steps == 0: + raise ValueError('Steps should be a non-zero integer') + + buffer.rotate(steps) + steps = -steps + n = len(buffer) + n //= math.gcd(n, steps) + + for _ in repeat(None, n): + buffer.rotate(steps) + yield tuple(buffer) + + +def make_decorator(wrapping_func, result_index=0): + """Return a decorator version of *wrapping_func*, which is a function that + modifies an iterable. *result_index* is the position in that function's + signature where the iterable goes. + + This lets you use itertools on the "production end," i.e. at function + definition. This can augment what the function returns without changing the + function's code. + + For example, to produce a decorator version of :func:`chunked`: + + >>> from more_itertools import chunked + >>> chunker = make_decorator(chunked, result_index=0) + >>> @chunker(3) + ... def iter_range(n): + ... return iter(range(n)) + ... + >>> list(iter_range(9)) + [[0, 1, 2], [3, 4, 5], [6, 7, 8]] + + To only allow truthy items to be returned: + + >>> truth_serum = make_decorator(filter, result_index=1) + >>> @truth_serum(bool) + ... def boolean_test(): + ... return [0, 1, '', ' ', False, True] + ... + >>> list(boolean_test()) + [1, ' ', True] + + The :func:`peekable` and :func:`seekable` wrappers make for practical + decorators: + + >>> from more_itertools import peekable + >>> peekable_function = make_decorator(peekable) + >>> @peekable_function() + ... def str_range(*args): + ... return (str(x) for x in range(*args)) + ... + >>> it = str_range(1, 20, 2) + >>> next(it), next(it), next(it) + ('1', '3', '5') + >>> it.peek() + '7' + >>> next(it) + '7' + + """ + + # See https://sites.google.com/site/bbayles/index/decorator_factory for + # notes on how this works. + def decorator(*wrapping_args, **wrapping_kwargs): + def outer_wrapper(f): + def inner_wrapper(*args, **kwargs): + result = f(*args, **kwargs) + wrapping_args_ = list(wrapping_args) + wrapping_args_.insert(result_index, result) + return wrapping_func(*wrapping_args_, **wrapping_kwargs) + + return inner_wrapper + + return outer_wrapper + + return decorator + + +def map_reduce(iterable, keyfunc, valuefunc=None, reducefunc=None): + """Return a dictionary that maps the items in *iterable* to categories + defined by *keyfunc*, transforms them with *valuefunc*, and + then summarizes them by category with *reducefunc*. + + *valuefunc* defaults to the identity function if it is unspecified. + If *reducefunc* is unspecified, no summarization takes place: + + >>> keyfunc = lambda x: x.upper() + >>> result = map_reduce('abbccc', keyfunc) + >>> sorted(result.items()) + [('A', ['a']), ('B', ['b', 'b']), ('C', ['c', 'c', 'c'])] + + Specifying *valuefunc* transforms the categorized items: + + >>> keyfunc = lambda x: x.upper() + >>> valuefunc = lambda x: 1 + >>> result = map_reduce('abbccc', keyfunc, valuefunc) + >>> sorted(result.items()) + [('A', [1]), ('B', [1, 1]), ('C', [1, 1, 1])] + + Specifying *reducefunc* summarizes the categorized items: + + >>> keyfunc = lambda x: x.upper() + >>> valuefunc = lambda x: 1 + >>> reducefunc = sum + >>> result = map_reduce('abbccc', keyfunc, valuefunc, reducefunc) + >>> sorted(result.items()) + [('A', 1), ('B', 2), ('C', 3)] + + You may want to filter the input iterable before applying the map/reduce + procedure: + + >>> all_items = range(30) + >>> items = [x for x in all_items if 10 <= x <= 20] # Filter + >>> keyfunc = lambda x: x % 2 # Evens map to 0; odds to 1 + >>> categories = map_reduce(items, keyfunc=keyfunc) + >>> sorted(categories.items()) + [(0, [10, 12, 14, 16, 18, 20]), (1, [11, 13, 15, 17, 19])] + >>> summaries = map_reduce(items, keyfunc=keyfunc, reducefunc=sum) + >>> sorted(summaries.items()) + [(0, 90), (1, 75)] + + Note that all items in the iterable are gathered into a list before the + summarization step, which may require significant storage. + + The returned object is a :obj:`collections.defaultdict` with the + ``default_factory`` set to ``None``, such that it behaves like a normal + dictionary. + + .. seealso:: :func:`bucket`, :func:`groupby_transform` + + If storage is a concern, :func:`bucket` can be used without consuming the + entire iterable right away. If the elements with the same key are already + adjacent, :func:`groupby_transform` or :func:`itertools.groupby` can be + used without any caching overhead. + + """ + + ret = defaultdict(list) + + if valuefunc is None: + for item in iterable: + key = keyfunc(item) + ret[key].append(item) + + else: + for item in iterable: + key = keyfunc(item) + value = valuefunc(item) + ret[key].append(value) + + if reducefunc is not None: + for key, value_list in ret.items(): + ret[key] = reducefunc(value_list) + + ret.default_factory = None + return ret + + +def rlocate(iterable, pred=bool, window_size=None): + """Yield the index of each item in *iterable* for which *pred* returns + ``True``, starting from the right and moving left. + + *pred* defaults to :func:`bool`, which will select truthy items: + + >>> list(rlocate([0, 1, 1, 0, 1, 0, 0])) # Truthy at 1, 2, and 4 + [4, 2, 1] + + Set *pred* to a custom function to, e.g., find the indexes for a particular + item: + + >>> iterator = iter('abcb') + >>> pred = lambda x: x == 'b' + >>> list(rlocate(iterator, pred)) + [3, 1] + + If *window_size* is given, then the *pred* function will be called with + that many items. This enables searching for sub-sequences: + + >>> iterable = [0, 1, 2, 3, 0, 1, 2, 3, 0, 1, 2, 3] + >>> pred = lambda *args: args == (1, 2, 3) + >>> list(rlocate(iterable, pred=pred, window_size=3)) + [9, 5, 1] + + Beware, this function won't return anything for infinite iterables. + If *iterable* is reversible, ``rlocate`` will reverse it and search from + the right. Otherwise, it will search from the left and return the results + in reverse order. + + See :func:`locate` to for other example applications. + + """ + if window_size is None: + try: + len_iter = len(iterable) + return (len_iter - i - 1 for i in locate(reversed(iterable), pred)) + except TypeError: + pass + + return reversed(list(locate(iterable, pred, window_size))) + + +def replace(iterable, pred, substitutes, count=None, window_size=1): + """Yield the items from *iterable*, replacing the items for which *pred* + returns ``True`` with the items from the iterable *substitutes*. + + >>> iterable = [1, 1, 0, 1, 1, 0, 1, 1] + >>> pred = lambda x: x == 0 + >>> substitutes = (2, 3) + >>> list(replace(iterable, pred, substitutes)) + [1, 1, 2, 3, 1, 1, 2, 3, 1, 1] + + If *count* is given, the number of replacements will be limited: + + >>> iterable = [1, 1, 0, 1, 1, 0, 1, 1, 0] + >>> pred = lambda x: x == 0 + >>> substitutes = [None] + >>> list(replace(iterable, pred, substitutes, count=2)) + [1, 1, None, 1, 1, None, 1, 1, 0] + + Use *window_size* to control the number of items passed as arguments to + *pred*. This allows for locating and replacing subsequences. + + >>> iterable = [0, 1, 2, 5, 0, 1, 2, 5] + >>> window_size = 3 + >>> pred = lambda *args: args == (0, 1, 2) # 3 items passed to pred + >>> substitutes = [3, 4] # Splice in these items + >>> list(replace(iterable, pred, substitutes, window_size=window_size)) + [3, 4, 5, 3, 4, 5] + + *pred* may receive fewer than *window_size* arguments at the end of + the iterable and should be able to handle this. + + """ + if window_size < 1: + raise ValueError('window_size must be at least 1') + + # Save the substitutes iterable, since it's used more than once + substitutes = tuple(substitutes) + + # Add padding such that the number of windows matches the length of the + # iterable + it = chain(iterable, repeat(_marker, window_size - 1)) + windows = windowed(it, window_size) + + n = 0 + for w in windows: + # Strip any _marker padding so pred never sees internal sentinels. + # Near the end of the iterable, pred will receive fewer arguments. + args = tuple(x for x in w if x is not _marker) + + # If the current window matches our predicate (and we haven't hit + # our maximum number of replacements), splice in the substitutes + # and then consume the following windows that overlap with this one. + # For example, if the iterable is (0, 1, 2, 3, 4...) + # and the window size is 2, we have (0, 1), (1, 2), (2, 3)... + # If the predicate matches on (0, 1), we need to zap (0, 1) and (1, 2) + if args and pred(*args): + if (count is None) or (n < count): + n += 1 + yield from substitutes + consume(windows, window_size - 1) + continue + + # If there was no match (or we've reached the replacement limit), + # yield the first item from the window. + if args: + yield args[0] + + +def partitions(iterable): + """Yield all possible order-preserving partitions of *iterable*. + + >>> iterable = 'abc' + >>> for part in partitions(iterable): + ... print([''.join(p) for p in part]) + ['abc'] + ['a', 'bc'] + ['ab', 'c'] + ['a', 'b', 'c'] + + This is unrelated to :func:`partition`. + + """ + sequence = list(iterable) + n = len(sequence) + for i in powerset(range(1, n)): + yield [sequence[i:j] for i, j in zip((0,) + i, i + (n,))] + + +def set_partitions(iterable, k=None, min_size=None, max_size=None): + """ + Yield the set partitions of *iterable* into *k* parts. Set partitions are + not order-preserving. + + >>> iterable = 'abc' + >>> for part in set_partitions(iterable, 2): + ... print([''.join(p) for p in part]) + ['a', 'bc'] + ['ab', 'c'] + ['b', 'ac'] + + + If *k* is not given, every set partition is generated. + + >>> iterable = 'abc' + >>> for part in set_partitions(iterable): + ... print([''.join(p) for p in part]) + ['abc'] + ['a', 'bc'] + ['ab', 'c'] + ['b', 'ac'] + ['a', 'b', 'c'] + + if *min_size* and/or *max_size* are given, the minimum and/or maximum size + per block in partition is set. + + >>> iterable = 'abc' + >>> for part in set_partitions(iterable, min_size=2): + ... print([''.join(p) for p in part]) + ['abc'] + >>> for part in set_partitions(iterable, max_size=2): + ... print([''.join(p) for p in part]) + ['a', 'bc'] + ['ab', 'c'] + ['b', 'ac'] + ['a', 'b', 'c'] + + """ + L = list(iterable) + n = len(L) + if k is not None: + if k < 1: + raise ValueError( + "Can't partition in a negative or zero number of groups" + ) + elif k > n: + return + + min_size = min_size if min_size is not None else 0 + max_size = max_size if max_size is not None else n + if min_size > max_size: + return + + def set_partitions_helper(L, k): + n = len(L) + if k == 1: + yield [L] + elif n == k: + yield [[s] for s in L] + else: + e, *M = L + for p in set_partitions_helper(M, k - 1): + yield [[e], *p] + for p in set_partitions_helper(M, k): + for i in range(len(p)): + yield p[:i] + [[e] + p[i]] + p[i + 1 :] + + if k is None: + for k in range(1, n + 1): + yield from filter( + lambda z: all(min_size <= len(bk) <= max_size for bk in z), + set_partitions_helper(L, k), + ) + else: + yield from filter( + lambda z: all(min_size <= len(bk) <= max_size for bk in z), + set_partitions_helper(L, k), + ) + + +class time_limited: + """ + Yield items from *iterable* until *limit_seconds* have passed. + If the time limit expires before all items have been yielded, the + ``timed_out`` parameter will be set to ``True``. + + >>> from time import sleep + >>> def generator(): + ... yield 1 + ... yield 2 + ... sleep(0.2) + ... yield 3 + >>> iterable = time_limited(0.1, generator()) + >>> list(iterable) + [1, 2] + >>> iterable.timed_out + True + + Note that the time is checked before each item is yielded, and iteration + stops if the time elapsed is greater than *limit_seconds*. If your time + limit is 1 second, but it takes 2 seconds to generate the first item from + the iterable, the function will run for 2 seconds and not yield anything. + As a special case, when *limit_seconds* is zero, the iterator never + returns anything. + + """ + + def __init__(self, limit_seconds, iterable): + if limit_seconds < 0: + raise ValueError('limit_seconds must be positive') + self.limit_seconds = limit_seconds + self._iterator = iter(iterable) + self._start_time = monotonic() + self.timed_out = False + + def __iter__(self): + return self + + def __next__(self): + if self.limit_seconds == 0: + self.timed_out = True + raise StopIteration + item = next(self._iterator) + if monotonic() - self._start_time > self.limit_seconds: + self.timed_out = True + raise StopIteration + + return item + + +def only(iterable, default=None, too_long=None): + """If *iterable* has only one item, return it. + If it has zero items, return *default*. + If it has more than one item, raise the exception given by *too_long*, + which is ``ValueError`` by default. + + >>> only([], default='missing') + 'missing' + >>> only([1]) + 1 + >>> only([1, 2]) # doctest: +IGNORE_EXCEPTION_DETAIL + Traceback (most recent call last): + ... + ValueError: Expected exactly one item in iterable, but got 1, 2, + and perhaps more.' + >>> only([1, 2], too_long=TypeError) # doctest: +IGNORE_EXCEPTION_DETAIL + Traceback (most recent call last): + ... + TypeError + + Note that :func:`only` attempts to advance *iterable* twice to ensure there + is only one item. See :func:`spy` or :func:`peekable` to check + iterable contents less destructively. + + """ + iterator = iter(iterable) + for first in iterator: + for second in iterator: + if too_long is not None: + raise too_long + raise ValueError( + f'Expected exactly one item in iterable, but got {first!r}, ' + f'{second!r}, and perhaps more.' + ) + return first + return default + + +def ichunked(iterable, n): + """Break *iterable* into sub-iterables with *n* elements each. + :func:`ichunked` is like :func:`chunked`, but it yields iterables + instead of lists. + + If the sub-iterables are read in order, the elements of *iterable* + won't be stored in memory. + If they are read out of order, :func:`itertools.tee` is used to cache + elements as necessary. + + >>> from itertools import count + >>> all_chunks = ichunked(count(), 4) + >>> c_1, c_2, c_3 = next(all_chunks), next(all_chunks), next(all_chunks) + >>> list(c_2) # c_1's elements have been cached; c_3's haven't been + [4, 5, 6, 7] + >>> list(c_1) + [0, 1, 2, 3] + >>> list(c_3) + [8, 9, 10, 11] + + """ + if n is not None and n < 0: + raise ValueError('n must be at least 0') + if n == 0: + return + + iterator = iter(iterable) + for first in iterator: + rest = islice(iterator, n - 1) + cache, cacher = tee(rest) + yield chain([first], rest, cache) + consume(cacher) + + +def iequals(*iterables): + """Return ``True`` if all given *iterables* are equal to each other, + which means that they contain the same elements in the same order. + + The function is useful for comparing iterables of different data types + or iterables that do not support equality checks. + + >>> iequals("abc", ['a', 'b', 'c'], ('a', 'b', 'c'), iter("abc")) + True + + >>> iequals("abc", "acb") + False + + Not to be confused with :func:`all_equal`, which checks whether all + elements of iterable are equal to each other. + + """ + try: + return all(map(all_equal, zip(*iterables, strict=True))) + except ValueError: + return False + + +def distinct_combinations(iterable, r): + """Yield the distinct combinations of *r* items taken from *iterable*. + + >>> list(distinct_combinations([0, 0, 1], 2)) + [(0, 0), (0, 1)] + + Equivalent to ``set(combinations(iterable))``, except duplicates are not + generated and thrown away. For larger input sequences this is much more + efficient. + + """ + if r < 0: + raise ValueError('r must be non-negative') + elif r == 0: + yield () + return + pool = tuple(iterable) + generators = [unique_everseen(enumerate(pool), key=itemgetter(1))] + current_combo = [None] * r + level = 0 + while generators: + try: + cur_idx, p = next(generators[-1]) + except StopIteration: + generators.pop() + level -= 1 + continue + current_combo[level] = p + if level + 1 == r: + yield tuple(current_combo) + else: + generators.append( + unique_everseen( + enumerate(pool[cur_idx + 1 :], cur_idx + 1), + key=itemgetter(1), + ) + ) + level += 1 + + +def filter_except(validator, iterable, *exceptions): + """Yield the items from *iterable* for which the *validator* function does + not raise one of the specified *exceptions*. + + *validator* is called for each item in *iterable*. + It should be a function that accepts one argument and raises an exception + if that item is not valid. + + >>> iterable = ['1', '2', 'three', '4', None] + >>> list(filter_except(int, iterable, ValueError, TypeError)) + ['1', '2', '4'] + + If an exception other than one given by *exceptions* is raised by + *validator*, it is raised like normal. + """ + for item in iterable: + try: + validator(item) + except exceptions: + pass + else: + yield item + + +def map_except(function, iterable, *exceptions): + """Transform each item from *iterable* with *function* and yield the + result, unless *function* raises one of the specified *exceptions*. + + *function* is called to transform each item in *iterable*. + It should accept one argument. + + >>> iterable = ['1', '2', 'three', '4', None] + >>> list(map_except(int, iterable, ValueError, TypeError)) + [1, 2, 4] + + If an exception other than one given by *exceptions* is raised by + *function*, it is raised like normal. + """ + for item in iterable: + try: + yield function(item) + except exceptions: + pass + + +def map_if(iterable, pred, func, func_else=None): + """Evaluate each item from *iterable* using *pred*. If the result is + equivalent to ``True``, transform the item with *func* and yield it. + Otherwise, transform the item with *func_else* and yield it. + + *pred*, *func*, and *func_else* should each be functions that accept + one argument. By default, *func_else* is the identity function. + + >>> from math import sqrt + >>> iterable = list(range(-5, 5)) + >>> iterable + [-5, -4, -3, -2, -1, 0, 1, 2, 3, 4] + >>> list(map_if(iterable, lambda x: x > 3, lambda x: 'toobig')) + [-5, -4, -3, -2, -1, 0, 1, 2, 3, 'toobig'] + >>> list(map_if(iterable, lambda x: x >= 0, + ... lambda x: f'{sqrt(x):.2f}', lambda x: None)) + [None, None, None, None, None, '0.00', '1.00', '1.41', '1.73', '2.00'] + """ + + if func_else is None: + for item in iterable: + yield func(item) if pred(item) else item + + else: + for item in iterable: + yield func(item) if pred(item) else func_else(item) + + +def _sample_unweighted(iterator, k, strict): + # Algorithm L in the 1994 paper by Kim-Hung Li: + # "Reservoir-Sampling Algorithms of Time Complexity O(n(1+log(N/n)))". + + reservoir = list(islice(iterator, k)) + if strict and len(reservoir) < k: + raise ValueError('Sample larger than population') + W = 1.0 + + with suppress(StopIteration): + while True: + W *= random() ** (1 / k) + skip = floor(log(random()) / log1p(-W)) + element = next(islice(iterator, skip, None)) + reservoir[randrange(k)] = element + + shuffle(reservoir) + return reservoir + + +def _sample_weighted(iterator, k, weights, strict): + # Implementation of "A-ExpJ" from the 2006 paper by Efraimidis et al. : + # "Weighted random sampling with a reservoir". + + # Log-transform for numerical stability for weights that are small/large + weight_keys = (log(random()) / weight for weight in weights) + + # Fill up the reservoir (collection of samples) with the first `k` + # weight-keys and elements, then heapify the list. + reservoir = take(k, zip(weight_keys, iterator)) + if strict and len(reservoir) < k: + raise ValueError('Sample larger than population') + + heapify(reservoir) + + # The number of jumps before changing the reservoir is a random variable + # with an exponential distribution. Sample it using random() and logs. + smallest_weight_key, _ = reservoir[0] + weights_to_skip = log(random()) / smallest_weight_key + + for weight, element in zip(weights, iterator): + if weight >= weights_to_skip: + # The notation here is consistent with the paper, but we store + # the weight-keys in log-space for better numerical stability. + smallest_weight_key, _ = reservoir[0] + t_w = exp(weight * smallest_weight_key) + r_2 = uniform(t_w, 1) # generate U(t_w, 1) + weight_key = log(r_2) / weight + heapreplace(reservoir, (weight_key, element)) + smallest_weight_key, _ = reservoir[0] + weights_to_skip = log(random()) / smallest_weight_key + else: + weights_to_skip -= weight + + ret = [element for weight_key, element in reservoir] + shuffle(ret) + return ret + + +def _sample_counted(population, k, counts, strict): + element = None + remaining = 0 + + def feed(i): + # Advance *i* steps ahead and consume an element + nonlocal element, remaining + + while i + 1 > remaining: + i = i - remaining + element = next(population) + remaining = next(counts) + remaining -= i + 1 + return element + + with suppress(StopIteration): + reservoir = [] + for _ in range(k): + reservoir.append(feed(0)) + + if strict and len(reservoir) < k: + raise ValueError('Sample larger than population') + + with suppress(StopIteration): + W = 1.0 + while True: + W *= random() ** (1 / k) + skip = floor(log(random()) / log1p(-W)) + element = feed(skip) + reservoir[randrange(k)] = element + + shuffle(reservoir) + return reservoir + + +def sample(iterable, k, weights=None, *, counts=None, strict=False): + """Return a *k*-length list of elements chosen (without replacement) + from the *iterable*. + + Similar to :func:`random.sample`, but works on inputs that aren't + indexable (such as sets and dictionaries) and on inputs where the + size isn't known in advance (such as generators). + + >>> iterable = range(100) + >>> sample(iterable, 5) # doctest: +SKIP + [81, 60, 96, 16, 4] + + For iterables with repeated elements, you may supply *counts* to + indicate the repeats. + + >>> iterable = ['a', 'b'] + >>> counts = [3, 4] # Equivalent to 'a', 'a', 'a', 'b', 'b', 'b', 'b' + >>> sample(iterable, k=3, counts=counts) # doctest: +SKIP + ['a', 'a', 'b'] + + An iterable with *weights* may be given: + + >>> iterable = range(100) + >>> weights = (i * i + 1 for i in range(100)) + >>> sampled = sample(iterable, 5, weights=weights) # doctest: +SKIP + [79, 67, 74, 66, 78] + + Weighted selections are made without replacement. + After an element is selected, it is removed from the pool and the + relative weights of the other elements increase (this + does not match the behavior of :func:`random.sample`'s *counts* + parameter). Note that *weights* may not be used with *counts*. + + If the length of *iterable* is less than *k*, + ``ValueError`` is raised if *strict* is ``True`` and + all elements are returned (in shuffled order) if *strict* is ``False``. + + By default, the `Algorithm L `__ reservoir sampling + technique is used. When *weights* are provided, + `Algorithm A-ExpJ `__ is used instead. + + Notes on reproducibility: + + * The algorithms rely on inexact floating-point functions provided + by the underlying math library (e.g. ``log``, ``log1p``, and ``pow``). + Those functions can `produce slightly different results + `_ on + different builds. Accordingly, selections can vary across builds + even for the same seed. + + * The algorithms loop over the input and make selections based on + ordinal position, so selections from unordered collections (such as + sets) won't reproduce across sessions on the same platform using the + same seed. For example, this won't reproduce:: + + >> seed(8675309) + >> sample(set('abcdefghijklmnopqrstuvwxyz'), 10) + ['c', 'p', 'e', 'w', 's', 'a', 'j', 'd', 'n', 't'] + + """ + iterator = iter(iterable) + + if k < 0: + raise ValueError('k must be non-negative') + + if k == 0: + return [] + + if weights is not None and counts is not None: + raise TypeError('weights and counts are mutually exclusive') + + elif weights is not None: + weights = iter(weights) + return _sample_weighted(iterator, k, weights, strict) + + elif counts is not None: + counts = iter(counts) + return _sample_counted(iterator, k, counts, strict) + + else: + return _sample_unweighted(iterator, k, strict) + + +def is_sorted(iterable, key=None, reverse=False, strict=False): + """Returns ``True`` if the items of iterable are in sorted order, and + ``False`` otherwise. *key* and *reverse* have the same meaning that they do + in the built-in :func:`sorted` function. + + >>> is_sorted(['1', '2', '3', '4', '5'], key=int) + True + >>> is_sorted([5, 4, 3, 1, 2], reverse=True) + False + + If *strict*, tests for strict sorting, that is, returns ``False`` if equal + elements are found: + + >>> is_sorted([1, 2, 2]) + True + >>> is_sorted([1, 2, 2], strict=True) + False + + The function returns ``False`` after encountering the first out-of-order + item, which means it may produce results that differ from the built-in + :func:`sorted` function for objects with unusual comparison dynamics + (like ``math.nan``). If there are no out-of-order items, the iterable is + exhausted. + """ + it = iterable if (key is None) else map(key, iterable) + a, b = tee(it) + next(b, None) + if reverse: + b, a = a, b + return all(map(lt, a, b)) if strict else not any(map(lt, b, a)) + + +class AbortThread(BaseException): + pass + + +class callback_iter: + """Convert a function that uses callbacks to an iterator. + + .. deprecated:: 11.0.0 + Will be removed in a future major release. + + Let *func* be a function that takes a `callback` keyword argument. + For example: + + >>> def func(callback=None): + ... for i, c in [(1, 'a'), (2, 'b'), (3, 'c')]: + ... if callback: + ... callback(i, c) + ... return 4 + + + Use ``with callback_iter(func)`` to get an iterator over the parameters + that are delivered to the callback. + + >>> with callback_iter(func) as it: + ... for args, kwargs in it: + ... print(args) + (1, 'a') + (2, 'b') + (3, 'c') + + The function will be called in a background thread. The ``done`` property + indicates whether it has completed execution. + + >>> it.done + True + + If it completes successfully, its return value will be available + in the ``result`` property. + + >>> it.result + 4 + + Notes: + + * If the function uses some keyword argument besides ``callback``, supply + *callback_kwd*. + * If it finished executing, but raised an exception, accessing the + ``result`` property will raise the same exception. + * If it hasn't finished executing, accessing the ``result`` + property from within the ``with`` block will raise ``RuntimeError``. + * If it hasn't finished executing, accessing the ``result`` property from + outside the ``with`` block will raise a + ``more_itertools.AbortThread`` exception. + * Provide *wait_seconds* to adjust how frequently the it is polled for + output. + + """ + + def __init__(self, func, callback_kwd='callback', wait_seconds=0.1): + self._func = func + self._callback_kwd = callback_kwd + self._aborted = False + self._future = None + self._wait_seconds = wait_seconds + + # Lazily import concurrent.future + self._module = __import__('concurrent.futures').futures + self._executor = self._module.ThreadPoolExecutor(max_workers=1) + self._iterator = self._reader() + + def __enter__(self): + return self + + def __exit__(self, exc_type, exc_value, traceback): + self._aborted = True + self._executor.shutdown() + + def __iter__(self): + return self + + def __next__(self): + return next(self._iterator) + + @property + def done(self): + if self._future is None: + return False + return self._future.done() + + @property + def result(self): + if self._future: + try: + return self._future.result(timeout=0) + except self._module.TimeoutError: + pass + + raise RuntimeError('Function has not yet completed') + + def _reader(self): + q = Queue() + + def callback(*args, **kwargs): + if self._aborted: + raise AbortThread('canceled by user') + + q.put((args, kwargs)) + + self._future = self._executor.submit( + self._func, **{self._callback_kwd: callback} + ) + + while True: + try: + item = q.get(timeout=self._wait_seconds) + except Empty: + pass + else: + q.task_done() + yield item + + if self._future.done(): + break + + remaining = [] + while True: + try: + item = q.get_nowait() + except Empty: + break + else: + q.task_done() + remaining.append(item) + q.join() + yield from remaining + + +def windowed_complete(iterable, n): + """ + Yield ``(beginning, middle, end)`` tuples, where: + + * Each ``middle`` has *n* items from *iterable* + * Each ``beginning`` has the items before the ones in ``middle`` + * Each ``end`` has the items after the ones in ``middle`` + + >>> iterable = range(7) + >>> n = 3 + >>> for beginning, middle, end in windowed_complete(iterable, n): + ... print(beginning, middle, end) + () (0, 1, 2) (3, 4, 5, 6) + (0,) (1, 2, 3) (4, 5, 6) + (0, 1) (2, 3, 4) (5, 6) + (0, 1, 2) (3, 4, 5) (6,) + (0, 1, 2, 3) (4, 5, 6) () + + Note that *n* must be at least 0 and most equal to the length of + *iterable*. + + This function will exhaust the iterable and may require significant + storage. + """ + if n < 0: + raise ValueError('n must be >= 0') + + seq = tuple(iterable) + size = len(seq) + + if n > size: + raise ValueError('n must be <= len(seq)') + + for i in range(size - n + 1): + beginning = seq[:i] + middle = seq[i : i + n] + end = seq[i + n :] + yield beginning, middle, end + + +def all_unique(iterable, key=None): + """ + Returns ``True`` if all the elements of *iterable* are unique (no two + elements are equal). + + >>> all_unique('ABCB') + False + + If a *key* function is specified, it will be used to make comparisons. + + >>> all_unique('ABCb') + True + >>> all_unique('ABCb', str.lower) + False + + The function returns as soon as the first non-unique element is + encountered. Iterables with a mix of hashable and unhashable items can + be used, but the function will be slower for unhashable items. + """ + seenset = set() + seenset_add = seenset.add + seenlist = [] + seenlist_add = seenlist.append + for element in map(key, iterable) if key else iterable: + try: + if element in seenset: + return False + seenset_add(element) + except TypeError: + if element in seenlist: + return False + seenlist_add(element) + return True + + +def nth_product(index, *iterables, repeat=1): + """Equivalent to ``list(product(*iterables, repeat=repeat))[index]``. + + The products of *iterables* can be ordered lexicographically. + :func:`nth_product` computes the product at sort position *index* without + computing the previous products. + + >>> nth_product(8, range(2), range(2), range(2), range(2)) + (1, 0, 0, 0) + + The *repeat* keyword argument specifies the number of repetitions + of the iterables. The above example is equivalent to:: + + >>> nth_product(8, range(2), repeat=4) + (1, 0, 0, 0) + + ``IndexError`` will be raised if the given *index* is invalid. + """ + pools = tuple(map(tuple, reversed(iterables))) * repeat + ns = tuple(map(len, pools)) + + c = prod(ns) + + if index < 0: + index += c + if not 0 <= index < c: + raise IndexError + + result = [] + for pool, n in zip(pools, ns): + result.append(pool[index % n]) + index //= n + + return tuple(reversed(result)) + + +def nth_permutation(iterable, r, index): + """Equivalent to ``list(permutations(iterable, r))[index]``` + + The subsequences of *iterable* that are of length *r* where order is + important can be ordered lexicographically. :func:`nth_permutation` + computes the subsequence at sort position *index* directly, without + computing the previous subsequences. + + >>> nth_permutation('ghijk', 2, 5) + ('h', 'i') + + ``ValueError`` will be raised If *r* is negative. + ``IndexError`` will be raised if the given *index* is invalid. + """ + pool = list(iterable) + n = len(pool) + if r is None: + r = n + c = perm(n, r) + + if index < 0: + index += c + if not 0 <= index < c: + raise IndexError + + result = [0] * r + q = index + for d in range(n - r + 1, n + 1): + q, i = divmod(q, d) + result[n - d] = i + if q == 0: + break + + return tuple(map(pool.pop, result)) + + +def nth_combination_with_replacement(iterable, r, index): + """Equivalent to + ``list(combinations_with_replacement(iterable, r))[index]``. + + + The subsequences with repetition of *iterable* that are of length *r* can + be ordered lexicographically. :func:`nth_combination_with_replacement` + computes the subsequence at sort position *index* directly, without + computing the previous subsequences with replacement. + + >>> nth_combination_with_replacement(range(5), 3, 5) + (0, 1, 1) + + ``ValueError`` will be raised If *r* is negative. + ``IndexError`` will be raised if the given *index* is invalid. + """ + pool = tuple(iterable) + n = len(pool) + if r < 0: + raise ValueError + c = comb(n + r - 1, r) if n else 0 if r else 1 + + if index < 0: + index += c + if not 0 <= index < c: + raise IndexError + + result = [] + i = 0 + while r: + r -= 1 + while n >= 0: + num_combs = comb(n + r - 1, r) + if index < num_combs: + break + n -= 1 + i += 1 + index -= num_combs + result.append(pool[i]) + + return tuple(result) + + +def value_chain(*args): + """Yield all arguments passed to the function in the same order in which + they were passed. If an argument itself is iterable then iterate over its + values. + + >>> list(value_chain(1, 2, 3, [4, 5, 6])) + [1, 2, 3, 4, 5, 6] + + Binary and text strings are not considered iterable and are emitted + as-is: + + >>> list(value_chain('12', '34', ['56', '78'])) + ['12', '34', '56', '78'] + + Pre- or postpend a single element to an iterable: + + >>> list(value_chain(1, [2, 3, 4, 5, 6])) + [1, 2, 3, 4, 5, 6] + >>> list(value_chain([1, 2, 3, 4, 5], 6)) + [1, 2, 3, 4, 5, 6] + + Multiple levels of nesting are not flattened. + + """ + scalar_types = (str, bytes) + for value in args: + if isinstance(value, scalar_types): + yield value + continue + try: + it = iter(value) + except TypeError: + yield value + else: + yield from it + + +def product_index(element, *iterables, repeat=1): + """Equivalent to ``list(product(*iterables, repeat=repeat)).index(tuple(element))`` + + The products of *iterables* can be ordered lexicographically. + :func:`product_index` computes the first index of *element* without + computing the previous products. + + >>> product_index([8, 2], range(10), range(5)) + 42 + + The *repeat* keyword argument specifies the number of repetitions + of the iterables:: + + >>> product_index([8, 0, 7], range(10), repeat=3) + 807 + + ``ValueError`` will be raised if the given *element* isn't in the product + of *args*. + """ + elements = tuple(element) + pools = tuple(map(tuple, iterables)) * repeat + if len(elements) != len(pools): + raise ValueError('element is not a product of args') + + index = 0 + for elem, pool in zip(elements, pools): + index = index * len(pool) + pool.index(elem) + return index + + +def combination_index(element, iterable): + """Equivalent to ``list(combinations(iterable, r)).index(element)`` + + The subsequences of *iterable* that are of length *r* can be ordered + lexicographically. :func:`combination_index` computes the index of the + first *element*, without computing the previous combinations. + + >>> combination_index('adf', 'abcdefg') + 10 + + ``ValueError`` will be raised if the given *element* isn't one of the + combinations of *iterable*. + """ + element = enumerate(element) + k, y = next(element, (None, None)) + if k is None: + return 0 + + indexes = [] + pool = enumerate(iterable) + for n, x in pool: + if x == y: + indexes.append(n) + tmp, y = next(element, (None, None)) + if tmp is None: + break + else: + k = tmp + else: + raise ValueError('element is not a combination of iterable') + + n, _ = last(pool, default=(n, None)) + + index = 1 + for i, j in enumerate(reversed(indexes), start=1): + j = n - j + if i <= j: + index += comb(j, i) + + return comb(n + 1, k + 1) - index + + +def combination_with_replacement_index(element, iterable): + """Equivalent to + ``list(combinations_with_replacement(iterable, r)).index(element)`` + + The subsequences with repetition of *iterable* that are of length *r* can + be ordered lexicographically. :func:`combination_with_replacement_index` + computes the index of the first *element*, without computing the previous + combinations with replacement. + + >>> combination_with_replacement_index('adf', 'abcdefg') + 20 + + ``ValueError`` will be raised if the given *element* isn't one of the + combinations with replacement of *iterable*. + """ + element = tuple(element) + r = len(element) + pool = tuple(iterable) + n = len(pool) + + occupations = [0] * n + try: + i = 0 + for e in element: + i = pool.index(e, i) + occupations[i] += 1 + except ValueError: + raise ValueError( + 'element is not a combination with replacement of iterable' + ) + + index = 0 + cumulative_sum = 0 + for k in range(1, n): + cumulative_sum += occupations[k - 1] + j = (r - cumulative_sum) + (n - k) - 1 + i = n - k + if i <= j: + index += comb(j, i) + + return index + + +def permutation_index(element, iterable): + """Equivalent to ``list(permutations(iterable, r)).index(element)``` + + The subsequences of *iterable* that are of length *r* where order is + important can be ordered lexicographically. :func:`permutation_index` + computes the index of the first *element* directly, without computing + the previous permutations. + + >>> permutation_index([1, 3, 2], range(5)) + 19 + + ``ValueError`` will be raised if the given *element* isn't one of the + permutations of *iterable*. + """ + index = 0 + pool = list(iterable) + for i, x in zip(range(len(pool), -1, -1), element): + r = pool.index(x) + index = index * i + r + del pool[r] + + return index + + +class countable: + """Wrap *iterable* and keep a count of how many items have been consumed. + + The ``items_seen`` attribute starts at ``0`` and increments as the iterable + is consumed: + + >>> iterable = map(str, range(10)) + >>> it = countable(iterable) + >>> it.items_seen + 0 + >>> next(it), next(it) + ('0', '1') + >>> list(it) + ['2', '3', '4', '5', '6', '7', '8', '9'] + >>> it.items_seen + 10 + """ + + def __init__(self, iterable): + self._iterator = iter(iterable) + self.items_seen = 0 + + def __iter__(self): + return self + + def __next__(self): + item = next(self._iterator) + self.items_seen += 1 + + return item + + +def chunked_even(iterable, n): + """Break *iterable* into lists of approximately length *n*. + Items are distributed such the lengths of the lists differ by at most + 1 item. + + >>> iterable = [1, 2, 3, 4, 5, 6, 7] + >>> n = 3 + >>> list(chunked_even(iterable, n)) # List lengths: 3, 2, 2 + [[1, 2, 3], [4, 5], [6, 7]] + >>> list(chunked(iterable, n)) # List lengths: 3, 3, 1 + [[1, 2, 3], [4, 5, 6], [7]] + + """ + iterator = iter(iterable) + + # Initialize a buffer to process the chunks while keeping + # some back to fill any underfilled chunks + min_buffer = (n - 1) * (n - 2) + buffer = list(islice(iterator, min_buffer)) + + # Append items until we have a completed chunk + for _ in islice(map(buffer.append, iterator), n, None, n): + yield buffer[:n] + del buffer[:n] + + # Check if any chunks need addition processing + if not buffer: + return + length = len(buffer) + + # Chunks are either size `full_size <= n` or `partial_size = full_size - 1` + q, r = divmod(length, n) + num_lists = q + (1 if r > 0 else 0) + q, r = divmod(length, num_lists) + full_size = q + (1 if r > 0 else 0) + partial_size = full_size - 1 + num_full = length - partial_size * num_lists + + # Yield chunks of full size + partial_start_idx = num_full * full_size + if full_size > 0: + for i in range(0, partial_start_idx, full_size): + yield buffer[i : i + full_size] + + # Yield chunks of partial size + if partial_size > 0: + for i in range(partial_start_idx, length, partial_size): + yield buffer[i : i + partial_size] + + +def zip_broadcast(*objects, scalar_types=(str, bytes), strict=False): + """A version of :func:`zip` that "broadcasts" any scalar + (i.e., non-iterable) items into output tuples. + + >>> iterable_1 = [1, 2, 3] + >>> iterable_2 = ['a', 'b', 'c'] + >>> scalar = '_' + >>> list(zip_broadcast(iterable_1, iterable_2, scalar)) + [(1, 'a', '_'), (2, 'b', '_'), (3, 'c', '_')] + + The *scalar_types* keyword argument determines what types are considered + scalar. It is set to ``(str, bytes)`` by default. Set it to ``None`` to + treat strings and byte strings as iterable: + + >>> list(zip_broadcast('abc', 0, 'xyz', scalar_types=None)) + [('a', 0, 'x'), ('b', 0, 'y'), ('c', 0, 'z')] + + If the *strict* keyword argument is ``True``, then + ``ValueError`` will be raised if any of the iterables have + different lengths. + """ + + def get_iterator(obj): + if scalar_types and isinstance(obj, scalar_types): + return None + try: + return iter(obj) + except TypeError: + return None + + size = len(objects) + if not size: + return + + new_item = [None] * size + iterables, iterable_positions = [], [] + for i, obj in enumerate(objects): + iterator = get_iterator(obj) + if iterator is None: + new_item[i] = obj + else: + iterables.append(iterator) + iterable_positions.append(i) + + if not iterables: + yield tuple(objects) + return + + for item in zip(*iterables, strict=strict): + for i, new_item[i] in zip(iterable_positions, item): + pass + yield tuple(new_item) + + +def unique_in_window(iterable, n, key=None): + """Yield the items from *iterable* that haven't been seen recently. + *n* is the size of the sliding window. + + >>> iterable = [0, 1, 0, 2, 3, 0] + >>> n = 3 + >>> list(unique_in_window(iterable, n)) + [0, 1, 2, 3, 0] + + The *key* function, if provided, will be used to determine uniqueness: + + >>> list(unique_in_window('abAcda', 3, key=lambda x: x.lower())) + ['a', 'b', 'c', 'd', 'a'] + + Updates a sliding window no larger than n and yields a value + if the item only occurs once in the updated window. + + When `n == 1`, *unique_in_window* is memoryless: + + >>> list(unique_in_window('aab', n=1)) + ['a', 'a', 'b'] + + The items in *iterable* must be hashable. + + """ + if n <= 0: + raise ValueError('n must be greater than 0') + + window = deque(maxlen=n) + counts = Counter() + use_key = key is not None + + for item in iterable: + if len(window) == n: + to_discard = window[0] + if counts[to_discard] == 1: + del counts[to_discard] + else: + counts[to_discard] -= 1 + + k = key(item) if use_key else item + if k not in counts: + yield item + counts[k] += 1 + window.append(k) + + +def duplicates_everseen(iterable, key=None): + """Yield duplicate elements after their first appearance. + + >>> list(duplicates_everseen('mississippi')) + ['s', 'i', 's', 's', 'i', 'p', 'i'] + >>> list(duplicates_everseen('AaaBbbCccAaa', str.lower)) + ['a', 'a', 'b', 'b', 'c', 'c', 'A', 'a', 'a'] + + This function is analogous to :func:`unique_everseen` and is subject to + the same performance considerations. + + If you would like each duplicate to only appear once, much like ``uniq -d`` + in the Unix shell or ``Itertools::duplicates`` from the Rust ``itertools`` + crate, pass the return value of this function into :func:`unique_everseen` + with the same ``key``. + + """ + seen_set = set() + seen_list = [] + use_key = key is not None + + for element in iterable: + k = key(element) if use_key else element + try: + if k not in seen_set: + seen_set.add(k) + else: + yield element + except TypeError: + if k not in seen_list: + seen_list.append(k) + else: + yield element + + +def duplicates_justseen(iterable, key=None): + """Yields serially-duplicate elements after their first appearance. + + >>> list(duplicates_justseen('mississippi')) + ['s', 's', 'p'] + >>> list(duplicates_justseen('AaaBbbCccAaa', str.lower)) + ['a', 'a', 'b', 'b', 'c', 'c', 'a', 'a'] + + This function is analogous to :func:`unique_justseen`. + + """ + return flatten(g for _, g in groupby(iterable, key) for _ in g) + + +def classify_unique(iterable, key=None): + """Classify each element in terms of its uniqueness. + + For each element in the input iterable, return a 3-tuple consisting of: + + 1. The element itself + 2. ``False`` if the element is equal to the one preceding it in the input, + ``True`` otherwise (i.e. the equivalent of :func:`unique_justseen`) + 3. ``False`` if this element has been seen anywhere in the input before, + ``True`` otherwise (i.e. the equivalent of :func:`unique_everseen`) + + >>> list(classify_unique('otto')) # doctest: +NORMALIZE_WHITESPACE + [('o', True, True), + ('t', True, True), + ('t', False, False), + ('o', True, False)] + + This function is analogous to :func:`unique_everseen` and is subject to + the same performance considerations. + + """ + seen_set = set() + seen_list = [] + use_key = key is not None + previous = None + + for i, element in enumerate(iterable): + k = key(element) if use_key else element + is_unique_justseen = not i or previous != k + previous = k + is_unique_everseen = False + try: + if k not in seen_set: + seen_set.add(k) + is_unique_everseen = True + except TypeError: + if k not in seen_list: + seen_list.append(k) + is_unique_everseen = True + yield element, is_unique_justseen, is_unique_everseen + + +def minmax(iterable_or_value, *others, key=None, default=_marker): + """Returns both the smallest and largest items from an iterable + or from two or more arguments. + + >>> minmax([3, 1, 5]) + (1, 5) + + >>> minmax(4, 2, 6) + (2, 6) + + If a *key* function is provided, it will be used to transform the input + items for comparison. + + >>> minmax([5, 30], key=str) # '30' sorts before '5' + (30, 5) + + If a *default* value is provided, it will be returned if there are no + input items. + + >>> minmax([], default=(0, 0)) + (0, 0) + + Otherwise ``ValueError`` is raised. + + This function makes a single pass over the input elements and takes care to + minimize the number of comparisons made during processing. + + Note that unlike the builtin ``max`` function, which always returns the first + item with the maximum value, this function may return another item when there are + ties. + + This function is based on the + `recipe `__ by + Raymond Hettinger. + """ + iterable = (iterable_or_value, *others) if others else iterable_or_value + + it = iter(iterable) + + try: + lo = hi = next(it) + except StopIteration as exc: + if default is _marker: + raise ValueError( + '`minmax()` argument is an empty iterable. ' + 'Provide a `default` value to suppress this error.' + ) from exc + return default + + # Different branches depending on the presence of key. This saves a lot + # of unimportant copies which would slow the "key=None" branch + # significantly down. + if key is None: + for x, y in zip_longest(it, it, fillvalue=lo): + if y < x: + if y < lo: + lo = y + if hi < x: + hi = x + else: + if x < lo: + lo = x + if hi < y: + hi = y + + else: + lo_key = hi_key = key(lo) + + for x, y in zip_longest(it, it, fillvalue=lo): + x_key, y_key = key(x), key(y) + + if y_key < x_key: + if y_key < lo_key: + lo, lo_key = y, y_key + if hi_key < x_key: + hi, hi_key = x, x_key + else: + if x_key < lo_key: + lo, lo_key = x, x_key + if hi_key < y_key: + hi, hi_key = y, y_key + + return lo, hi + + +def constrained_batches( + iterable, max_size, max_count=None, get_len=len, strict=True +): + """Yield batches of items from *iterable* with a combined size limited by + *max_size*. + + >>> iterable = [b'12345', b'123', b'12345678', b'1', b'1', b'12', b'1'] + >>> list(constrained_batches(iterable, 10)) + [(b'12345', b'123'), (b'12345678', b'1', b'1'), (b'12', b'1')] + + If a *max_count* is supplied, the number of items per batch is also + limited. It must be greater than zero: + + >>> iterable = [b'12345', b'123', b'12345678', b'1', b'1', b'12', b'1'] + >>> list(constrained_batches(iterable, 10, max_count = 2)) + [(b'12345', b'123'), (b'12345678', b'1'), (b'1', b'12'), (b'1',)] + + If a *get_len* function is supplied, use that instead of :func:`len` to + determine item size. + + If *strict* is ``True``, raise ``ValueError`` if any single item is bigger + than *max_size*. Otherwise, allow single items to exceed *max_size*. + """ + if max_size <= 0: + raise ValueError('maximum size must be greater than zero') + if max_count is not None and max_count <= 0: + raise ValueError('maximum count must be greater than zero') + + batch = [] + batch_size = 0 + batch_count = 0 + for item in iterable: + item_len = get_len(item) + if strict and item_len > max_size: + raise ValueError('item size exceeds maximum size') + + reached_count = batch_count == max_count + reached_size = item_len + batch_size > max_size + if batch_count and (reached_size or reached_count): + yield tuple(batch) + batch.clear() + batch_size = 0 + batch_count = 0 + + batch.append(item) + batch_size += item_len + batch_count += 1 + + if batch: + yield tuple(batch) + + +def gray_product(*iterables, repeat=1): + """Like :func:`itertools.product`, but return tuples in an order such + that only one element in the generated tuple changes from one iteration + to the next. + + >>> list(gray_product('AB','CD')) + [('A', 'C'), ('B', 'C'), ('B', 'D'), ('A', 'D')] + + The *repeat* keyword argument specifies the number of repetitions + of the iterables. For example, ``gray_product('AB', repeat=3)`` is + equivalent to ``gray_product('AB', 'AB', 'AB')``. + + This function consumes all of the input iterables before producing output. + If any of the input iterables have fewer than two items, ``ValueError`` + is raised. + + For information on the algorithm, see + `this section `__ + of Donald Knuth's *The Art of Computer Programming*. + """ + all_iterables = tuple(map(tuple, iterables)) * repeat + iterable_count = len(all_iterables) + for iterable in all_iterables: + if len(iterable) < 2: + raise ValueError("each iterable must have two or more items") + + # This is based on "Algorithm H" from section 7.2.1.1, page 20. + # a holds the indexes of the source iterables for the n-tuple to be yielded + # f is the array of "focus pointers" + # o is the array of "directions" + a = [0] * iterable_count + f = list(range(iterable_count + 1)) + o = [1] * iterable_count + while True: + yield tuple(all_iterables[i][a[i]] for i in range(iterable_count)) + j = f[0] + f[0] = 0 + if j == iterable_count: + break + a[j] = a[j] + o[j] + if a[j] == 0 or a[j] == len(all_iterables[j]) - 1: + o[j] = -o[j] + f[j] = f[j + 1] + f[j + 1] = j + 1 + + +def partial_product(*iterables, repeat=1): + """Yields tuples containing one item from each iterator, with subsequent + tuples changing a single item at a time by advancing each iterator until it + is exhausted. This sequence guarantees every value in each iterable is + output at least once without generating all possible combinations. + + This may be useful, for example, when testing an expensive function. + + >>> list(partial_product('AB', 'C', 'DEF')) + [('A', 'C', 'D'), ('B', 'C', 'D'), ('B', 'C', 'E'), ('B', 'C', 'F')] + + The *repeat* keyword argument specifies the number of repetitions + of the iterables. For example, ``partial_product('AB', repeat=3)`` is + equivalent to ``partial_product('AB', 'AB', 'AB')``. + """ + + all_iterables = tuple(map(tuple, iterables)) * repeat + iterators = tuple(map(iter, all_iterables)) + + try: + prod = [next(it) for it in iterators] + except StopIteration: + return + yield tuple(prod) + + for i, it in enumerate(iterators): + for prod[i] in it: + yield tuple(prod) + + +def takewhile_inclusive(predicate, iterable): + """A variant of :func:`takewhile` that yields one additional element. + + >>> list(takewhile_inclusive(lambda x: x < 5, [1, 4, 6, 4, 1])) + [1, 4, 6] + + :func:`takewhile` would return ``[1, 4]``. + """ + for x in iterable: + yield x + if not predicate(x): + break + + +def outer_product(func, xs, ys, *args, **kwargs): + """A generalized outer product that applies a binary function to all + pairs of items. Returns a 2D matrix with ``len(xs)`` rows and ``len(ys)`` + columns. + Also accepts ``*args`` and ``**kwargs`` that are passed to ``func``. + + Multiplication table: + + >>> from operator import mul + >>> list(outer_product(mul, range(1, 4), range(1, 6))) + [(1, 2, 3, 4, 5), (2, 4, 6, 8, 10), (3, 6, 9, 12, 15)] + + Cross tabulation: + + >>> xs = ['A', 'B', 'A', 'A', 'B', 'B', 'A', 'A', 'B', 'B'] + >>> ys = ['X', 'X', 'X', 'Y', 'Z', 'Z', 'Y', 'Y', 'Z', 'Z'] + >>> pair_counts = Counter(zip(xs, ys)) + >>> count_rows = lambda x, y: pair_counts[x, y] + >>> list(outer_product(count_rows, sorted(set(xs)), sorted(set(ys)))) + [(2, 3, 0), (1, 0, 4)] + + Usage with ``*args`` and ``**kwargs``: + + >>> animals = ['cat', 'wolf', 'mouse'] + >>> list(outer_product(min, animals, animals, key=len)) + [('cat', 'cat', 'cat'), ('cat', 'wolf', 'wolf'), ('cat', 'wolf', 'mouse')] + """ + ys = tuple(ys) + return batched( + starmap(lambda x, y: func(x, y, *args, **kwargs), product(xs, ys)), + n=len(ys), + ) + + +def iter_suppress(iterable, *exceptions): + """Yield each of the items from *iterable*. If the iteration raises one of + the specified *exceptions*, that exception will be suppressed and iteration + will stop. + + >>> from itertools import chain + >>> def breaks_at_five(x): + ... while True: + ... if x >= 5: + ... raise RuntimeError + ... yield x + ... x += 1 + >>> it_1 = iter_suppress(breaks_at_five(1), RuntimeError) + >>> it_2 = iter_suppress(breaks_at_five(2), RuntimeError) + >>> list(chain(it_1, it_2)) + [1, 2, 3, 4, 2, 3, 4] + """ + try: + yield from iterable + except exceptions: + return + + +def filter_map(func, iterable): + """Apply *func* to every element of *iterable*, yielding only those which + are not ``None``. + + >>> elems = ['1', 'a', '2', 'b', '3'] + >>> list(filter_map(lambda s: int(s) if s.isnumeric() else None, elems)) + [1, 2, 3] + """ + for x in iterable: + y = func(x) + if y is not None: + yield y + + +def powerset_of_sets(iterable, *, baseset=set): + """Yields all possible subsets of the iterable. + + >>> list(powerset_of_sets([1, 2, 3])) # doctest: +SKIP + [set(), {1}, {2}, {3}, {1, 2}, {1, 3}, {2, 3}, {1, 2, 3}] + >>> list(powerset_of_sets([1, 1, 0])) # doctest: +SKIP + [set(), {1}, {0}, {0, 1}] + + :func:`powerset_of_sets` takes care to minimize the number + of hash operations performed. + + The *baseset* parameter determines what kind of sets are + constructed, either *set* or *frozenset*. + """ + sets = tuple(dict.fromkeys(map(frozenset, zip(iterable)))) + union = baseset().union + return chain.from_iterable( + starmap(union, combinations(sets, r)) for r in range(len(sets) + 1) + ) + + +def join_mappings(**field_to_map): + """ + Joins multiple mappings together using their common keys. + + >>> user_scores = {'elliot': 50, 'claris': 60} + >>> user_times = {'elliot': 30, 'claris': 40} + >>> join_mappings(score=user_scores, time=user_times) + {'elliot': {'score': 50, 'time': 30}, 'claris': {'score': 60, 'time': 40}} + """ + ret = defaultdict(dict) + + for field_name, mapping in field_to_map.items(): + for key, value in mapping.items(): + ret[key][field_name] = value + + return dict(ret) + + +def _complex_sumprod(v1, v2): + """High precision sumprod() for complex numbers. + Used by :func:`dft` and :func:`idft`. + """ + + real = attrgetter('real') + imag = attrgetter('imag') + r1 = chain(map(real, v1), map(neg, map(imag, v1))) + r2 = chain(map(real, v2), map(imag, v2)) + i1 = chain(map(real, v1), map(imag, v1)) + i2 = chain(map(imag, v2), map(real, v2)) + return complex(_fsumprod(r1, r2), _fsumprod(i1, i2)) + + +def dft(xarr): + """Discrete Fourier Transform. *xarr* is a sequence of complex numbers. + Yields the components of the corresponding transformed output vector. + + >>> import cmath + >>> xarr = [1, 2-1j, -1j, -1+2j] # time domain + >>> Xarr = [2, -2-2j, -2j, 4+4j] # frequency domain + >>> magnitudes, phases = zip(*map(cmath.polar, Xarr)) + >>> all(map(cmath.isclose, dft(xarr), Xarr)) + True + + Inputs are restricted to numeric types that can add and multiply + with a complex number. This includes int, float, complex, and + Fraction, but excludes Decimal. + + See :func:`idft` for the inverse Discrete Fourier Transform. + """ + N = len(xarr) + roots_of_unity = [e ** (n / N * tau * -1j) for n in range(N)] + for k in range(N): + coeffs = [roots_of_unity[k * n % N] for n in range(N)] + yield _complex_sumprod(xarr, coeffs) + + +def idft(Xarr): + """Inverse Discrete Fourier Transform. *Xarr* is a sequence of + complex numbers. Yields the components of the corresponding + inverse-transformed output vector. + + >>> import cmath + >>> xarr = [1, 2-1j, -1j, -1+2j] # time domain + >>> Xarr = [2, -2-2j, -2j, 4+4j] # frequency domain + >>> all(map(cmath.isclose, idft(Xarr), xarr)) + True + + Inputs are restricted to numeric types that can add and multiply + with a complex number. This includes int, float, complex, and + Fraction, but excludes Decimal. + + See :func:`dft` for the Discrete Fourier Transform. + """ + N = len(Xarr) + roots_of_unity = [e ** (n / N * tau * 1j) for n in range(N)] + for k in range(N): + coeffs = [roots_of_unity[k * n % N] for n in range(N)] + yield _complex_sumprod(Xarr, coeffs) / N + + +def doublestarmap(func, iterable): + """Apply *func* to every item of *iterable* by dictionary unpacking + the item into *func*. + + The difference between :func:`itertools.starmap` and :func:`doublestarmap` + parallels the distinction between ``func(*a)`` and ``func(**a)``. + + >>> iterable = [{'a': 1, 'b': 2}, {'a': 40, 'b': 60}] + >>> list(doublestarmap(lambda a, b: a + b, iterable)) + [3, 100] + + ``TypeError`` will be raised if *func*'s signature doesn't match the + mapping contained in *iterable* or if *iterable* does not contain mappings. + """ + for item in iterable: + yield func(**item) + + +def _nth_prime_bounds(n): + """Bounds for the nth prime (counting from 1): lb < p_n < ub.""" + # At and above 688,383, the lb/ub spread is under 0.003 * p_n. + + if n < 1: + raise ValueError + + if n < 6: + return (n, 2.25 * n) + + # https://en.wikipedia.org/wiki/Prime-counting_function#Inequalities + upper_bound = n * log(n * log(n)) + lower_bound = upper_bound - n + if n >= 688_383: + upper_bound -= n * (1.0 - (log(log(n)) - 2.0) / log(n)) + + return lower_bound, upper_bound + + +def nth_prime(n, *, approximate=False): + """Return the nth prime (counting from 0). + + >>> nth_prime(0) + 2 + >>> nth_prime(100) + 547 + + If *approximate* is set to True, will return a prime close + to the nth prime. The estimation is much faster than computing + an exact result. + + >>> nth_prime(200_000_000, approximate=True) # Exact result is 4222234763 + 4217820427 + + """ + lb, ub = _nth_prime_bounds(n + 1) + + if not approximate or n <= 1_000_000: + return nth(sieve(ceil(ub)), n) + + # Search from the midpoint and return the first odd prime + odd = floor((lb + ub) / 2) | 1 + return first_true(count(odd, step=2), pred=is_prime) + + +def argmin(iterable, *, key=None): + """ + Index of the first occurrence of a minimum value in an iterable. + + >>> argmin('efghabcdijkl') + 4 + >>> argmin([3, 2, 1, 0, 4, 2, 1, 0]) + 3 + + For example, look up a label corresponding to the position + of a value that minimizes a cost function:: + + >>> def cost(x): + ... "Days for a wound to heal given a subject's age." + ... return x**2 - 20*x + 150 + ... + >>> labels = ['homer', 'marge', 'bart', 'lisa', 'maggie'] + >>> ages = [ 35, 30, 10, 9, 1 ] + + # Fastest healing family member + >>> labels[argmin(ages, key=cost)] + 'bart' + + # Age with fastest healing + >>> min(ages, key=cost) + 10 + + """ + if key is not None: + iterable = map(key, iterable) + return min(enumerate(iterable), key=itemgetter(1))[0] + + +def argmax(iterable, *, key=None): + """ + Index of the first occurrence of a maximum value in an iterable. + + >>> argmax('abcdefghabcd') + 7 + >>> argmax([0, 1, 2, 3, 3, 2, 1, 0]) + 3 + + For example, identify the best machine learning model:: + + >>> models = ['svm', 'random forest', 'knn', 'naïve bayes'] + >>> accuracy = [ 68, 61, 84, 72 ] + + # Most accurate model + >>> models[argmax(accuracy)] + 'knn' + + # Best accuracy + >>> max(accuracy) + 84 + + """ + if key is not None: + iterable = map(key, iterable) + return max(enumerate(iterable), key=itemgetter(1))[0] + + +def _extract_monotonic(iterator, indices): + 'Non-decreasing indices, lazily consumed' + num_read = 0 + for index in indices: + advance = index - num_read + try: + value = next(islice(iterator, advance, None)) + except ValueError: + if advance != -1 or index < 0: + raise ValueError(f'Invalid index: {index}') from None + except StopIteration: + raise IndexError(index) from None + else: + num_read += advance + 1 + yield value + + +def _extract_buffered(iterator, index_and_position): + 'Arbitrary index order, greedily consumed' + buffer = {} + iterator_position = -1 + next_to_emit = 0 + + for index, order in index_and_position: + advance = index - iterator_position + if advance: + try: + value = next(islice(iterator, advance - 1, None)) + except StopIteration: + raise IndexError(index) from None + iterator_position = index + + buffer[order] = value + + while next_to_emit in buffer: + yield buffer.pop(next_to_emit) + next_to_emit += 1 + + +def extract(iterable, indices, *, monotonic=False): + """Yield values at the specified indices. + + Example: + + >>> data = 'abcdefghijklmnopqrstuvwxyz' + >>> list(extract(data, [7, 4, 11, 11, 14])) + ['h', 'e', 'l', 'l', 'o'] + + The *iterable* is consumed lazily and can be infinite. + + When *monotonic* is false, the *indices* are consumed immediately + and must be finite. When *monotonic* is true, *indices* are consumed + lazily and can be infinite but must be non-decreasing. + + Raises ``IndexError`` if an index lies beyond the iterable. + Raises ``ValueError`` for a negative index or for a decreasing + index when *monotonic* is true. + """ + + iterator = iter(iterable) + indices = iter(indices) + + if monotonic: + return _extract_monotonic(iterator, indices) + + index_and_position = sorted(zip(indices, count())) + if index_and_position and index_and_position[0][0] < 0: + raise ValueError('Indices must be non-negative') + return _extract_buffered(iterator, index_and_position) + + +class serialize: + """Wrap a non-concurrent iterator with a lock to enforce sequential access. + + Applies a non-reentrant lock around calls to ``__next__``, allowing + iterator and generator instances to be shared by multiple consumer + threads. + """ + + __slots__ = ('_iterator', '_lock') + + def __init__(self, iterable): + self._iterator = iter(iterable) + self._lock = Lock() + + def __iter__(self): + return self + + def __next__(self): + with self._lock: + return next(self._iterator) + + def send(self, value, /): + """Send a value to a generator. + + Raises AttributeError if not a generator. + """ + with self._lock: + return self._iterator.send(value) + + def throw(self, *args): + """Call throw() on a generator. + + Raises AttributeError if not a generator. + """ + with self._lock: + return self._iterator.throw(*args) + + def close(self): + """Call close() on a generator. + + Raises AttributeError if not a generator. + """ + with self._lock: + return self._iterator.close() + + +def synchronized(func): + """Wrap an iterator-returning callable to make its iterators thread-safe. + + Existing itertools and more-itertools can be wrapped so that their + iterator instances are serialized. + + For example, ``itertools.count`` does not make thread-safe instances, + but that is easily fixed with:: + + atomic_counter = synchronized(itertools.count) + + Can also be used as a decorator for generator functions definitions + so that the generator instances are serialized:: + + @synchronized + def enumerate_and_timestamp(iterable): + for count, value in enumerate(iterable): + yield count, time_ns(), value + + """ + + @wraps(func) + def inner(*args, **kwargs): + iterator = func(*args, **kwargs) + return serialize(iterator) + + return inner + + +def concurrent_tee(iterable, n=2): + """Variant of itertools.tee() but with guaranteed threading semantics. + + Takes a non-threadsafe iterator as an input and creates concurrent + tee objects for other threads to have reliable independent copies of + the data stream. + + The new iterators are only thread-safe if consumed within a single thread. + To share just one of the new iterators across multiple threads, wrap it + with :func:`serialize`. + """ + + if n < 0: + raise ValueError + if n == 0: + return () + iterator = _concurrent_tee(iterable) + result = [iterator] + for _ in range(n - 1): + result.append(_concurrent_tee(iterator)) + return tuple(result) + + +class _concurrent_tee: + __slots__ = ('iterator', 'link', 'lock') + + def __init__(self, iterable): + if isinstance(iterable, _concurrent_tee): + self.iterator = iterable.iterator + self.link = iterable.link + self.lock = iterable.lock + else: + self.iterator = iter(iterable) + self.link = [None, None] + self.lock = Lock() + + def __iter__(self): + return self + + def __next__(self): + link = self.link + if link[1] is None: + with self.lock: + if link[1] is None: + link[0] = next(self.iterator) + link[1] = [None, None] + value, self.link = link + return value + + +def subfactorial(n): + """Number of permutations of *n* elements with no fixed points. + + The :func:`subfactorial` function computes the length of + :func:`derangements`. For example, there are 1,854 ways to + rearrange the letters in word "epsilon" without leaving any + letter in its original position: + + >>> from more_itertools import derangements, ilen + >>> ilen(derangements('epsilon')) + 1854 + >>> subfactorial(len('epsilon')) + 1854 + + Reference: https://oeis.org/A000166 + + """ + if n < 0: + raise ValueError + sf = adj = 1 + for i in range(n + 1): + sf = sf * i + adj + adj = -adj + return sf + + +def _full_period_lcg(n): + "Returns values from range(n) in randomly shuffled order." + + # The algorithm is a fluxed, full period linear congruential generator. + # Size *m* is a power of two. The *mask* speeds-up modulo calculations. + # The multiplier *a* and addend *c* are Hull-Dobell constants. See: + # https://jackgiffin.com/main/pdfs/Random-Number-Generators-T-E-Hull-and-A-R-Dobell.pdf + # The value *x* is a random starting point. + # The *flux* bitfield breaks-up patterns in the LCG. + + m = 1 << n.bit_length() + mask = m - 1 + a = randrange(5, m, 4) if m > 4 else 1 + c = randrange(1, m, 2) if m > 1 else 1 + x = randrange(m) + flux = randrange(m) + + for _ in repeat(None, m): + index = x ^ flux + if index < n: + yield index + x = (a * x + c) & mask + + +def _random_ordered_indices(n): + "Shuffle batches to mitigate the small state space of the LCG." + + # Batched variation of Knuth's Algorithm M in §3.2.2 of TAOCP. + batch_size = 1024 + iterable = range(n) if n <= batch_size else _full_period_lcg(n) + for batch in map(list, batched(iterable, batch_size)): + shuffle(batch) + yield from batch + + +def random_ordered_range(*args): + "Return :func:`range` values in randomly shuffled order." + range_object = range(*args) + n = len(range_object) + for index in _random_ordered_indices(n): + yield range_object[index] diff --git a/tests/test_more.py b/tests/test_more.py new file mode 100644 index 0000000..9ecdd06 --- /dev/null +++ b/tests/test_more.py @@ -0,0 +1,7444 @@ +from __future__ import annotations + +import cmath +import gc +import platform +import types +import weakref + +from collections import Counter, deque +from collections.abc import Set, Sequence, Iterable, Iterator, Hashable +from datetime import datetime, timedelta +from decimal import Decimal +from doctest import DocTestSuite +from fractions import Fraction +from functools import partial, reduce +from io import StringIO +from itertools import ( + accumulate, + chain, + combinations, + combinations_with_replacement, + count, + cycle, + groupby, + islice, + permutations, + product, + repeat, +) +from math import factorial +from operator import add, mul, itemgetter, not_ +from pickle import loads, dumps +from random import Random, choices, random, randrange, seed +from statistics import mean, median +from string import ascii_letters +from threading import Event, Thread, Lock +from time import sleep +from typing import NamedTuple +from unittest import TestCase, mock, skipUnless + +import more_itertools as mi + + +def load_tests(loader, tests, ignore): + # Add the doctests + tests.addTests(DocTestSuite('more_itertools.more')) + return tests + + +class ChunkedTests(TestCase): + """Tests for ``chunked()``""" + + def test_even(self): + """Test when ``n`` divides evenly into the length of the iterable.""" + self.assertEqual( + list(mi.chunked('ABCDEF', 3)), [['A', 'B', 'C'], ['D', 'E', 'F']] + ) + + def test_odd(self): + """Test when ``n`` does not divide evenly into the length of the + iterable. + + """ + self.assertEqual( + list(mi.chunked('ABCDE', 3)), [['A', 'B', 'C'], ['D', 'E']] + ) + + def test_none(self): + """Test when ``n`` has the value ``None``.""" + self.assertEqual( + list(mi.chunked('ABCDE', None)), [['A', 'B', 'C', 'D', 'E']] + ) + + def test_negative(self): + """Test that a negative ``n`` raises a clear ``ValueError``, matching + the behavior of :func:`sliced`.""" + self.assertRaisesRegex( + ValueError, + "n must be at least 0", + lambda: list(mi.chunked('ABCDE', -1)), + ) + + def test_strict_false(self): + """Test when ``n`` does not divide evenly into the length of the + iterable and strict is false. + + """ + self.assertEqual( + list(mi.chunked('ABCDE', 3, strict=False)), + [['A', 'B', 'C'], ['D', 'E']], + ) + + def test_strict_being_true(self): + """Test when ``n`` does not divide evenly into the length of the + iterable and strict is True (raising an exception). + + """ + + def f(): + return list(mi.chunked('ABCDE', 3, strict=True)) + + self.assertRaisesRegex(ValueError, "iterable is not divisible by n", f) + self.assertEqual( + list(mi.chunked('ABCDEF', 3, strict=True)), + [['A', 'B', 'C'], ['D', 'E', 'F']], + ) + + def test_strict_being_true_with_size_none(self): + """Test when ``n`` has value ``None`` and the keyword strict is True + (raising an exception). + + """ + + def f(): + return list(mi.chunked('ABCDE', None, strict=True)) + + self.assertRaisesRegex( + ValueError, "n must not be None when using strict mode.", f + ) + + +class FirstTests(TestCase): + def test_many(self): + # Also try it on a generator expression to make sure it works on + # whatever those return, across Python versions. + self.assertEqual(mi.first(x for x in range(4)), 0) + + def test_one(self): + self.assertEqual(mi.first([3]), 3) + + def test_empty(self): + with self.assertRaises(ValueError): + mi.first([]) + + def test_default(self): + self.assertEqual(mi.first([], 'boo'), 'boo') + + +class IterOnlyRange: + """User-defined iterable class which only support __iter__. + + >>> r = IterOnlyRange(5) + >>> r[0] + AttributeError: IterOnlyRange instance has no attribute '__getitem__' + + Note: In Python 3, ``TypeError`` will be raised because ``object`` is + inherited implicitly by default. + + >>> r[0] + TypeError: 'IterOnlyRange' object does not support indexing + """ + + def __init__(self, n): + """Set the length of the range.""" + self.n = n + + def __iter__(self): + """Works same as range().""" + return iter(range(self.n)) + + +class LastTests(TestCase): + def test_basic(self): + cases = [ + (range(4), 3), + (iter(range(4)), 3), + (range(1), 0), + (iter(range(1)), 0), + (IterOnlyRange(5), 4), + ({n: str(n) for n in range(5)}, 4), + ({0: '0', -1: '-1', 2: '-2'}, 2), + ] + + for iterable, expected in cases: + with self.subTest(iterable=iterable): + self.assertEqual(mi.last(iterable), expected) + + def test_default(self): + for iterable, default, expected in [ + (range(1), None, 0), + ([], None, None), + ({}, None, None), + (iter([]), None, None), + ]: + with self.subTest(args=(iterable, default)): + self.assertEqual(mi.last(iterable, default=default), expected) + + def test_empty(self): + for iterable in ([], iter(range(0))): + with self.subTest(iterable=iterable): + with self.assertRaises(ValueError): + mi.last(iterable) + + def test_reversed_is_none(self): + # See https://github.com/more-itertools/more-itertools/issues/1001 + class ReversedIsNone: + __reversed__ = None + + def __iter__(self): + return iter([1]) + + self.assertEqual(mi.last(ReversedIsNone()), 1) + + +class NthOrLastTests(TestCase): + """Tests for ``nth_or_last()``""" + + def test_basic(self): + self.assertEqual(mi.nth_or_last(range(3), 1), 1) + self.assertEqual(mi.nth_or_last(range(3), 3), 2) + + def test_default_value(self): + default = 42 + self.assertEqual(mi.nth_or_last(range(0), 3, default), default) + + def test_empty_iterable_no_default(self): + self.assertRaises(ValueError, lambda: mi.nth_or_last(range(0), 0)) + + +class PeekableMixinTests: + """Common tests for ``peekable()`` and ``seekable()`` behavior""" + + cls = None + + def test_passthrough(self): + """Iterating a peekable without using ``peek()`` or ``prepend()`` + should just give the underlying iterable's elements (a trivial test but + useful to set a baseline in case something goes wrong)""" + expected = [1, 2, 3, 4, 5] + actual = list(self.cls(expected)) + self.assertEqual(actual, expected) + + def test_peek_default(self): + """Make sure passing a default into ``peek()`` works.""" + p = self.cls([]) + self.assertEqual(p.peek(7), 7) + + def test_truthiness(self): + """Make sure a ``peekable`` tests true iff there are items remaining in + the iterable. + + """ + p = self.cls([]) + self.assertFalse(p) + + p = self.cls(range(3)) + self.assertTrue(p) + + def test_simple_peeking(self): + """Make sure ``next`` and ``peek`` advance and don't advance the + iterator, respectively. + + """ + p = self.cls(range(10)) + self.assertEqual(next(p), 0) + self.assertEqual(p.peek(), 1) + self.assertEqual(p.peek(), 1) + self.assertEqual(next(p), 1) + + +class PeekableTests(PeekableMixinTests, TestCase): + cls = mi.peekable + + def test_indexing(self): + """ + Indexing into the peekable shouldn't advance the iterator. + """ + p = mi.peekable('abcdefghijkl') + + # The 0th index is what ``next()`` will return + self.assertEqual(p[0], 'a') + self.assertEqual(next(p), 'a') + + # Indexing further into the peekable shouldn't advance the iterator + self.assertEqual(p[2], 'd') + self.assertEqual(next(p), 'b') + + # The 0th index moves up with the iterator; the last index follows + self.assertEqual(p[0], 'c') + self.assertEqual(p[9], 'l') + + self.assertEqual(next(p), 'c') + self.assertEqual(p[8], 'l') + + # Negative indexing should work too + self.assertEqual(p[-2], 'k') + self.assertEqual(p[-9], 'd') + self.assertRaises(IndexError, lambda: p[-10]) + + def test_slicing(self): + """Slicing the peekable shouldn't advance the iterator.""" + seq = list('abcdefghijkl') + p = mi.peekable(seq) + + # Slicing the peekable should just be like slicing a re-iterable + self.assertEqual(p[1:4], seq[1:4]) + + # Advancing the iterator moves the slices up also + self.assertEqual(next(p), 'a') + self.assertEqual(p[1:4], seq[1:][1:4]) + + # Implicit starts and stop should work + self.assertEqual(p[:5], seq[1:][:5]) + self.assertEqual(p[:], seq[1:][:]) + + # Indexing past the end should work + self.assertEqual(p[:100], seq[1:][:100]) + + # Steps should work, including negative + self.assertEqual(p[::2], seq[1:][::2]) + self.assertEqual(p[::-1], seq[1:][::-1]) + + def test_slicing_reset(self): + """Test slicing on a fresh iterable each time""" + iterable = ['0', '1', '2', '3', '4', '5'] + indexes = list(range(-4, len(iterable) + 4)) + [None] + steps = [1, 2, 3, 4, -1, -2, -3, 4] + for slice_args in product(indexes, indexes, steps): + it = iter(iterable) + p = mi.peekable(it) + next(p) + index = slice(*slice_args) + actual = p[index] + expected = iterable[1:][index] + self.assertEqual(actual, expected, slice_args) + + def test_slicing_error(self): + iterable = '01234567' + p = mi.peekable(iter(iterable)) + + # Prime the cache + p.peek() + old_cache = list(p._cache) + + # Illegal slice + with self.assertRaises(ValueError): + p[1:-1:0] + + # Neither the cache nor the iteration should be affected + self.assertEqual(old_cache, list(p._cache)) + self.assertEqual(list(p), list(iterable)) + + # prepend() behavior tests + + def test_prepend(self): + """Tests interspersed ``prepend()`` and ``next()`` calls""" + it = mi.peekable(range(2)) + actual = [] + + # Test prepend() before next() + it.prepend(10) + actual += [next(it), next(it)] + + # Test prepend() between next()s + it.prepend(11) + actual += [next(it), next(it)] + + # Test prepend() after source iterable is consumed + it.prepend(12) + actual += [next(it)] + + expected = [10, 0, 11, 1, 12] + self.assertEqual(actual, expected) + + def test_multi_prepend(self): + """Tests prepending multiple items and getting them in proper order""" + it = mi.peekable(range(5)) + actual = [next(it), next(it)] + it.prepend(10, 11, 12) + it.prepend(20, 21) + actual += list(it) + expected = [0, 1, 20, 21, 10, 11, 12, 2, 3, 4] + self.assertEqual(actual, expected) + + def test_empty(self): + """Tests prepending in front of an empty iterable""" + it = mi.peekable([]) + it.prepend(10) + actual = list(it) + expected = [10] + self.assertEqual(actual, expected) + + def test_prepend_truthiness(self): + """Tests that ``__bool__()`` or ``__nonzero__()`` works properly + with ``prepend()``""" + it = mi.peekable(range(5)) + self.assertTrue(it) + actual = list(it) + self.assertFalse(it) + it.prepend(10) + self.assertTrue(it) + actual += [next(it)] + self.assertFalse(it) + expected = [0, 1, 2, 3, 4, 10] + self.assertEqual(actual, expected) + + def test_multi_prepend_peek(self): + """Tests prepending multiple elements and getting them in reverse order + while peeking""" + it = mi.peekable(range(5)) + actual = [next(it), next(it)] + self.assertEqual(it.peek(), 2) + it.prepend(10, 11, 12) + self.assertEqual(it.peek(), 10) + it.prepend(20, 21) + self.assertEqual(it.peek(), 20) + actual += list(it) + self.assertFalse(it) + expected = [0, 1, 20, 21, 10, 11, 12, 2, 3, 4] + self.assertEqual(actual, expected) + + def test_prepend_after_stop(self): + """Test resuming iteration after a previous exhaustion""" + it = mi.peekable(range(3)) + self.assertEqual(list(it), [0, 1, 2]) + self.assertRaises(StopIteration, lambda: next(it)) + it.prepend(10) + self.assertEqual(next(it), 10) + self.assertRaises(StopIteration, lambda: next(it)) + + def test_prepend_slicing(self): + """Tests interaction between prepending and slicing""" + seq = list(range(20)) + p = mi.peekable(seq) + + p.prepend(30, 40, 50) + pseq = [30, 40, 50] + seq # pseq for prepended_seq + + # adapt the specific tests from test_slicing + self.assertEqual(p[0], 30) + self.assertEqual(p[1:8], pseq[1:8]) + self.assertEqual(p[1:], pseq[1:]) + self.assertEqual(p[:5], pseq[:5]) + self.assertEqual(p[:], pseq[:]) + self.assertEqual(p[:100], pseq[:100]) + self.assertEqual(p[::2], pseq[::2]) + self.assertEqual(p[::-1], pseq[::-1]) + + def test_prepend_indexing(self): + """Tests interaction between prepending and indexing""" + seq = list(range(20)) + p = mi.peekable(seq) + + p.prepend(30, 40, 50) + + self.assertEqual(p[0], 30) + self.assertEqual(next(p), 30) + self.assertEqual(p[2], 0) + self.assertEqual(next(p), 40) + self.assertEqual(p[0], 50) + self.assertEqual(p[9], 8) + self.assertEqual(next(p), 50) + self.assertEqual(p[8], 8) + self.assertEqual(p[-2], 18) + self.assertEqual(p[-9], 11) + self.assertRaises(IndexError, lambda: p[-21]) + + def test_prepend_iterable(self): + """Tests prepending from an iterable""" + it = mi.peekable(range(5)) + # Don't directly use the range() object to avoid any range-specific + # optimizations + it.prepend(*(x for x in range(5))) + actual = list(it) + expected = list(chain(range(5), range(5))) + self.assertEqual(actual, expected) + + def test_prepend_many(self): + """Tests that prepending a huge number of elements works""" + it = mi.peekable(range(5)) + # Don't directly use the range() object to avoid any range-specific + # optimizations + it.prepend(*(x for x in range(20000))) + actual = list(it) + expected = list(chain(range(20000), range(5))) + self.assertEqual(actual, expected) + + def test_prepend_reversed(self): + """Tests prepending from a reversed iterable""" + it = mi.peekable(range(3)) + it.prepend(*reversed((10, 11, 12))) + actual = list(it) + expected = [12, 11, 10, 0, 1, 2] + self.assertEqual(actual, expected) + + def test_class_getitem(self): + """peekable[T] should return a GenericAlias, matching the behaviour of + list[T], allowing use in generic type annotations at runtime.""" + alias = mi.peekable[str] + self.assertIsInstance(alias, types.GenericAlias) + self.assertIs(alias.__origin__, mi.peekable) + self.assertEqual(alias.__args__, (str,)) + + +class ConsumerTests(TestCase): + """Tests for ``consumer()``""" + + def test_consumer(self): + @mi.consumer + def eater(): + while True: + x = yield # noqa + + e = eater() + e.send('hi') # without @consumer, would raise TypeError + + +class DistinctPermutationsTests(TestCase): + def test_basic(self): + iterable = ['z', 'a', 'a', 'q', 'q', 'q', 'y'] + actual = list(mi.distinct_permutations(iterable)) + expected = set(permutations(iterable)) + self.assertCountEqual(actual, expected) + + def test_r(self): + for iterable, r in ( + ('mississippi', 0), + ('mississippi', 1), + ('mississippi', 6), + ('mississippi', 7), + ('mississippi', 12), + ([0, 1, 1, 0], 0), + ([0, 1, 1, 0], 1), + ([0, 1, 1, 0], 2), + ([0, 1, 1, 0], 3), + ([0, 1, 1, 0], 4), + (['a'], 0), + (['a'], 1), + (['a'], 5), + ([], 0), + ([], 1), + ([], 4), + ): + with self.subTest(iterable=iterable, r=r): + expected = set(permutations(iterable, r)) + actual = list(mi.distinct_permutations(iter(iterable), r)) + self.assertCountEqual(actual, expected) + + def test_unsortable(self): + iterable = ['1', 2, 2, 3, 3, 3] + actual = list(mi.distinct_permutations(iterable)) + expected = set(permutations(iterable)) + self.assertCountEqual(actual, expected) + + def test_unsortable_r(self): + iterable = ['1', 2, 2, 3, 3, 3] + for r in range(len(iterable) + 1): + with self.subTest(iterable=iterable, r=r): + actual = list(mi.distinct_permutations(iterable, r=r)) + expected = set(permutations(iterable, r=r)) + self.assertCountEqual(actual, expected) + + def test_unsorted_equivalent(self): + iterable = [1, True, '3'] + actual = list(mi.distinct_permutations(iterable)) + expected = set(permutations(iterable)) + self.assertCountEqual(actual, expected) + + def test_unhashable(self): + iterable = ([1], [1], 2) + actual = list(mi.distinct_permutations(iterable)) + expected = list(mi.unique_everseen(permutations(iterable), key=str)) + self.assertCountEqual(actual, expected) + + def test_sortable_input_produces_sorted_output(self): + for iterable in [ + [4, 3, 2, 1], + [2, 2, 3, 1], + ]: + for r in range(len(iterable) + 1): + with self.subTest(iterable=iterable, r=r): + actual = list(mi.distinct_permutations(iterable, r=r)) + expected = sorted(set(permutations(iterable, r=r))) + self.assertEqual(actual, expected) + + +class DerangementsTests(TestCase): + def test_unique_values(self): + n = 8 + expected = set( + x + for x in permutations(range(n)) + if not any(x[i] == i for i in range(n)) + ) + for i, iterable in enumerate( + [ + range(n), + list(range(n)), + set(range(n)), + ] + ): + actual = set(mi.derangements(iterable)) + self.assertEqual(actual, expected) + + def test_repeated_values(self): + self.assertEqual( + [''.join(x) for x in mi.derangements('AACD')], + [ + 'AADC', + 'ACDA', + 'ADAC', + 'CADA', + 'CDAA', + 'CDAA', + 'DAAC', + 'DCAA', + 'DCAA', + ], + ) + + def test_unsortable_unhashable(self): + iterable = (0, True, ['Carol']) + actual = list(mi.derangements(iterable)) + expected = [(True, ['Carol'], 0), (['Carol'], 0, True)] + self.assertListEqual(actual, expected) + + def test_r(self): + s = 'ABCD' + for r, expected in [ + (0, ['']), + (1, ['B', 'C', 'D']), + (2, ['BA', 'BC', 'BD', 'CA', 'CD', 'DA', 'DC']), + ( + 3, + [ + 'BAD', + 'BCA', + 'BCD', + 'BDA', + 'CAB', + 'CAD', + 'CDA', + 'CDB', + 'DAB', + 'DCA', + 'DCB', + ], + ), + ( + 4, + [ + 'BADC', + 'BCDA', + 'BDAC', + 'CADB', + 'CDAB', + 'CDBA', + 'DABC', + 'DCAB', + 'DCBA', + ], + ), + ]: + with self.subTest(r=r): + actual = [''.join(x) for x in mi.derangements(s, r=r)] + self.assertEqual(actual, expected) + + +class IlenTests(TestCase): + def test_ilen(self): + """Sanity-checks for ``ilen()``.""" + # Non-empty + self.assertEqual( + mi.ilen(filter(lambda x: x % 10 == 0, range(101))), 11 + ) + + # Empty + self.assertEqual(mi.ilen(x for x in range(0)), 0) + + # Iterable with __len__ + self.assertEqual(mi.ilen(list(range(6))), 6) + + +class MinMaxTests(TestCase): + def test_basic(self): + for iterable, expected in ( + # easy case + ([0, 1, 2, 3], (0, 3)), + # min and max are not in the extremes + we have `int`s and `float`s + ([3, 5.5, -1, 2], (-1, 5.5)), + # unordered collection + ({3, 5.5, -1, 2}, (-1, 5.5)), + # with repetitions + ([3, 5.5, float('-Inf'), 5.5], (float('-Inf'), 5.5)), + # other collections + ('banana', ('a', 'n')), + ({0: 1, 2: 100, 1: 10}, (0, 2)), + (range(3, 14), (3, 13)), + ): + with self.subTest(iterable=iterable, expected=expected): + # check for expected results + self.assertTupleEqual(mi.minmax(iterable), expected) + # check for equality with built-in `min` and `max` + self.assertTupleEqual( + mi.minmax(iterable), (min(iterable), max(iterable)) + ) + + def test_unpacked(self): + self.assertTupleEqual(mi.minmax(2, 3, 1), (1, 3)) + self.assertTupleEqual(mi.minmax(12, 3, 4, key=str), (12, 4)) + + def test_iterables(self): + self.assertTupleEqual(mi.minmax(x for x in [0, 1, 2, 3]), (0, 3)) + self.assertTupleEqual( + mi.minmax(map(str, [3, 5.5, 'a', 2])), ('2', 'a') + ) + self.assertTupleEqual( + mi.minmax(filter(None, [0, 3, '', None, 10])), (3, 10) + ) + + def test_key(self): + self.assertTupleEqual( + mi.minmax({(), (1, 4, 2), 'abcde', range(4)}, key=len), + ((), 'abcde'), + ) + self.assertTupleEqual( + mi.minmax((x for x in [10, 3, 25]), key=str), (10, 3) + ) + + def test_key_matches_builtins(self): + # The keyed code path compares items pairwise, so a single example + # only reaches some of its branches. Exhausting the short orderings + # covers each one. `x % 7` keeps every key distinct, so there are no + # ties and `min`/`max` pick the same items `minmax` does. + key = lambda x: x % 7 + for n in range(1, 5): + for values in permutations(range(1, 8), n): + with self.subTest(values=values): + self.assertTupleEqual( + mi.minmax(values, key=key), + (min(values, key=key), max(values, key=key)), + ) + + def test_default(self): + with self.assertRaises(ValueError): + mi.minmax([]) + + self.assertIs(mi.minmax([], default=None), None) + self.assertListEqual(mi.minmax([], default=[1, 'a']), [1, 'a']) + + +class WithIterTests(TestCase): + def test_with_iter(self): + s = StringIO('One fish\nTwo fish') + initial_words = [line.split()[0] for line in mi.with_iter(s)] + + # Iterable's items should be faithfully represented + self.assertEqual(initial_words, ['One', 'Two']) + # The file object should be closed + self.assertTrue(s.closed) + + +class SizedIteratorTests(TestCase): + def test_sized_iterator(self): + gen = (x for x in range(3)) + sized_gen = mi.sized_iterator(gen, 3) + + # The sized_generator should have the correct length + self.assertEqual(len(sized_gen), 3) + # The sized_generator elements should be the same as the original + self.assertListEqual(list(sized_gen), list(range(3))) + + +class OneTests(TestCase): + def test_basic(self): + it = iter(['item']) + self.assertEqual(mi.one(it), 'item') + + def test_too_short_new(self): + it = iter([]) + self.assertRaises(ValueError, lambda: mi.one(it)) + self.assertRaises( + OverflowError, lambda: mi.one(it, too_short=OverflowError) + ) + + def test_too_long(self): + it = count() + self.assertRaises(ValueError, lambda: mi.one(it)) # burn 0 and 1 + self.assertEqual(next(it), 2) + self.assertRaises( + OverflowError, lambda: mi.one(it, too_long=OverflowError) + ) + + def test_too_long_default_message(self): + it = count() + self.assertRaisesRegex( + ValueError, + "Expected exactly one item in " + "iterable, but got 0, 1, and " + "perhaps more.", + lambda: mi.one(it), + ) + + def test_falsy_custom_exception(self): + # An exception whose instances are falsy must still be raised, rather + # than being treated as "no custom exception given". + class FalsyError(Exception): + def __bool__(self): + return False + + too_long = FalsyError('too many') + self.assertRaises( + FalsyError, lambda: mi.one(count(), too_long=too_long) + ) + + too_short = FalsyError('too few') + self.assertRaises(FalsyError, lambda: mi.one([], too_short=too_short)) + + def test_too_long_does_not_evaluate_repr(self): + # When too_long is given, the default message (which reprs the first two + # items) should never be built. + class NoRepr: + def __repr__(self): + raise RuntimeError('repr should not be called') + + it = (NoRepr() for _ in count()) + self.assertRaises( + OverflowError, lambda: mi.one(it, too_long=OverflowError) + ) + + +class IntersperseTest(TestCase): + """Tests for intersperse()""" + + def test_even(self): + iterable = (x for x in '01') + self.assertEqual( + list(mi.intersperse(None, iterable)), ['0', None, '1'] + ) + + def test_odd(self): + iterable = (x for x in '012') + self.assertEqual( + list(mi.intersperse(None, iterable)), ['0', None, '1', None, '2'] + ) + + def test_nested(self): + element = ('a', 'b') + iterable = (x for x in '012') + actual = list(mi.intersperse(element, iterable)) + expected = ['0', ('a', 'b'), '1', ('a', 'b'), '2'] + self.assertEqual(actual, expected) + + def test_not_iterable(self): + self.assertRaises(TypeError, lambda: mi.intersperse('x', 1)) + + def test_n(self): + for n, element, expected in [ + (1, '_', ['0', '_', '1', '_', '2', '_', '3', '_', '4', '_', '5']), + (2, '_', ['0', '1', '_', '2', '3', '_', '4', '5']), + (3, '_', ['0', '1', '2', '_', '3', '4', '5']), + (4, '_', ['0', '1', '2', '3', '_', '4', '5']), + (5, '_', ['0', '1', '2', '3', '4', '_', '5']), + (6, '_', ['0', '1', '2', '3', '4', '5']), + (7, '_', ['0', '1', '2', '3', '4', '5']), + (3, ['a', 'b'], ['0', '1', '2', ['a', 'b'], '3', '4', '5']), + ]: + iterable = (x for x in '012345') + actual = list(mi.intersperse(element, iterable, n=n)) + self.assertEqual(actual, expected) + + def test_n_zero(self): + self.assertRaises( + ValueError, lambda: list(mi.intersperse('x', '012', n=0)) + ) + + +class UniqueToEachTests(TestCase): + """Tests for ``unique_to_each()``""" + + def test_all_unique(self): + """When all the input iterables are unique the output should match + the input.""" + iterables = [[1, 2], [3, 4, 5], [6, 7, 8]] + self.assertEqual(mi.unique_to_each(*iterables), iterables) + + def test_duplicates(self): + """When there are duplicates in any of the input iterables that aren't + in the rest, those duplicates should be emitted.""" + iterables = ["mississippi", "missouri"] + self.assertEqual( + mi.unique_to_each(*iterables), [['p', 'p'], ['o', 'u', 'r']] + ) + + def test_mixed(self): + """When the input iterables contain different types the function should + still behave properly""" + iterables = ['x', (i for i in range(3)), [1, 2, 3], tuple()] + self.assertEqual(mi.unique_to_each(*iterables), [['x'], [0], [3], []]) + + +class WindowedTests(TestCase): + def test_basic(self): + iterable = [1, 2, 3, 4, 5] + + for n, expected in ( + (6, [(1, 2, 3, 4, 5, None)]), + (5, [(1, 2, 3, 4, 5)]), + (4, [(1, 2, 3, 4), (2, 3, 4, 5)]), + (3, [(1, 2, 3), (2, 3, 4), (3, 4, 5)]), + (2, [(1, 2), (2, 3), (3, 4), (4, 5)]), + (1, [(1,), (2,), (3,), (4,), (5,)]), + ): + with self.subTest(n=n): + actual = list(mi.windowed(iterable, n)) + self.assertEqual(actual, expected) + + def test_fillvalue(self): + actual = list(mi.windowed([1, 2, 3, 4, 5], 6, fillvalue='!')) + expected = [(1, 2, 3, 4, 5, '!')] + self.assertEqual(actual, expected) + + def test_step(self): + iterable = [1, 2, 3, 4, 5, 6, 7] + for n, step, expected in [ + (3, 2, [(1, 2, 3), (3, 4, 5), (5, 6, 7)]), # n > step + (3, 3, [(1, 2, 3), (4, 5, 6), (7, None, None)]), # n == step + (3, 4, [(1, 2, 3), (5, 6, 7)]), # lines up nicely + (3, 5, [(1, 2, 3), (6, 7, None)]), # off by one + (3, 6, [(1, 2, 3), (7, None, None)]), # off by two + (3, 7, [(1, 2, 3)]), # step past the end + (7, 8, [(1, 2, 3, 4, 5, 6, 7)]), # step > len(iterable) + ]: + with self.subTest(n=n, step=step): + actual = list(mi.windowed(iterable, n, step=step)) + self.assertEqual(actual, expected) + + def test_invalid_step(self): + # Step must be greater than or equal to 1 + with self.assertRaises(ValueError): + list(mi.windowed([1, 2, 3, 4, 5], 3, step=0)) + + def test_fillvalue_step(self): + actual = list(mi.windowed([1, 2, 3, 4, 5], 3, fillvalue='!', step=3)) + expected = [(1, 2, 3), (4, 5, '!')] + self.assertEqual(actual, expected) + + def test_invalid_n(self): + with self.assertRaises(ValueError): + list(mi.windowed([1, 2, 3, 4, 5], 0)) # n is zero + with self.assertRaises(ValueError): + list(mi.windowed([1, 2, 3, 4, 5], -1)) # n is negative + + def test_empty_seq(self): + actual = list(mi.windowed([], 3)) + expected = [] + self.assertEqual(actual, expected) + + +class SubstringsTests(TestCase): + def test_basic(self): + iterable = (x for x in range(4)) + actual = list(mi.substrings(iterable)) + expected = [ + (0,), + (1,), + (2,), + (3,), + (0, 1), + (1, 2), + (2, 3), + (0, 1, 2), + (1, 2, 3), + (0, 1, 2, 3), + ] + self.assertEqual(actual, expected) + + def test_strings(self): + iterable = 'abc' + actual = list(mi.substrings(iterable)) + expected = [ + ('a',), + ('b',), + ('c',), + ('a', 'b'), + ('b', 'c'), + ('a', 'b', 'c'), + ] + self.assertEqual(actual, expected) + + def test_empty(self): + iterable = iter([]) + actual = list(mi.substrings(iterable)) + expected = [] + self.assertEqual(actual, expected) + + def test_order(self): + iterable = [2, 0, 1] + actual = list(mi.substrings(iterable)) + expected = [(2,), (0,), (1,), (2, 0), (0, 1), (2, 0, 1)] + self.assertEqual(actual, expected) + + +class SubstringsIndexesTests(TestCase): + def test_basic(self): + sequence = [x for x in range(4)] + actual = list(mi.substrings_indexes(sequence)) + expected = [ + ([0], 0, 1), + ([1], 1, 2), + ([2], 2, 3), + ([3], 3, 4), + ([0, 1], 0, 2), + ([1, 2], 1, 3), + ([2, 3], 2, 4), + ([0, 1, 2], 0, 3), + ([1, 2, 3], 1, 4), + ([0, 1, 2, 3], 0, 4), + ] + self.assertEqual(actual, expected) + + def test_strings(self): + sequence = 'abc' + actual = list(mi.substrings_indexes(sequence)) + expected = [ + ('a', 0, 1), + ('b', 1, 2), + ('c', 2, 3), + ('ab', 0, 2), + ('bc', 1, 3), + ('abc', 0, 3), + ] + self.assertEqual(actual, expected) + + def test_empty(self): + sequence = [] + actual = list(mi.substrings_indexes(sequence)) + expected = [] + self.assertEqual(actual, expected) + + def test_order(self): + sequence = [2, 0, 1] + actual = list(mi.substrings_indexes(sequence)) + expected = [ + ([2], 0, 1), + ([0], 1, 2), + ([1], 2, 3), + ([2, 0], 0, 2), + ([0, 1], 1, 3), + ([2, 0, 1], 0, 3), + ] + self.assertEqual(actual, expected) + + def test_reverse(self): + sequence = [2, 0, 1] + actual = list(mi.substrings_indexes(sequence, reverse=True)) + expected = [ + ([2, 0, 1], 0, 3), + ([2, 0], 0, 2), + ([0, 1], 1, 3), + ([2], 0, 1), + ([0], 1, 2), + ([1], 2, 3), + ] + self.assertEqual(actual, expected) + + +class BucketTests(TestCase): + def test_basic(self): + iterable = [10, 20, 30, 11, 21, 31, 12, 22, 23, 33] + D = mi.bucket(iterable, key=lambda x: 10 * (x // 10)) + + # In-order access + self.assertEqual(list(D[10]), [10, 11, 12]) + + # Out of order access + self.assertEqual(list(D[30]), [30, 31, 33]) + self.assertEqual(list(D[20]), [20, 21, 22, 23]) + + self.assertEqual(list(D[40]), []) # Nothing in here! + + def test_in(self): + iterable = [10, 20, 30, 11, 21, 31, 12, 22, 23, 33] + D = mi.bucket(iterable, key=lambda x: 10 * (x // 10)) + + self.assertIn(10, D) + self.assertNotIn(40, D) + self.assertIn(20, D) + self.assertNotIn(21, D) + + # Checking in-ness shouldn't advance the iterator + self.assertEqual(next(D[10]), 10) + + def test_validator(self): + iterable = count(0) + key = lambda x: int(str(x)[0]) # First digit of each number + validator = lambda x: 0 < x < 10 # No leading zeros + D = mi.bucket(iterable, key, validator=validator) + self.assertEqual(mi.take(3, D[1]), [1, 10, 11]) + self.assertNotIn(0, D) # Non-valid entries don't return True + self.assertNotIn(0, D._cache) # Don't store non-valid entries + self.assertEqual(list(D[0]), []) + + def test_list(self): + iterable = [10, 20, 30, 11, 21, 31, 12, 22, 23, 33] + D = mi.bucket(iterable, key=lambda x: 10 * (x // 10)) + self.assertEqual(list(D[10]), [10, 11, 12]) + self.assertEqual(list(D[20]), [20, 21, 22, 23]) + self.assertEqual(list(D[30]), [30, 31, 33]) + self.assertEqual(set(D), {10, 20, 30}) + + def test_list_validator(self): + iterable = [10, 20, 30, 11, 21, 31, 12, 22, 23, 33] + key = lambda x: 10 * (x // 10) + validator = lambda x: x != 20 + D = mi.bucket(iterable, key, validator=validator) + self.assertEqual(set(D), {10, 30}) + self.assertEqual(list(D[10]), [10, 11, 12]) + self.assertEqual(list(D[20]), []) + self.assertEqual(list(D[30]), [30, 31, 33]) + + def test_falsey_validator_is_honoured(self): + # A falsey callable is still a valid validator and must not be + # replaced by the default (see #1308). + class AllowedKeys(set): + def __call__(self, key): + return key in self + + allowed = AllowedKeys() + self.assertFalse(allowed) # Sanity: the validator itself is falsey + D = mi.bucket( + iter([10, 20, 11, 21]), key=lambda x: x, validator=allowed + ) + self.assertEqual(list(D[10]), []) + self.assertNotIn(10, D) + self.assertEqual(set(D), set()) + + def test_in_does_not_add_key(self): + # A failed membership test must not invent a key (see #1284) + iterable = [10, 20, 11, 21] + D = mi.bucket(iterable, key=lambda x: 10 * (x // 10)) + self.assertFalse(30 in D) + self.assertEqual(set(D), {10, 20}) + self.assertEqual(list(D[30]), []) + + def test_getitem_miss_does_not_add_key(self): + # Selecting an absent bucket must not invent a key either + iterable = [10, 20, 11, 21] + D = mi.bucket(iterable, key=lambda x: 10 * (x // 10)) + self.assertEqual(list(D[30]), []) + self.assertEqual(set(D), {10, 20}) + + def test_in_does_not_add_key_with_validator(self): + # A key the validator accepts but the iterable never produces + iterable = [10, 20, 11, 21] + key = lambda x: 10 * (x // 10) + validator = lambda x: x in {10, 20, 30} + D = mi.bucket(iterable, key, validator=validator) + self.assertFalse(30 in D) + self.assertEqual(set(D), {10, 20}) + + def test_keys_are_only_those_seen(self): + # Every reported key must have had at least one matching item + iterable = [10, 20, 11] + D = mi.bucket(iterable, key=lambda x: 10 * (x // 10)) + for missing in (30, 40, 50): + self.assertFalse(missing in D) + self.assertEqual(set(D), {10, 20}) + + +class SpyTests(TestCase): + """Tests for ``spy()``""" + + def test_basic(self): + original_iterable = iter('abcdefg') + head, new_iterable = mi.spy(original_iterable) + self.assertEqual(head, ['a']) + self.assertEqual( + list(new_iterable), ['a', 'b', 'c', 'd', 'e', 'f', 'g'] + ) + + def test_unpacking(self): + original_iterable = iter('abcdefg') + (first, second, third), new_iterable = mi.spy(original_iterable, 3) + self.assertEqual(first, 'a') + self.assertEqual(second, 'b') + self.assertEqual(third, 'c') + self.assertEqual( + list(new_iterable), ['a', 'b', 'c', 'd', 'e', 'f', 'g'] + ) + + def test_too_many(self): + original_iterable = iter('abc') + head, new_iterable = mi.spy(original_iterable, 4) + self.assertEqual(head, ['a', 'b', 'c']) + self.assertEqual(list(new_iterable), ['a', 'b', 'c']) + + def test_zero(self): + original_iterable = iter('abc') + head, new_iterable = mi.spy(original_iterable, 0) + self.assertEqual(head, []) + self.assertEqual(list(new_iterable), ['a', 'b', 'c']) + + def test_immutable(self): + original_iterable = iter('abcdefg') + head, new_iterable = mi.spy(original_iterable, 3) + head[0] = 'A' + self.assertEqual(head, ['A', 'b', 'c']) + self.assertEqual( + list(new_iterable), ['a', 'b', 'c', 'd', 'e', 'f', 'g'] + ) + + +class InterleaveTests(TestCase): + def test_even(self): + actual = list(mi.interleave([1, 4, 7], [2, 5, 8], [3, 6, 9])) + expected = [1, 2, 3, 4, 5, 6, 7, 8, 9] + self.assertEqual(actual, expected) + + def test_short(self): + actual = list(mi.interleave([1, 4], [2, 5, 7], [3, 6, 8])) + expected = [1, 2, 3, 4, 5, 6] + self.assertEqual(actual, expected) + + def test_mixed_types(self): + it_list = ['a', 'b', 'c', 'd'] + it_str = '12345' + it_inf = count() + actual = list(mi.interleave(it_list, it_str, it_inf)) + expected = ['a', '1', 0, 'b', '2', 1, 'c', '3', 2, 'd', '4', 3] + self.assertEqual(actual, expected) + + +class InterleaveLongestTests(TestCase): + def test_even(self): + actual = list(mi.interleave_longest([1, 4, 7], [2, 5, 8], [3, 6, 9])) + expected = [1, 2, 3, 4, 5, 6, 7, 8, 9] + self.assertEqual(actual, expected) + + def test_short(self): + actual = list(mi.interleave_longest([1, 4], [2, 5, 7], [3, 6, 8])) + expected = [1, 2, 3, 4, 5, 6, 7, 8] + self.assertEqual(actual, expected) + + def test_mixed_types(self): + it_list = ['a', 'b', 'c', 'd'] + it_str = '12345' + it_gen = (x for x in range(3)) + actual = list(mi.interleave_longest(it_list, it_str, it_gen)) + expected = ['a', '1', 0, 'b', '2', 1, 'c', '3', 2, 'd', '4', '5'] + self.assertEqual(actual, expected) + + +class InterleaveEvenlyTests(TestCase): + def test_equal_lengths(self): + # when lengths are equal, the relative order shouldn't change + a = [1, 2, 3] + b = [5, 6, 7] + actual = list(mi.interleave_evenly([a, b])) + expected = [1, 5, 2, 6, 3, 7] + self.assertEqual(actual, expected) + + def test_proportional(self): + # easy case where the iterables have proportional length + a = [1, 2, 3, 4] + b = [5, 6] + actual = list(mi.interleave_evenly([a, b])) + expected = [1, 2, 5, 3, 4, 6] + self.assertEqual(actual, expected) + + # swapping a and b should yield the same result + actual_swapped = list(mi.interleave_evenly([b, a])) + self.assertEqual(actual_swapped, expected) + + def test_not_proportional(self): + a = [1, 2, 3, 4, 5, 6, 7] + b = [8, 9, 10] + expected = [1, 2, 8, 3, 4, 9, 5, 6, 10, 7] + actual = list(mi.interleave_evenly([a, b])) + self.assertEqual(actual, expected) + + def test_degenerate_one(self): + a = [0, 1, 2, 3, 4] + b = [5] + expected = [0, 1, 2, 5, 3, 4] + actual = list(mi.interleave_evenly([a, b])) + self.assertEqual(actual, expected) + + def test_degenerate_empty(self): + a = [1, 2, 3] + b = [] + expected = [1, 2, 3] + actual = list(mi.interleave_evenly([a, b])) + self.assertEqual(actual, expected) + + def test_no_iterables(self): + self.assertEqual(list(mi.interleave_evenly([])), []) + self.assertEqual(list(mi.interleave_evenly([], lengths=[])), []) + + def test_three_iters(self): + a = ["a1", "a2", "a3", "a4", "a5"] + b = ["b1", "b2", "b3"] + c = ["c1"] + actual = list(mi.interleave_evenly([a, b, c])) + expected = ["a1", "b1", "a2", "c1", "a3", "b2", "a4", "b3", "a5"] + self.assertEqual(actual, expected) + + def test_many_iters(self): + # smoke test with many iterables: create iterables with a random + # number of elements starting with a character ("a0", "a1", ...) + rng = Random(0) + iterables = [] + for ch in ascii_letters: + length = rng.randint(0, 100) + iterable = [f"{ch}{i}" for i in range(length)] + iterables.append(iterable) + + interleaved = list(mi.interleave_evenly(iterables)) + + # for each iterable, check that the result contains all its items + for iterable, ch_expect in zip(iterables, ascii_letters): + interleaved_actual = [ + e for e in interleaved if e.startswith(ch_expect) + ] + assert len(set(interleaved_actual)) == len(iterable) + + def test_manual_lengths(self): + a = combinations(range(4), 2) + len_a = 4 * (4 - 1) // 2 # == 6 + b = combinations(range(4), 3) + len_b = 4 + + expected = [ + (0, 1), + (0, 1, 2), + (0, 2), + (0, 3), + (0, 1, 3), + (1, 2), + (0, 2, 3), + (1, 3), + (2, 3), + (1, 2, 3), + ] + actual = list(mi.interleave_evenly([a, b], lengths=[len_a, len_b])) + self.assertEqual(expected, actual) + + def test_no_length_raises(self): + # combinations doesn't have __len__, should trigger ValueError + iterables = [range(5), combinations(range(5), 2)] + with self.assertRaises(ValueError): + list(mi.interleave_evenly(iterables)) + + def test_argument_mismatch_raises(self): + # pass mismatching number of iterables and lengths + iterables = [range(3)] + lengths = [3, 4] + with self.assertRaises(ValueError): + list(mi.interleave_evenly(iterables, lengths=lengths)) + + +class InterleaveRandomlyTests(TestCase): + def test_basic(self): + seed(0) # For reproducibility + iterables = [1, 2, 3], 'abc', (True, False, None) + self.assertEqual( + list(mi.interleave_randomly(*iterables)), + ['a', 'b', 1, 'c', True, False, None, 2, 3], + ) + + def test_some_empty(self): + self.assertEqual( + list(mi.interleave_randomly([1, 2, 3], [], [])), + [1, 2, 3], + ) + self.assertEqual( + list(mi.interleave_randomly([], [1, 2, 3], [])), + [1, 2, 3], + ) + self.assertEqual( + list(mi.interleave_randomly([], [], [1, 2, 3])), + [1, 2, 3], + ) + + def test_all_empty(self): + iterables = [], [], [] + self.assertEqual(list(mi.interleave_randomly(*iterables)), []) + + def test_no_args(self): + self.assertEqual(list(mi.interleave_randomly()), []) + + def test_bad_type(self): + # Should raise TypeError if not all arguments are iterable + with self.assertRaises(TypeError): + list(mi.interleave_randomly(1, [2, 3], 'abc')) + + +class TestCollapse(TestCase): + """Tests for ``collapse()``""" + + def test_collapse(self): + l = [[1], 2, [[3], 4], [[[5]]]] + self.assertEqual(list(mi.collapse(l)), [1, 2, 3, 4, 5]) + + def test_collapse_to_string(self): + l = [["s1"], "s2", [["s3"], "s4"], [[["s5"]]]] + self.assertEqual(list(mi.collapse(l)), ["s1", "s2", "s3", "s4", "s5"]) + + def test_collapse_to_bytes(self): + l = [[b"s1"], b"s2", [[b"s3"], b"s4"], [[[b"s5"]]]] + self.assertEqual( + list(mi.collapse(l)), [b"s1", b"s2", b"s3", b"s4", b"s5"] + ) + + def test_collapse_flatten(self): + l = [[1], [2], [[3], 4], [[[5]]]] + self.assertEqual(list(mi.collapse(l, levels=1)), list(mi.flatten(l))) + + def test_collapse_to_level(self): + l = [[1], 2, [[3], 4], [[[5]]]] + self.assertEqual(list(mi.collapse(l, levels=2)), [1, 2, 3, 4, [5]]) + self.assertEqual( + list(mi.collapse(mi.collapse(l, levels=1), levels=1)), + list(mi.collapse(l, levels=2)), + ) + + def test_collapse_to_list(self): + l = (1, [2], (3, [4, (5,)], 'ab')) + actual = list(mi.collapse(l, base_type=list)) + expected = [1, [2], 3, [4, (5,)], 'ab'] + self.assertEqual(actual, expected) + + +class SideEffectTests(TestCase): + """Tests for ``side_effect()``""" + + def test_individual(self): + # The function increments the counter for each call + counter = [0] + + def func(arg): + counter[0] += 1 + + result = list(mi.side_effect(func, range(10))) + self.assertEqual(result, list(range(10))) + self.assertEqual(counter[0], 10) + + def test_chunked(self): + # The function increments the counter for each call + counter = [0] + + def func(arg): + counter[0] += 1 + + result = list(mi.side_effect(func, range(10), 2)) + self.assertEqual(result, list(range(10))) + self.assertEqual(counter[0], 5) + + def test_before_after(self): + f = StringIO() + collector = [] + + def func(item): + print(item, file=f) + collector.append(f.getvalue()) + + def it(): + yield 'a' + yield 'b' + raise RuntimeError('kaboom') + + before = lambda: print('HEADER', file=f) + after = f.close + + try: + mi.consume(mi.side_effect(func, it(), before=before, after=after)) + except RuntimeError: + pass + + # The iterable should have been written to the file + self.assertEqual(collector, ['HEADER\na\n', 'HEADER\na\nb\n']) + + # The file should be closed even though something bad happened + self.assertTrue(f.closed) + + def test_before_fails(self): + f = StringIO() + func = lambda x: print(x, file=f) + + def before(): + raise RuntimeError('ouch') + + try: + mi.consume( + mi.side_effect(func, 'abc', before=before, after=f.close) + ) + except RuntimeError: + pass + + # The file should be closed even though something bad happened in the + # before function + self.assertTrue(f.closed) + + +class SlicedTests(TestCase): + """Tests for ``sliced()``""" + + def test_even(self): + """Test when the length of the sequence is divisible by *n*""" + seq = 'ABCDEFGHI' + self.assertEqual(list(mi.sliced(seq, 3)), ['ABC', 'DEF', 'GHI']) + + def test_odd(self): + """Test when the length of the sequence is not divisible by *n*""" + seq = 'ABCDEFGHI' + self.assertEqual(list(mi.sliced(seq, 4)), ['ABCD', 'EFGH', 'I']) + + def test_not_sliceable(self): + seq = (x for x in 'ABCDEFGHI') + + with self.assertRaises(TypeError): + list(mi.sliced(seq, 3)) + + def test_odd_and_strict(self): + seq = [x for x in 'ABCDEFGHI'] + + with self.assertRaises(ValueError): + list(mi.sliced(seq, 4, strict=True)) + + def test_negative(self): + """Negative slice sizes should raise instead of silently + yielding a wrong result.""" + seq = 'ABCDEFG' + self.assertRaises(ValueError, lambda: list(mi.sliced(seq, -1))) + self.assertRaises( + ValueError, lambda: list(mi.sliced(seq, -1, strict=True)) + ) + + def test_numpy_like_array(self): + # Numpy arrays don't behave like Python lists - calling bool() + # on them doesn't return False for empty lists and True for non-empty + # ones. Emulate that behavior. + class FalseList(list): + def __getitem__(self, key): + ret = super().__getitem__(key) + if isinstance(key, slice): + return FalseList(ret) + + return ret + + def __bool__(self): + return False + + seq = FalseList(range(9)) + actual = list(mi.sliced(seq, 3)) + expected = [[0, 1, 2], [3, 4, 5], [6, 7, 8]] + self.assertEqual(actual, expected) + + +class SplitAtTests(TestCase): + def test_basic(self): + for iterable, separator in [ + ('a,bb,ccc,dddd', ','), + (',a,bb,ccc,dddd', ','), + ('a,bb,ccc,dddd,', ','), + ('a,bb,ccc,,dddd', ','), + ('', ','), + (',', ','), + ('a,bb,ccc,dddd', ';'), + ]: + with self.subTest(iterable=iterable, separator=separator): + it = iter(iterable) + pred = lambda x: x == separator + actual = [''.join(x) for x in mi.split_at(it, pred)] + expected = iterable.split(separator) + self.assertEqual(actual, expected) + + def test_maxsplit(self): + separator = ',' + pred = lambda x: x == separator + + for iterable in ['a,bb,ccc,dddd', '']: + for maxsplit in range(-1, 4): + with self.subTest(iterable=iterable, maxsplit=maxsplit): + it = iter(iterable) + result = mi.split_at(it, pred, maxsplit=maxsplit) + actual = [''.join(x) for x in result] + expected = iterable.split(separator, maxsplit) + self.assertEqual(actual, expected) + + def test_keep_separator(self): + separator = ',' + pred = lambda x: x == separator + + for iterable, expected in [ + ('a,bb,ccc', ['a', ',', 'bb', ',', 'ccc']), + (',a,bb,ccc', ['', ',', 'a', ',', 'bb', ',', 'ccc']), + ('a,bb,ccc,', ['a', ',', 'bb', ',', 'ccc', ',', '']), + ]: + with self.subTest(iterable=iterable): + it = iter(iterable) + result = mi.split_at(it, pred, keep_separator=True) + actual = [''.join(x) for x in result] + self.assertEqual(actual, expected) + + def test_combination(self): + iterable = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10] + pred = lambda x: x % 3 == 0 + actual = list( + mi.split_at(iterable, pred, maxsplit=2, keep_separator=True) + ) + expected = [[1, 2], [3], [4, 5], [6], [7, 8, 9, 10]] + self.assertEqual(actual, expected) + + +class SplitBeforeTest(TestCase): + """Tests for ``split_before()``""" + + def test_starts_with_sep(self): + actual = list(mi.split_before('xooxoo', lambda c: c == 'x')) + expected = [['x', 'o', 'o'], ['x', 'o', 'o']] + self.assertEqual(actual, expected) + + def test_ends_with_sep(self): + actual = list(mi.split_before('ooxoox', lambda c: c == 'x')) + expected = [['o', 'o'], ['x', 'o', 'o'], ['x']] + self.assertEqual(actual, expected) + + def test_no_sep(self): + actual = list(mi.split_before('ooo', lambda c: c == 'x')) + expected = [['o', 'o', 'o']] + self.assertEqual(actual, expected) + + def test_empty_collection(self): + for maxsplit in range(-1, 4): + with self.subTest(maxsplit=maxsplit): + actual = list( + mi.split_before([], lambda c: bool(c), maxsplit=maxsplit) + ) + expected = [] + self.assertEqual(actual, expected) + + def test_max_split(self): + for args, expected in [ + ( + ('a,b,c,d', lambda c: c == ',', -1), + [['a'], [',', 'b'], [',', 'c'], [',', 'd']], + ), + ( + ('a,b,c,d', lambda c: c == ',', 0), + [['a', ',', 'b', ',', 'c', ',', 'd']], + ), + ( + ('a,b,c,d', lambda c: c == ',', 1), + [['a'], [',', 'b', ',', 'c', ',', 'd']], + ), + ( + ('a,b,c,d', lambda c: c == ',', 2), + [['a'], [',', 'b'], [',', 'c', ',', 'd']], + ), + ( + ('a,b,c,d', lambda c: c == ',', 10), + [['a'], [',', 'b'], [',', 'c'], [',', 'd']], + ), + ( + ('a,b,c,d', lambda c: c == '@', 2), + [['a', ',', 'b', ',', 'c', ',', 'd']], + ), + ( + ('a,b,c,d', lambda c: c != ',', 2), + [['a', ','], ['b', ','], ['c', ',', 'd']], + ), + ]: + actual = list(mi.split_before(*args)) + self.assertEqual(actual, expected) + + +class SplitAfterTest(TestCase): + """Tests for ``split_after()``""" + + def test_starts_with_sep(self): + actual = list(mi.split_after('xooxoo', lambda c: c == 'x')) + expected = [['x'], ['o', 'o', 'x'], ['o', 'o']] + self.assertEqual(actual, expected) + + def test_ends_with_sep(self): + actual = list(mi.split_after('ooxoox', lambda c: c == 'x')) + expected = [['o', 'o', 'x'], ['o', 'o', 'x']] + self.assertEqual(actual, expected) + + def test_no_sep(self): + actual = list(mi.split_after('ooo', lambda c: c == 'x')) + expected = [['o', 'o', 'o']] + self.assertEqual(actual, expected) + + def test_empty_collection(self): + for maxsplit in range(-1, 4): + with self.subTest(maxsplit=maxsplit): + actual = list( + mi.split_after([], lambda c: bool(c), maxsplit=maxsplit) + ) + expected = [] + self.assertEqual(actual, expected) + + def test_max_split(self): + for args, expected in [ + ( + ('a,b,c,d', lambda c: c == ',', -1), + [['a', ','], ['b', ','], ['c', ','], ['d']], + ), + ( + ('a,b,c,d', lambda c: c == ',', 0), + [['a', ',', 'b', ',', 'c', ',', 'd']], + ), + ( + ('a,b,c,d', lambda c: c == ',', 1), + [['a', ','], ['b', ',', 'c', ',', 'd']], + ), + ( + ('a,b,c,d', lambda c: c == ',', 2), + [['a', ','], ['b', ','], ['c', ',', 'd']], + ), + ( + ('a,b,c,d', lambda c: c == ',', 10), + [['a', ','], ['b', ','], ['c', ','], ['d']], + ), + ( + ('a,b,c,d', lambda c: c == '@', 2), + [['a', ',', 'b', ',', 'c', ',', 'd']], + ), + ( + ('a,b,c,d', lambda c: c != ',', 2), + [['a'], [',', 'b'], [',', 'c', ',', 'd']], + ), + ( + ([1], lambda x: x == 1, 1), + [[1]], + ), + ]: + actual = list(mi.split_after(*args)) + self.assertEqual(actual, expected) + + +class SplitWhenTests(TestCase): + """Tests for ``split_when()``""" + + @staticmethod + def _split_when_before(iterable, pred): + return mi.split_when(iterable, lambda _, c: pred(c)) + + @staticmethod + def _split_when_after(iterable, pred): + return mi.split_when(iterable, lambda c, _: pred(c)) + + # split_before emulation + def test_before_emulation_starts_with_sep(self): + actual = list(self._split_when_before('xooxoo', lambda c: c == 'x')) + expected = [['x', 'o', 'o'], ['x', 'o', 'o']] + self.assertEqual(actual, expected) + + def test_before_emulation_ends_with_sep(self): + actual = list(self._split_when_before('ooxoox', lambda c: c == 'x')) + expected = [['o', 'o'], ['x', 'o', 'o'], ['x']] + self.assertEqual(actual, expected) + + def test_before_emulation_no_sep(self): + actual = list(self._split_when_before('ooo', lambda c: c == 'x')) + expected = [['o', 'o', 'o']] + self.assertEqual(actual, expected) + + # split_after emulation + def test_after_emulation_starts_with_sep(self): + actual = list(self._split_when_after('xooxoo', lambda c: c == 'x')) + expected = [['x'], ['o', 'o', 'x'], ['o', 'o']] + self.assertEqual(actual, expected) + + def test_after_emulation_ends_with_sep(self): + actual = list(self._split_when_after('ooxoox', lambda c: c == 'x')) + expected = [['o', 'o', 'x'], ['o', 'o', 'x']] + self.assertEqual(actual, expected) + + def test_after_emulation_no_sep(self): + actual = list(self._split_when_after('ooo', lambda c: c == 'x')) + expected = [['o', 'o', 'o']] + self.assertEqual(actual, expected) + + # edge cases + def test_empty_iterable(self): + for maxsplit in range(-1, 4): + with self.subTest(maxsplit=maxsplit): + actual = list( + mi.split_when('', lambda a, b: a != b, maxsplit=maxsplit) + ) + expected = [] + self.assertEqual(actual, expected) + + def test_one_element(self): + actual = list(mi.split_when('o', lambda a, b: a == b)) + expected = [['o']] + self.assertEqual(actual, expected) + + def test_one_element_is_second_item(self): + actual = list(self._split_when_before('x', lambda c: c == 'x')) + expected = [['x']] + self.assertEqual(actual, expected) + + def test_one_element_is_first_item(self): + actual = list(self._split_when_after('x', lambda c: c == 'x')) + expected = [['x']] + self.assertEqual(actual, expected) + + def test_max_split(self): + for args, expected in [ + ( + ('a,b,c,d', lambda a, _: a == ',', -1), + [['a', ','], ['b', ','], ['c', ','], ['d']], + ), + ( + ('a,b,c,d', lambda a, _: a == ',', 0), + [['a', ',', 'b', ',', 'c', ',', 'd']], + ), + ( + ('a,b,c,d', lambda _, b: b == ',', 1), + [['a'], [',', 'b', ',', 'c', ',', 'd']], + ), + ( + ('a,b,c,d', lambda a, _: a == ',', 2), + [['a', ','], ['b', ','], ['c', ',', 'd']], + ), + ( + ('0124376', lambda a, b: a > b, -1), + [['0', '1', '2', '4'], ['3', '7'], ['6']], + ), + ( + ('0124376', lambda a, b: a > b, 0), + [['0', '1', '2', '4', '3', '7', '6']], + ), + ( + ('0124376', lambda a, b: a > b, 1), + [['0', '1', '2', '4'], ['3', '7', '6']], + ), + ( + ('0124376', lambda a, b: a > b, 2), + [['0', '1', '2', '4'], ['3', '7'], ['6']], + ), + ]: + actual = list(mi.split_when(*args)) + self.assertEqual(actual, expected, str(args)) + + +class SplitIntoTests(TestCase): + """Tests for ``split_into()``""" + + def test_iterable_just_right(self): + """Size of ``iterable`` equals the sum of ``sizes``.""" + iterable = [1, 2, 3, 4, 5, 6, 7, 8, 9] + sizes = [2, 3, 4] + expected = [[1, 2], [3, 4, 5], [6, 7, 8, 9]] + actual = list(mi.split_into(iterable, sizes)) + self.assertEqual(actual, expected) + + def test_iterable_too_small(self): + """Size of ``iterable`` is smaller than sum of ``sizes``. Last return + list is shorter as a result.""" + iterable = [1, 2, 3, 4, 5, 6, 7] + sizes = [2, 3, 4] + expected = [[1, 2], [3, 4, 5], [6, 7]] + actual = list(mi.split_into(iterable, sizes)) + self.assertEqual(actual, expected) + + def test_iterable_too_small_extra(self): + """Size of ``iterable`` is smaller than sum of ``sizes``. Second last + return list is shorter and last return list is empty as a result.""" + iterable = [1, 2, 3, 4, 5, 6, 7] + sizes = [2, 3, 4, 5] + expected = [[1, 2], [3, 4, 5], [6, 7], []] + actual = list(mi.split_into(iterable, sizes)) + self.assertEqual(actual, expected) + + def test_iterable_too_large(self): + """Size of ``iterable`` is larger than sum of ``sizes``. Not all + items of iterable are returned.""" + iterable = [1, 2, 3, 4, 5, 6, 7, 8, 9] + sizes = [2, 3, 2] + expected = [[1, 2], [3, 4, 5], [6, 7]] + actual = list(mi.split_into(iterable, sizes)) + self.assertEqual(actual, expected) + + def test_using_none_with_leftover(self): + """Last item of ``sizes`` is None when items still remain in + ``iterable``. Last list returned stretches to fit all remaining items + of ``iterable``.""" + iterable = [1, 2, 3, 4, 5, 6, 7, 8, 9] + sizes = [2, 3, None] + expected = [[1, 2], [3, 4, 5], [6, 7, 8, 9]] + actual = list(mi.split_into(iterable, sizes)) + self.assertEqual(actual, expected) + + def test_using_none_without_leftover(self): + """Last item of ``sizes`` is None when no items remain in + ``iterable``. Last list returned is empty.""" + iterable = [1, 2, 3, 4, 5, 6, 7, 8, 9] + sizes = [2, 3, 4, None] + expected = [[1, 2], [3, 4, 5], [6, 7, 8, 9], []] + actual = list(mi.split_into(iterable, sizes)) + self.assertEqual(actual, expected) + + def test_using_none_mid_sizes(self): + """None is present in ``sizes`` but is not the last item. Last list + returned stretches to fit all remaining items of ``iterable`` but + all items in ``sizes`` after None are ignored.""" + iterable = [1, 2, 3, 4, 5, 6, 7, 8, 9] + sizes = [2, 3, None, 4] + expected = [[1, 2], [3, 4, 5], [6, 7, 8, 9]] + actual = list(mi.split_into(iterable, sizes)) + self.assertEqual(actual, expected) + + def test_iterable_empty(self): + """``iterable`` argument is empty but ``sizes`` is not. An empty + list is returned for each item in ``sizes``.""" + iterable = [] + sizes = [2, 4, 2] + expected = [[], [], []] + actual = list(mi.split_into(iterable, sizes)) + self.assertEqual(actual, expected) + + def test_iterable_empty_using_none(self): + """``iterable`` argument is empty but ``sizes`` is not. An empty + list is returned for each item in ``sizes`` that is not after a + None item.""" + iterable = [] + sizes = [2, 4, None, 2] + expected = [[], [], []] + actual = list(mi.split_into(iterable, sizes)) + self.assertEqual(actual, expected) + + def test_sizes_empty(self): + """``sizes`` argument is empty but ``iterable`` is not. An empty + generator is returned.""" + iterable = [1, 2, 3, 4, 5, 6, 7, 8, 9] + sizes = [] + expected = [] + actual = list(mi.split_into(iterable, sizes)) + self.assertEqual(actual, expected) + + def test_both_empty(self): + """Both ``sizes`` and ``iterable`` arguments are empty. An empty + generator is returned.""" + iterable = [] + sizes = [] + expected = [] + actual = list(mi.split_into(iterable, sizes)) + self.assertEqual(actual, expected) + + def test_bool_in_sizes(self): + """A bool object is present in ``sizes`` is treated as a 1 or 0 for + ``True`` or ``False`` due to bool being an instance of int.""" + iterable = [1, 2, 3, 4, 5, 6, 7, 8, 9] + sizes = [3, True, 2, False] + expected = [[1, 2, 3], [4], [5, 6], []] + actual = list(mi.split_into(iterable, sizes)) + self.assertEqual(actual, expected) + + def test_invalid_in_sizes(self): + """A ValueError is raised if an object in ``sizes`` is neither ``None`` + or an integer.""" + iterable = [1, 2, 3, 4, 5, 6, 7, 8, 9] + sizes = [1, [], 3] + with self.assertRaises(ValueError): + list(mi.split_into(iterable, sizes)) + + def test_invalid_in_sizes_after_none(self): + """A item in ``sizes`` that is invalid will not raise a TypeError if it + comes after a ``None`` item.""" + iterable = [1, 2, 3, 4, 5, 6, 7, 8, 9] + sizes = [3, 4, None, []] + expected = [[1, 2, 3], [4, 5, 6, 7], [8, 9]] + actual = list(mi.split_into(iterable, sizes)) + self.assertEqual(actual, expected) + + def test_generator_iterable_integrity(self): + """Check that if ``iterable`` is an iterator, it is consumed only by as + many items as the sum of ``sizes``.""" + iterable = (i for i in range(10)) + sizes = [2, 3] + + expected = [[0, 1], [2, 3, 4]] + actual = list(mi.split_into(iterable, sizes)) + self.assertEqual(actual, expected) + + iterable_expected = [5, 6, 7, 8, 9] + iterable_actual = list(iterable) + self.assertEqual(iterable_actual, iterable_expected) + + def test_generator_sizes_integrity(self): + """Check that if ``sizes`` is an iterator, it is consumed only until a + ``None`` item is reached""" + iterable = [1, 2, 3, 4, 5, 6, 7, 8, 9] + sizes = (i for i in [1, 2, None, 3, 4]) + + expected = [[1], [2, 3], [4, 5, 6, 7, 8, 9]] + actual = list(mi.split_into(iterable, sizes)) + self.assertEqual(actual, expected) + + sizes_expected = [3, 4] + sizes_actual = list(sizes) + self.assertEqual(sizes_actual, sizes_expected) + + +class PaddedTest(TestCase): + """Tests for ``padded()``""" + + def test_no_n(self): + seq = [1, 2, 3] + + # No fillvalue + self.assertEqual(mi.take(5, mi.padded(seq)), [1, 2, 3, None, None]) + + # With fillvalue + self.assertEqual( + mi.take(5, mi.padded(seq, fillvalue='')), [1, 2, 3, '', ''] + ) + + def test_invalid_n(self): + self.assertRaises(ValueError, lambda: list(mi.padded([1, 2, 3], n=-1))) + self.assertRaises(ValueError, lambda: list(mi.padded([1, 2, 3], n=0))) + + def test_valid_n(self): + seq = [1, 2, 3, 4, 5] + + # No need for padding: len(seq) <= n + self.assertEqual(list(mi.padded(seq, n=4)), [1, 2, 3, 4, 5]) + self.assertEqual(list(mi.padded(seq, n=5)), [1, 2, 3, 4, 5]) + + # No fillvalue + self.assertEqual( + list(mi.padded(seq, n=7)), [1, 2, 3, 4, 5, None, None] + ) + + # With fillvalue + self.assertEqual( + list(mi.padded(seq, fillvalue='', n=7)), [1, 2, 3, 4, 5, '', ''] + ) + + def test_next_multiple(self): + seq = [1, 2, 3, 4, 5, 6] + + # No need for padding: len(seq) % n == 0 + self.assertEqual( + list(mi.padded(seq, n=3, next_multiple=True)), [1, 2, 3, 4, 5, 6] + ) + + # Padding needed: len(seq) < n + self.assertEqual( + list(mi.padded(seq, n=8, next_multiple=True)), + [1, 2, 3, 4, 5, 6, None, None], + ) + + # No padding needed: len(seq) == n + self.assertEqual( + list(mi.padded(seq, n=6, next_multiple=True)), [1, 2, 3, 4, 5, 6] + ) + + # Padding needed: len(seq) > n + self.assertEqual( + list(mi.padded(seq, n=4, next_multiple=True)), + [1, 2, 3, 4, 5, 6, None, None], + ) + + # With fillvalue + self.assertEqual( + list(mi.padded(seq, fillvalue='', n=4, next_multiple=True)), + [1, 2, 3, 4, 5, 6, '', ''], + ) + + +class RepeatEachTests(TestCase): + """Tests for repeat_each()""" + + def test_default(self): + actual = list(mi.repeat_each('ABC')) + expected = ['A', 'A', 'B', 'B', 'C', 'C'] + self.assertEqual(actual, expected) + + def test_basic(self): + actual = list(mi.repeat_each('ABC', 3)) + expected = ['A', 'A', 'A', 'B', 'B', 'B', 'C', 'C', 'C'] + self.assertEqual(actual, expected) + + def test_empty(self): + actual = list(mi.repeat_each('')) + expected = [] + self.assertEqual(actual, expected) + + def test_no_repeat(self): + actual = list(mi.repeat_each('ABC', 0)) + expected = [] + self.assertEqual(actual, expected) + + def test_negative_repeat(self): + actual = list(mi.repeat_each('ABC', -1)) + expected = [] + self.assertEqual(actual, expected) + + def test_infinite_input(self): + repeater = mi.repeat_each(cycle('AB')) + actual = mi.take(6, repeater) + expected = ['A', 'A', 'B', 'B', 'A', 'A'] + self.assertEqual(actual, expected) + + +class RepeatLastTests(TestCase): + def test_empty_iterable(self): + slice_length = 3 + iterable = iter([]) + actual = mi.take(slice_length, mi.repeat_last(iterable)) + expected = [None] * slice_length + self.assertEqual(actual, expected) + + def test_default_value(self): + slice_length = 3 + iterable = iter([]) + default = '3' + actual = mi.take(slice_length, mi.repeat_last(iterable, default)) + expected = ['3'] * slice_length + self.assertEqual(actual, expected) + + def test_basic(self): + slice_length = 10 + iterable = (str(x) for x in range(5)) + actual = mi.take(slice_length, mi.repeat_last(iterable)) + expected = ['0', '1', '2', '3', '4', '4', '4', '4', '4', '4'] + self.assertEqual(actual, expected) + + +class DistributeTest(TestCase): + """Tests for distribute()""" + + def test_invalid_n(self): + self.assertRaises(ValueError, lambda: mi.distribute(-1, [1, 2, 3])) + self.assertRaises(ValueError, lambda: mi.distribute(0, [1, 2, 3])) + + def test_basic(self): + iterable = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10] + + for n, expected in [ + (1, [iterable]), + (2, [[1, 3, 5, 7, 9], [2, 4, 6, 8, 10]]), + (3, [[1, 4, 7, 10], [2, 5, 8], [3, 6, 9]]), + (10, [[n] for n in range(1, 10 + 1)]), + ]: + self.assertEqual( + [list(x) for x in mi.distribute(n, iterable)], expected + ) + + def test_large_n(self): + iterable = [1, 2, 3, 4] + self.assertEqual( + [list(x) for x in mi.distribute(6, iterable)], + [[1], [2], [3], [4], [], []], + ) + + +class StaggerTest(TestCase): + """Tests for ``stagger()``""" + + def test_default(self): + iterable = [0, 1, 2, 3] + actual = list(mi.stagger(iterable)) + expected = [(None, 0, 1), (0, 1, 2), (1, 2, 3)] + self.assertEqual(actual, expected) + + def test_offsets(self): + iterable = [0, 1, 2, 3] + for offsets, expected in [ + ((-2, 0, 2), [('', 0, 2), ('', 1, 3)]), + ((-2, -1), [('', ''), ('', 0), (0, 1), (1, 2), (2, 3)]), + ((1, 2), [(1, 2), (2, 3)]), + ]: + all_groups = mi.stagger(iterable, offsets=offsets, fillvalue='') + self.assertEqual(list(all_groups), expected) + + def test_longest(self): + iterable = [0, 1, 2, 3] + for offsets, expected in [ + ( + (-1, 0, 1), + [('', 0, 1), (0, 1, 2), (1, 2, 3), (2, 3, ''), (3, '', '')], + ), + ((-2, -1), [('', ''), ('', 0), (0, 1), (1, 2), (2, 3), (3, '')]), + ((1, 2), [(1, 2), (2, 3), (3, '')]), + ]: + all_groups = mi.stagger( + iterable, offsets=offsets, fillvalue='', longest=True + ) + self.assertEqual(list(all_groups), expected) + + +class ZipOffsetTest(TestCase): + """Tests for ``zip_offset()``""" + + def test_shortest(self): + a_1 = [0, 1, 2, 3] + a_2 = [0, 1, 2, 3, 4, 5] + a_3 = [0, 1, 2, 3, 4, 5, 6, 7] + actual = list( + mi.zip_offset(a_1, a_2, a_3, offsets=(-1, 0, 1), fillvalue='') + ) + expected = [('', 0, 1), (0, 1, 2), (1, 2, 3), (2, 3, 4), (3, 4, 5)] + self.assertEqual(actual, expected) + + def test_longest(self): + a_1 = [0, 1, 2, 3] + a_2 = [0, 1, 2, 3, 4, 5] + a_3 = [0, 1, 2, 3, 4, 5, 6, 7] + actual = list( + mi.zip_offset(a_1, a_2, a_3, offsets=(-1, 0, 1), longest=True) + ) + expected = [ + (None, 0, 1), + (0, 1, 2), + (1, 2, 3), + (2, 3, 4), + (3, 4, 5), + (None, 5, 6), + (None, None, 7), + ] + self.assertEqual(actual, expected) + + def test_mismatch(self): + iterables = [0, 1, 2], [2, 3, 4] + offsets = (-1, 0, 1) + self.assertRaises( + ValueError, + lambda: list(mi.zip_offset(*iterables, offsets=offsets)), + ) + + +class UnzipTests(TestCase): + """Tests for unzip()""" + + def test_empty_iterable(self): + self.assertEqual(list(mi.unzip([])), []) + # in reality zip([], [], []) is equivalent to iter([]) + # but it doesn't hurt to test both + self.assertEqual(list(mi.unzip(zip([], [], []))), []) + + def test_length_one_iterable(self): + xs, ys, zs = mi.unzip(zip([1], [2], [3])) + self.assertEqual(list(xs), [1]) + self.assertEqual(list(ys), [2]) + self.assertEqual(list(zs), [3]) + + def test_normal_case(self): + xs, ys, zs = range(10), range(1, 11), range(2, 12) + zipped = zip(xs, ys, zs) + xs, ys, zs = mi.unzip(zipped) + self.assertEqual(list(xs), list(range(10))) + self.assertEqual(list(ys), list(range(1, 11))) + self.assertEqual(list(zs), list(range(2, 12))) + + def test_improperly_zipped(self): + zipped = iter([(1, 2, 3), (4, 5), (6,)]) + xs, ys, zs = mi.unzip(zipped) + self.assertEqual(list(xs), [1, 4, 6]) + self.assertEqual(list(ys), [2, 5]) + self.assertEqual(list(zs), [3]) + + def test_increasingly_zipped(self): + zipped = iter([(1, 2), (3, 4, 5), (6, 7, 8, 9)]) + unzipped = mi.unzip(zipped) + # from the docstring: + # len(first tuple) is the number of iterables zipped + self.assertEqual(len(unzipped), 2) + xs, ys = unzipped + self.assertEqual(list(xs), [1, 3, 6]) + self.assertEqual(list(ys), [2, 4, 7]) + + +class SortTogetherTest(TestCase): + """Tests for sort_together()""" + + def test_key_list(self): + """tests `key_list` including default, iterables include duplicates""" + iterables = [ + ['GA', 'GA', 'GA', 'CT', 'CT', 'CT'], + ['May', 'Aug.', 'May', 'June', 'July', 'July'], + [97, 20, 100, 70, 100, 20], + ] + + self.assertEqual( + mi.sort_together(iterables), + [ + ('CT', 'CT', 'CT', 'GA', 'GA', 'GA'), + ('June', 'July', 'July', 'May', 'Aug.', 'May'), + (70, 100, 20, 97, 20, 100), + ], + ) + + self.assertEqual( + mi.sort_together(iterables, key_list=(0, 1)), + [ + ('CT', 'CT', 'CT', 'GA', 'GA', 'GA'), + ('July', 'July', 'June', 'Aug.', 'May', 'May'), + (100, 20, 70, 20, 97, 100), + ], + ) + + self.assertEqual( + mi.sort_together(iterables, key_list=(0, 1, 2)), + [ + ('CT', 'CT', 'CT', 'GA', 'GA', 'GA'), + ('July', 'July', 'June', 'Aug.', 'May', 'May'), + (20, 100, 70, 20, 97, 100), + ], + ) + + self.assertEqual( + mi.sort_together(iterables, key_list=(2,)), + [ + ('GA', 'CT', 'CT', 'GA', 'GA', 'CT'), + ('Aug.', 'July', 'June', 'May', 'May', 'July'), + (20, 20, 70, 97, 100, 100), + ], + ) + + def test_invalid_key_list(self): + """tests `key_list` for indexes not available in `iterables`""" + iterables = [ + ['GA', 'GA', 'GA', 'CT', 'CT', 'CT'], + ['May', 'Aug.', 'May', 'June', 'July', 'July'], + [97, 20, 100, 70, 100, 20], + ] + + self.assertRaises( + IndexError, lambda: mi.sort_together(iterables, key_list=(5,)) + ) + + def test_key_function(self): + """tests `key` function, including interaction with `key_list`""" + iterables = [ + ['GA', 'GA', 'GA', 'CT', 'CT', 'CT'], + ['May', 'Aug.', 'May', 'June', 'July', 'July'], + [97, 20, 100, 70, 100, 20], + ] + self.assertEqual( + mi.sort_together(iterables, key=lambda x: x), + [ + ('CT', 'CT', 'CT', 'GA', 'GA', 'GA'), + ('June', 'July', 'July', 'May', 'Aug.', 'May'), + (70, 100, 20, 97, 20, 100), + ], + ) + self.assertEqual( + mi.sort_together(iterables, key=lambda x: x[::-1]), + [ + ('GA', 'GA', 'GA', 'CT', 'CT', 'CT'), + ('May', 'Aug.', 'May', 'June', 'July', 'July'), + (97, 20, 100, 70, 100, 20), + ], + ) + self.assertEqual( + mi.sort_together( + iterables, + key_list=(0, 2), + key=lambda state, number: ( + number if state == 'CT' else 2 * number + ), + ), + [ + ('CT', 'GA', 'CT', 'CT', 'GA', 'GA'), + ('July', 'Aug.', 'June', 'July', 'May', 'May'), + (20, 20, 70, 100, 97, 100), + ], + ) + + def test_reverse(self): + """tests `reverse` to ensure a reverse sort for `key_list` iterables""" + iterables = [ + ['GA', 'GA', 'GA', 'CT', 'CT', 'CT'], + ['May', 'Aug.', 'May', 'June', 'July', 'July'], + [97, 20, 100, 70, 100, 20], + ] + + self.assertEqual( + mi.sort_together(iterables, key_list=(0, 1, 2), reverse=True), + [ + ('GA', 'GA', 'GA', 'CT', 'CT', 'CT'), + ('May', 'May', 'Aug.', 'June', 'July', 'July'), + (100, 97, 20, 70, 100, 20), + ], + ) + + def test_uneven_iterables(self): + """tests trimming of iterables to the shortest length before sorting""" + iterables = [ + ['GA', 'GA', 'GA', 'CT', 'CT', 'CT', 'MA'], + ['May', 'Aug.', 'May', 'June', 'July', 'July'], + [97, 20, 100, 70, 100, 20, 0], + ] + + self.assertEqual( + mi.sort_together(iterables), + [ + ('CT', 'CT', 'CT', 'GA', 'GA', 'GA'), + ('June', 'July', 'July', 'May', 'Aug.', 'May'), + (70, 100, 20, 97, 20, 100), + ], + ) + + def test_strict(self): + # Test for list of lists or tuples + self.assertRaises( + ValueError, + lambda: mi.sort_together( + [(4, 3, 2, 1), ('a', 'b', 'c')], strict=True + ), + ) + + # Test for list of iterables + self.assertRaises( + ValueError, + lambda: mi.sort_together([range(4), range(5)], strict=True), + ) + + # Test for iterable of iterables + self.assertRaises( + ValueError, + lambda: mi.sort_together( + (range(i) for i in range(4)), strict=True + ), + ) + + +class DivideTest(TestCase): + """Tests for divide()""" + + def test_invalid_n(self): + self.assertRaises(ValueError, lambda: mi.divide(-1, [1, 2, 3])) + self.assertRaises(ValueError, lambda: mi.divide(0, [1, 2, 3])) + + def test_basic(self): + iterable = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10] + + for n, expected in [ + (1, [iterable]), + (2, [[1, 2, 3, 4, 5], [6, 7, 8, 9, 10]]), + (3, [[1, 2, 3, 4], [5, 6, 7], [8, 9, 10]]), + (10, [[n] for n in range(1, 10 + 1)]), + ]: + self.assertEqual( + [list(x) for x in mi.divide(n, iterable)], expected + ) + + def test_large_n(self): + self.assertEqual( + [list(x) for x in mi.divide(6, iter(range(1, 4 + 1)))], + [[1], [2], [3], [4], [], []], + ) + + +class TestAlwaysIterable(TestCase): + """Tests for always_iterable()""" + + def test_single(self): + self.assertEqual(list(mi.always_iterable(1)), [1]) + + def test_strings(self): + for obj in ['foo', b'bar', 'baz']: + actual = list(mi.always_iterable(obj)) + expected = [obj] + self.assertEqual(actual, expected) + + def test_base_type(self): + dict_obj = {'a': 1, 'b': 2} + str_obj = '123' + + # Default: dicts are iterable like they normally are + default_actual = list(mi.always_iterable(dict_obj)) + default_expected = list(dict_obj) + self.assertEqual(default_actual, default_expected) + + # Unitary types set: dicts are not iterable + custom_actual = list(mi.always_iterable(dict_obj, base_type=dict)) + custom_expected = [dict_obj] + self.assertEqual(custom_actual, custom_expected) + + # With unitary types set, strings are iterable + str_actual = list(mi.always_iterable(str_obj, base_type=None)) + str_expected = list(str_obj) + self.assertEqual(str_actual, str_expected) + + # base_type handles nested tuple (via isinstance). + base_type = ((dict,),) + custom_actual = list(mi.always_iterable(dict_obj, base_type=base_type)) + custom_expected = [dict_obj] + self.assertEqual(custom_actual, custom_expected) + + def test_iterables(self): + self.assertEqual(list(mi.always_iterable([0, 1])), [0, 1]) + self.assertEqual( + list(mi.always_iterable([0, 1], base_type=list)), [[0, 1]] + ) + self.assertEqual( + list(mi.always_iterable(iter('foo'))), ['f', 'o', 'o'] + ) + self.assertEqual(list(mi.always_iterable([])), []) + + def test_none(self): + self.assertEqual(list(mi.always_iterable(None)), []) + + def test_generator(self): + def _gen(): + yield 0 + yield 1 + + self.assertEqual(list(mi.always_iterable(_gen())), [0, 1]) + + +class AdjacentTests(TestCase): + def test_typical(self): + actual = list(mi.adjacent(lambda x: x % 5 == 0, range(10))) + expected = [ + (True, 0), + (True, 1), + (False, 2), + (False, 3), + (True, 4), + (True, 5), + (True, 6), + (False, 7), + (False, 8), + (False, 9), + ] + self.assertEqual(actual, expected) + + def test_empty_iterable(self): + actual = list(mi.adjacent(lambda x: x % 5 == 0, [])) + expected = [] + self.assertEqual(actual, expected) + + def test_length_one(self): + actual = list(mi.adjacent(lambda x: x % 5 == 0, [0])) + expected = [(True, 0)] + self.assertEqual(actual, expected) + + actual = list(mi.adjacent(lambda x: x % 5 == 0, [1])) + expected = [(False, 1)] + self.assertEqual(actual, expected) + + def test_consecutive_true(self): + """Test that when the predicate matches multiple consecutive elements + it doesn't repeat elements in the output""" + actual = list(mi.adjacent(lambda x: x % 5 < 2, range(10))) + expected = [ + (True, 0), + (True, 1), + (True, 2), + (False, 3), + (True, 4), + (True, 5), + (True, 6), + (True, 7), + (False, 8), + (False, 9), + ] + self.assertEqual(actual, expected) + + def test_distance(self): + actual = list(mi.adjacent(lambda x: x % 5 == 0, range(10), distance=2)) + expected = [ + (True, 0), + (True, 1), + (True, 2), + (True, 3), + (True, 4), + (True, 5), + (True, 6), + (True, 7), + (False, 8), + (False, 9), + ] + self.assertEqual(actual, expected) + + actual = list(mi.adjacent(lambda x: x % 5 == 0, range(10), distance=3)) + expected = [ + (True, 0), + (True, 1), + (True, 2), + (True, 3), + (True, 4), + (True, 5), + (True, 6), + (True, 7), + (True, 8), + (False, 9), + ] + self.assertEqual(actual, expected) + + def test_large_distance(self): + """Test distance larger than the length of the iterable""" + iterable = range(10) + actual = list(mi.adjacent(lambda x: x % 5 == 4, iterable, distance=20)) + expected = list(zip(repeat(True), iterable)) + self.assertEqual(actual, expected) + + actual = list(mi.adjacent(lambda x: False, iterable, distance=20)) + expected = list(zip(repeat(False), iterable)) + self.assertEqual(actual, expected) + + def test_zero_distance(self): + """Test that adjacent() reduces to zip+map when distance is 0""" + iterable = range(1000) + predicate = lambda x: x % 4 == 2 + actual = mi.adjacent(predicate, iterable, 0) + expected = zip(map(predicate, iterable), iterable) + self.assertTrue(all(a == e for a, e in zip(actual, expected))) + + def test_negative_distance(self): + """Test that adjacent() raises an error with negative distance""" + pred = lambda x: x + self.assertRaises( + ValueError, lambda: mi.adjacent(pred, range(1000), -1) + ) + self.assertRaises( + ValueError, lambda: mi.adjacent(pred, range(10), -10) + ) + + def test_grouping(self): + """Test interaction of adjacent() with groupby_transform()""" + iterable = mi.adjacent(lambda x: x % 5 == 0, range(10)) + grouper = mi.groupby_transform(iterable, itemgetter(0), itemgetter(1)) + actual = [(k, list(g)) for k, g in grouper] + expected = [ + (True, [0, 1]), + (False, [2, 3]), + (True, [4, 5, 6]), + (False, [7, 8, 9]), + ] + self.assertEqual(actual, expected) + + def test_call_once(self): + """Test that the predicate is only called once per item.""" + already_seen = set() + iterable = range(10) + + def predicate(item): + self.assertNotIn(item, already_seen) + already_seen.add(item) + return True + + actual = list(mi.adjacent(predicate, iterable)) + expected = [(True, x) for x in iterable] + self.assertEqual(actual, expected) + + +class GroupByTransformTests(TestCase): + def assertAllGroupsEqual(self, groupby1, groupby2): + for a, b in zip(groupby1, groupby2): + key1, group1 = a + key2, group2 = b + self.assertEqual(key1, key2) + self.assertListEqual(list(group1), list(group2)) + self.assertRaises(StopIteration, lambda: next(groupby1)) + self.assertRaises(StopIteration, lambda: next(groupby2)) + + def test_default_funcs(self): + iterable = [(x // 5, x) for x in range(1000)] + actual = mi.groupby_transform(iterable) + expected = groupby(iterable) + self.assertAllGroupsEqual(actual, expected) + + def test_valuefunc(self): + iterable = [(int(x / 5), int(x / 3), x) for x in range(10)] + + # Test the standard usage of grouping one iterable using another's keys + grouper = mi.groupby_transform( + iterable, keyfunc=itemgetter(0), valuefunc=itemgetter(-1) + ) + actual = [(k, list(g)) for k, g in grouper] + expected = [(0, [0, 1, 2, 3, 4]), (1, [5, 6, 7, 8, 9])] + self.assertEqual(actual, expected) + + grouper = mi.groupby_transform( + iterable, keyfunc=itemgetter(1), valuefunc=itemgetter(-1) + ) + actual = [(k, list(g)) for k, g in grouper] + expected = [(0, [0, 1, 2]), (1, [3, 4, 5]), (2, [6, 7, 8]), (3, [9])] + self.assertEqual(actual, expected) + + # and now for something a little different + d = dict(zip(range(10), 'abcdefghij')) + grouper = mi.groupby_transform( + range(10), keyfunc=lambda x: x // 5, valuefunc=d.get + ) + actual = [(k, ''.join(g)) for k, g in grouper] + expected = [(0, 'abcde'), (1, 'fghij')] + self.assertEqual(actual, expected) + + def test_no_valuefunc(self): + iterable = range(1000) + + def key(x): + return x // 5 + + actual = mi.groupby_transform(iterable, key, valuefunc=None) + expected = groupby(iterable, key) + self.assertAllGroupsEqual(actual, expected) + + actual = mi.groupby_transform(iterable, key) # default valuefunc + expected = groupby(iterable, key) + self.assertAllGroupsEqual(actual, expected) + + def test_reducefunc(self): + iterable = range(50) + keyfunc = lambda k: 10 * (k // 10) + valuefunc = lambda v: v + 1 + reducefunc = sum + actual = list( + mi.groupby_transform( + iterable, + keyfunc=keyfunc, + valuefunc=valuefunc, + reducefunc=reducefunc, + ) + ) + expected = [(0, 55), (10, 155), (20, 255), (30, 355), (40, 455)] + self.assertEqual(actual, expected) + + +class NumericRangeTests(TestCase): + def test_basic(self): + for args, expected in [ + ((4,), [0, 1, 2, 3]), + ((4.0,), [0.0, 1.0, 2.0, 3.0]), + ((1.0, 4), [1.0, 2.0, 3.0]), + ((1, 4.0), [1.0, 2.0, 3.0]), + ((1.0, 5), [1.0, 2.0, 3.0, 4.0]), + ((0, 20, 5), [0, 5, 10, 15]), + ((0, 20, 5.0), [0.0, 5.0, 10.0, 15.0]), + ((0, 10, 3), [0, 3, 6, 9]), + ((0, 10, 3.0), [0.0, 3.0, 6.0, 9.0]), + ((0, -5, -1), [0, -1, -2, -3, -4]), + ((0.0, -5, -1), [0.0, -1.0, -2.0, -3.0, -4.0]), + ((1, 2, Fraction(1, 2)), [Fraction(1, 1), Fraction(3, 2)]), + ((0,), []), + ((0.0,), []), + ((1, 0), []), + ((1.0, 0.0), []), + ((0.1, 0.30000000000000001, 0.2), [0.1]), # IEEE 754 ! + ( + ( + Decimal("0.1"), + Decimal("0.30000000000000001"), + Decimal("0.2"), + ), + [Decimal("0.1"), Decimal("0.3")], + ), # okay with Decimal + ( + ( + Fraction(1, 10), + Fraction(30000000000000001, 100000000000000000), + Fraction(2, 10), + ), + [Fraction(1, 10), Fraction(3, 10)], + ), # okay with Fraction + ((Fraction(2, 1),), [Fraction(0, 1), Fraction(1, 1)]), + ((Decimal('2.0'),), [Decimal('0.0'), Decimal('1.0')]), + ( + ( + datetime(2019, 3, 29, 12, 34, 56), + datetime(2019, 3, 29, 12, 37, 55), + timedelta(minutes=1), + ), + [ + datetime(2019, 3, 29, 12, 34, 56), + datetime(2019, 3, 29, 12, 35, 56), + datetime(2019, 3, 29, 12, 36, 56), + ], + ), + ]: + actual = list(mi.numeric_range(*args)) + self.assertEqual(expected, actual) + self.assertTrue( + all(type(a) is type(e) for a, e in zip(actual, expected)) + ) + + def test_arg_count(self): + for args, message in [ + ((), 'numeric_range expected at least 1 argument, got 0'), + ( + (0, 1, 2, 3), + 'numeric_range expected at most 3 arguments, got 4', + ), + ]: + with self.assertRaisesRegex(TypeError, message): + mi.numeric_range(*args) + + def test_zero_step(self): + for args in [ + (1, 2, 0), + ( + datetime(2019, 3, 29, 12, 34, 56), + datetime(2019, 3, 29, 12, 37, 55), + timedelta(minutes=0), + ), + (1.0, 2.0, 0.0), + (Decimal("1.0"), Decimal("2.0"), Decimal("0.0")), + (Fraction(2, 2), Fraction(4, 2), Fraction(0, 2)), + ]: + with self.assertRaises(ValueError): + list(mi.numeric_range(*args)) + + def test_bool(self): + for args, expected in [ + ((1.0, 3.0, 1.5), True), + ((1.0, 2.0, 1.5), True), + ((1.0, 1.0, 1.5), False), + ((1.0, 0.0, 1.5), False), + ((3.0, 1.0, -1.5), True), + ((2.0, 1.0, -1.5), True), + ((1.0, 1.0, -1.5), False), + ((0.0, 1.0, -1.5), False), + ((Decimal("1.0"), Decimal("2.0"), Decimal("1.5")), True), + ((Decimal("1.0"), Decimal("0.0"), Decimal("1.5")), False), + ((Fraction(2, 2), Fraction(4, 2), Fraction(3, 2)), True), + ((Fraction(2, 2), Fraction(0, 2), Fraction(3, 2)), False), + ( + ( + datetime(2019, 3, 29), + datetime(2019, 3, 30), + timedelta(hours=1), + ), + True, + ), + ( + ( + datetime(2019, 3, 29), + datetime(2019, 3, 28), + timedelta(hours=1), + ), + False, + ), + ]: + self.assertEqual(expected, bool(mi.numeric_range(*args))) + + def test_contains(self): + for args, expected_in, expected_not_in in [ + ((10,), range(10), (0.5,)), + ((1.0, 9.9, 1.5), (1.0, 2.5, 4.0, 5.5, 7.0, 8.5), (0.9,)), + ((9.0, 1.0, -1.5), (1.5, 3.0, 4.5, 6.0, 7.5, 9.0), (0.0, 0.9)), + ( + (Decimal("1.0"), Decimal("9.9"), Decimal("1.5")), + ( + Decimal("1.0"), + Decimal("2.5"), + Decimal("4.0"), + Decimal("5.5"), + Decimal("7.0"), + Decimal("8.5"), + ), + (Decimal("0.9"),), + ), + ( + (Fraction(0, 1), Fraction(5, 1), Fraction(1, 2)), + (Fraction(0, 1), Fraction(1, 2), Fraction(9, 2)), + (Fraction(10, 2),), + ), + ( + ( + datetime(2019, 3, 29), + datetime(2019, 3, 30), + timedelta(hours=1), + ), + (datetime(2019, 3, 29, 15),), + (datetime(2019, 3, 29, 15, 30),), + ), + ]: + r = mi.numeric_range(*args) + for v in expected_in: + self.assertTrue(v in r) + self.assertFalse(v not in r) + + for v in expected_not_in: + self.assertFalse(v in r) + self.assertTrue(v not in r) + + def test_contains_items_with_inexact_step(self): + # The items are produced by repeated multiplication, so with a step + # that has no exact float representation they do not land on an exact + # multiple of it. Every item must still be a member, and index() must + # locate it where __getitem__ puts it. + for args in [ + (0.0, 1.0, 0.1), + (0.0, 3.0, 0.7), + (1.0, 0.0, -0.1), + (2.1, 5.1, 0.5), + (-1.0, 1.0, 0.25), + ]: + r = mi.numeric_range(*args) + for i, v in enumerate(r): + self.assertIn(v, r) + self.assertEqual(i, r.index(v)) + self.assertEqual(v, r[i]) + self.assertEqual(1, r.count(v)) + + def test_contains_non_items_with_inexact_step(self): + # Values the range does not produce stay out of it. + r = mi.numeric_range(0.0, 1.0, 0.1) + items = set(r) + for i in range(1000): + v = i / 1000 + if v not in items: + self.assertNotIn(v, r) + self.assertRaises(ValueError, r.index, v) + + def test_eq(self): + # numeric_range objects are equal to themselves + range_1 = mi.numeric_range(2, 4, 1) + self.assertEqual(range_1, range_1) + + # ...and to distinct objects covering the same values + self.assertEqual(range_1, mi.numeric_range(2, 4, 1)) + + # numeric_range objects are not equal other types of objects + self.assertNotEqual(mi.numeric_range(7.0), 1) + self.assertNotEqual(mi.numeric_range(7.0), "abc") + + for args1, args2 in [ + ((0, 5, 2), (0, 6, 2)), + ((1.0, 9.9, 1.5), (1.0, 8.6, 1.5)), + ((8.5, 0.0, -1.5), (8.5, 0.7, -1.5)), + ((7.0, 0.0, 1.0), (17.0, 7.0, 0.5)), + ( + (Decimal("1.0"), Decimal("9.9"), Decimal("1.5")), + (Decimal("1.0"), Decimal("8.6"), Decimal("1.5")), + ), + ( + (Fraction(1, 1), Fraction(10, 1), Fraction(3, 2)), + (Fraction(1, 1), Fraction(9, 1), Fraction(3, 2)), + ), + ( + ( + datetime(2019, 3, 29), + datetime(2019, 3, 30), + timedelta(hours=10), + ), + ( + datetime(2019, 3, 29), + datetime(2019, 3, 30, 1), + timedelta(hours=10), + ), + ), + ((2, 3, 1), (2, 3, 5)), + ]: + range1 = mi.numeric_range(*args1) + range2 = mi.numeric_range(*args2) + self.assertEqual(range1, range2) + self.assertEqual(hash(range1), hash(range2)) + + for args1, args2 in [ + ((0, 5, 2), (0, 7, 2)), + ((1.0, 9.9, 1.5), (1.2, 9.9, 1.5)), + ((1.0, 9.9, 1.5), (1.0, 10.3, 1.5)), + ((1.0, 9.9, 1.5), (1.0, 9.9, 1.4)), + ((8.5, 0.0, -1.5), (8.4, 0.0, -1.5)), + ((8.5, 0.0, -1.5), (8.5, -0.7, -1.5)), + ((8.5, 0.0, -1.5), (8.5, 0.0, -1.4)), + ((0.0, 7.0, 1.0), (7.0, 0.0, 1.0)), + ( + (Decimal("1.0"), Decimal("10.0"), Decimal("1.5")), + (Decimal("1.0"), Decimal("10.5"), Decimal("1.5")), + ), + ( + (Fraction(1, 1), Fraction(10, 1), Fraction(3, 2)), + (Fraction(1, 1), Fraction(21, 2), Fraction(3, 2)), + ), + ( + ( + datetime(2019, 3, 29), + datetime(2019, 3, 30), + timedelta(hours=10), + ), + ( + datetime(2019, 3, 29), + datetime(2019, 3, 30, 15), + timedelta(hours=10), + ), + ), + ]: + range1 = mi.numeric_range(*args1) + range2 = mi.numeric_range(*args2) + self.assertNotEqual(range1, range2) + self.assertNotEqual(hash(range1), hash(range2)) + + def test_get_item_by_index(self): + for args, index, expected in [ + ((1, 6), 2, 3), + ((1.0, 6.0, 1.5), 0, 1.0), + ((1.0, 6.0, 1.5), 1, 2.5), + ((1.0, 6.0, 1.5), 2, 4.0), + ((1.0, 6.0, 1.5), 3, 5.5), + ((1.0, 6.0, 1.5), -1, 5.5), + ((1.0, 6.0, 1.5), -2, 4.0), + ( + (Decimal("1.0"), Decimal("9.0"), Decimal("1.5")), + -1, + Decimal("8.5"), + ), + ( + (Fraction(1, 1), Fraction(10, 1), Fraction(3, 2)), + 2, + Fraction(4, 1), + ), + ( + ( + datetime(2019, 3, 29), + datetime(2019, 3, 30), + timedelta(hours=10), + ), + 1, + datetime(2019, 3, 29, 10), + ), + ]: + self.assertEqual(expected, mi.numeric_range(*args)[index]) + + for args, index in [ + ((1.0, 6.0, 1.5), 4), + ((1.0, 6.0, 1.5), -5), + ((6.0, 1.0, 1.5), 0), + ((6.0, 1.0, 1.5), -1), + ((Decimal("1.0"), Decimal("9.0"), Decimal("-1.5")), -1), + ((Fraction(1, 1), Fraction(2, 1), Fraction(3, 2)), 2), + ( + ( + datetime(2019, 3, 29), + datetime(2019, 3, 30), + timedelta(hours=10), + ), + 8, + ), + ]: + with self.assertRaises(IndexError): + mi.numeric_range(*args)[index] + + def test_get_item_by_slice(self): + for args, sl, expected_args in [ + ((1.0, 9.0, 1.5), slice(None, None, None), (1.0, 9.0, 1.5)), + ((1.0, 9.0, 1.5), slice(None, 1, None), (1.0, 2.5, 1.5)), + ((1.0, 9.0, 1.5), slice(None, None, 2), (1.0, 9.0, 3.0)), + ((1.0, 9.0, 1.5), slice(None, 2, None), (1.0, 4.0, 1.5)), + ((1.0, 9.0, 1.5), slice(1, 2, None), (2.5, 4.0, 1.5)), + ((1.0, 9.0, 1.5), slice(1, -1, None), (2.5, 8.5, 1.5)), + ((1.0, 9.0, 1.5), slice(10, None, 3), (9.0, 9.0, 4.5)), + ((1.0, 9.0, 1.5), slice(-10, None, 3), (1.0, 9.0, 4.5)), + ((1.0, 9.0, 1.5), slice(None, -10, 3), (1.0, 1.0, 4.5)), + ((1.0, 9.0, 1.5), slice(None, 10, 3), (1.0, 9.0, 4.5)), + ( + (10.0, 0.0, -2.0), + slice(None, None, -1), + (2.0, 12.0, 2.0), + ), + ( + (Decimal("1.0"), Decimal("9.0"), Decimal("1.5")), + slice(1, -1, None), + (Decimal("2.5"), Decimal("8.5"), Decimal("1.5")), + ), + ( + (Fraction(1, 1), Fraction(5, 1), Fraction(3, 2)), + slice(1, -1, None), + (Fraction(5, 2), Fraction(4, 1), Fraction(3, 2)), + ), + ( + ( + datetime(2019, 3, 29), + datetime(2019, 3, 30), + timedelta(hours=10), + ), + slice(1, -1, None), + ( + datetime(2019, 3, 29, 10), + datetime(2019, 3, 29, 20), + timedelta(hours=10), + ), + ), + ]: + self.assertEqual( + mi.numeric_range(*expected_args), mi.numeric_range(*args)[sl] + ) + + def test_iter_twice(self): + r1 = mi.numeric_range(1.0, 9.9, 1.5) + r2 = mi.numeric_range(8.5, 0.0, -1.5) + self.assertEqual([1.0, 2.5, 4.0, 5.5, 7.0, 8.5], list(r1)) + self.assertEqual([1.0, 2.5, 4.0, 5.5, 7.0, 8.5], list(r1)) + self.assertEqual([8.5, 7.0, 5.5, 4.0, 2.5, 1.0], list(r2)) + self.assertEqual([8.5, 7.0, 5.5, 4.0, 2.5, 1.0], list(r2)) + + def test_len(self): + for args, expected in [ + ((1.0, 7.0, 1.5), 4), + ((1.0, 7.01, 1.5), 5), + ((7.0, 1.0, -1.5), 4), + ((7.01, 1.0, -1.5), 5), + ((0.1, 0.30000000000000001, 0.2), 1), # IEEE 754 ! + ( + ( + Decimal("0.1"), + Decimal("0.30000000000000001"), + Decimal("0.2"), + ), + 2, + ), # works with Decimal + ((Decimal("1.0"), Decimal("9.0"), Decimal("1.5")), 6), + ((Fraction(1, 1), Fraction(5, 1), Fraction(3, 2)), 3), + ( + ( + datetime(2019, 3, 29), + datetime(2019, 3, 30), + timedelta(hours=10), + ), + 3, + ), + ]: + self.assertEqual(expected, len(mi.numeric_range(*args))) + + def test_len_matches_iteration(self): + # len() is derived by division while the items come from repeated + # multiplication; the two must not disagree, or the last index reports + # an item the range never yields. + for args in [ + # the division overshoots: len() counted an item past the stop + (-8.732, -13.532, -2.4), + (4.51, 10.375, 0.255), + # the division falls short: len() missed an item the range yields + (-69.37276490678778, -29.8339073344664, 4.942357196540172), + (32.50227965239458, 9.003418351482358, -1.9582384417426848), + (0.0, 1.0, 0.1), + (0.0, 3.0, 0.7), + (10.0, -10.0, -0.7), + (2.1, 5.1, 0.5), + ]: + r = mi.numeric_range(*args) + items = list(r) + self.assertEqual(len(items), len(r)) + if items: + self.assertEqual(items[-1], r[len(r) - 1]) + + def test_repr(self): + for args, *expected in [ + ((7.0,), "numeric_range(0.0, 7.0)"), + ((1.0, 7.0), "numeric_range(1.0, 7.0)"), + ((7.0, 1.0, -1.5), "numeric_range(7.0, 1.0, -1.5)"), + ( + (Decimal("1.0"), Decimal("9.0"), Decimal("1.5")), + ( + "numeric_range(Decimal('1.0'), Decimal('9.0'), " + "Decimal('1.5'))" + ), + ), + ( + (Fraction(7, 7), Fraction(10, 2), Fraction(3, 2)), + ( + "numeric_range(Fraction(1, 1), Fraction(5, 1), " + "Fraction(3, 2))" + ), + ), + ( + ( + datetime(2019, 3, 29), + datetime(2019, 3, 30), + timedelta(hours=10), + ), + "numeric_range(datetime.datetime(2019, 3, 29, 0, 0), " + "datetime.datetime(2019, 3, 30, 0, 0), " + "datetime.timedelta(seconds=36000))", + "numeric_range(datetime.datetime(2019, 3, 29, 0, 0), " + "datetime.datetime(2019, 3, 30, 0, 0), " + "datetime.timedelta(0, 36000))", + ), + ]: + with self.subTest(args=args): + self.assertIn(repr(mi.numeric_range(*args)), expected) + + def test_reversed(self): + for args, expected in [ + ((7.0,), [6.0, 5.0, 4.0, 3.0, 2.0, 1.0, 0.0]), + ((1.0, 7.0), [6.0, 5.0, 4.0, 3.0, 2.0, 1.0]), + ((7.0, 1.0, -1.5), [2.5, 4.0, 5.5, 7.0]), + ((7.0, 0.9, -1.5), [1.0, 2.5, 4.0, 5.5, 7.0]), + ( + (Decimal("1.0"), Decimal("5.0"), Decimal("1.5")), + [Decimal('4.0'), Decimal('2.5'), Decimal('1.0')], + ), + ( + (Fraction(1, 1), Fraction(5, 1), Fraction(3, 2)), + [Fraction(4, 1), Fraction(5, 2), Fraction(1, 1)], + ), + ( + ( + datetime(2019, 3, 29), + datetime(2019, 3, 30), + timedelta(hours=10), + ), + [ + datetime(2019, 3, 29, 20), + datetime(2019, 3, 29, 10), + datetime(2019, 3, 29), + ], + ), + ]: + self.assertEqual(expected, list(reversed(mi.numeric_range(*args)))) + + def test_reversed_preserves_float_values(self): + for args in [ + (0.0, 1.0, 0.1), + (1.0, 0.0, -0.1), + (0.1, 0.5, 0.1), + ]: + with self.subTest(args=args): + values = mi.numeric_range(*args) + self.assertEqual(list(values)[::-1], list(reversed(values))) + + def test_reversed_datetime_limits(self): + for start, step in [ + (datetime.min, timedelta(days=1)), + (datetime.max, -timedelta(days=1)), + ]: + with self.subTest(start=start, step=step): + values = mi.numeric_range(start, start + 2 * step, step) + self.assertEqual(list(values)[::-1], list(reversed(values))) + + def test_count(self): + for args, v, c in [ + ((7.0,), 0.0, 1), + ((7.0,), 0.5, 0), + ((7.0,), 6.0, 1), + ((7.0,), 7.0, 0), + ((7.0,), 10.0, 0), + ( + (Decimal("1.0"), Decimal("5.0"), Decimal("1.5")), + Decimal('4.0'), + 1, + ), + ( + (Fraction(1, 1), Fraction(5, 1), Fraction(3, 2)), + Fraction(5, 2), + 1, + ), + ( + ( + datetime(2019, 3, 29), + datetime(2019, 3, 30), + timedelta(hours=10), + ), + datetime(2019, 3, 29, 20), + 1, + ), + ]: + self.assertEqual(c, mi.numeric_range(*args).count(v)) + + def test_index(self): + for args, v, i in [ + ((7.0,), 0.0, 0), + ((7.0,), 6.0, 6), + ((7.0, 0.0, -1.0), 7.0, 0), + ((7.0, 0.0, -1.0), 1.0, 6), + ( + (Decimal("1.0"), Decimal("5.0"), Decimal("1.5")), + Decimal('4.0'), + 2, + ), + ( + (Fraction(1, 1), Fraction(5, 1), Fraction(3, 2)), + Fraction(5, 2), + 1, + ), + ( + ( + datetime(2019, 3, 29), + datetime(2019, 3, 30), + timedelta(hours=10), + ), + datetime(2019, 3, 29, 20), + 2, + ), + ]: + self.assertEqual(i, mi.numeric_range(*args).index(v)) + + for args, v in [ + ((0.7,), 0.5), + ((0.7,), 7.0), + ((0.7,), 10.0), + ((7.0, 0.0, -1.0), 0.5), + ((7.0, 0.0, -1.0), 0.0), + ((7.0, 0.0, -1.0), 10.0), + ((7.0, 0.0), 5.0), + ((Decimal("1.0"), Decimal("5.0"), Decimal("1.5")), Decimal('4.5')), + ((Fraction(1, 1), Fraction(5, 1), Fraction(3, 2)), Fraction(5, 3)), + ( + ( + datetime(2019, 3, 29), + datetime(2019, 3, 30), + timedelta(hours=10), + ), + datetime(2019, 3, 30), + ), + ]: + with self.assertRaises(ValueError): + mi.numeric_range(*args).index(v) + + def test_parent_classes(self): + r = mi.numeric_range(7.0) + self.assertTrue(isinstance(r, Iterable)) + self.assertFalse(isinstance(r, Iterator)) + self.assertTrue(isinstance(r, Sequence)) + self.assertTrue(isinstance(r, Hashable)) + + def test_bad_key(self): + r = mi.numeric_range(7.0) + for arg, message in [ + ('a', 'numeric range indices must be integers or slices, not str'), + ( + (), + 'numeric range indices must be integers or slices, not tuple', + ), + ]: + with self.assertRaisesRegex(TypeError, message): + r[arg] + + def test_pickle(self): + for args in [ + (7.0,), + (5.0, 7.0), + (5.0, 7.0, 3.0), + (7.0, 5.0), + (7.0, 5.0, 4.0), + (7.0, 5.0, -1.0), + (Decimal("1.0"), Decimal("5.0"), Decimal("1.5")), + (Fraction(1, 1), Fraction(5, 1), Fraction(3, 2)), + (datetime(2019, 3, 29), datetime(2019, 3, 30)), + ]: + r = mi.numeric_range(*args) + self.assertTrue(dumps(r)) # assert not empty + self.assertEqual(r, loads(dumps(r))) + + def test_empty_reversed(self): + self.assertEqual(list(reversed(mi.numeric_range(0))), []) + + +class CountCycleTests(TestCase): + def test_basic(self): + expected = [ + (0, 'a'), + (0, 'b'), + (0, 'c'), + (1, 'a'), + (1, 'b'), + (1, 'c'), + (2, 'a'), + (2, 'b'), + (2, 'c'), + ] + for actual in [ + mi.take(9, mi.count_cycle('abc')), # n=None + list(mi.count_cycle('abc', 3)), # n=3 + ]: + self.assertEqual(actual, expected) + + def test_empty(self): + self.assertEqual(list(mi.count_cycle('')), []) + self.assertEqual(list(mi.count_cycle('', 2)), []) + + def test_negative(self): + self.assertEqual(list(mi.count_cycle('abc', -3)), []) + + +class MarkEndsTests(TestCase): + def test_basic(self): + for size, expected in [ + (0, []), + (1, [(True, True, '0')]), + (2, [(True, False, '0'), (False, True, '1')]), + (3, [(True, False, '0'), (False, False, '1'), (False, True, '2')]), + ( + 4, + [ + (True, False, '0'), + (False, False, '1'), + (False, False, '2'), + (False, True, '3'), + ], + ), + ]: + with self.subTest(size=size): + iterable = map(str, range(size)) + actual = list(mi.mark_ends(iterable)) + self.assertEqual(actual, expected) + + +class LocateTests(TestCase): + def test_default_pred(self): + iterable = [0, 1, 1, 0, 1, 0, 0] + actual = list(mi.locate(iterable)) + expected = [1, 2, 4] + self.assertEqual(actual, expected) + + def test_no_matches(self): + iterable = [0, 0, 0] + actual = list(mi.locate(iterable)) + expected = [] + self.assertEqual(actual, expected) + + def test_custom_pred(self): + iterable = ['0', 1, 1, '0', 1, '0', '0'] + pred = lambda x: x == '0' + actual = list(mi.locate(iterable, pred)) + expected = [0, 3, 5, 6] + self.assertEqual(actual, expected) + + def test_window_size(self): + iterable = ['0', 1, 1, '0', 1, '0', '0'] + pred = lambda *args: args == ('0', 1) + actual = list(mi.locate(iterable, pred, window_size=2)) + expected = [0, 3] + self.assertEqual(actual, expected) + + def test_window_size_large(self): + iterable = [1, 2, 3, 4] + pred = lambda *args: True + actual = list(mi.locate(iterable, pred, window_size=5)) + expected = [0] + self.assertEqual(actual, expected) + + def test_window_size_zero(self): + iterable = [1, 2, 3, 4] + pred = lambda: True + with self.assertRaises(ValueError): + list(mi.locate(iterable, pred, window_size=0)) + + +class StripFunctionTests(TestCase): + def test_hashable(self): + iterable = list('www.example.com') + pred = lambda x: x in set('cmowz.') + + self.assertEqual(list(mi.lstrip(iterable, pred)), list('example.com')) + self.assertEqual(list(mi.rstrip(iterable, pred)), list('www.example')) + self.assertEqual(list(mi.strip(iterable, pred)), list('example')) + + def test_not_hashable(self): + iterable = [ + list('http://'), + list('www'), + list('.example'), + list('.com'), + ] + pred = lambda x: x in [list('http://'), list('www'), list('.com')] + + self.assertEqual(list(mi.lstrip(iterable, pred)), iterable[2:]) + self.assertEqual(list(mi.rstrip(iterable, pred)), iterable[:3]) + self.assertEqual(list(mi.strip(iterable, pred)), iterable[2:3]) + + def test_math(self): + iterable = [0, 1, 2, 3, 0, 1, 2, 3, 0, 1, 2] + pred = lambda x: x <= 2 + + self.assertEqual(list(mi.lstrip(iterable, pred)), iterable[3:]) + self.assertEqual(list(mi.rstrip(iterable, pred)), iterable[:-3]) + self.assertEqual(list(mi.strip(iterable, pred)), iterable[3:-3]) + + +class IteratorWithWeakReferences: + class _AnObj: + pass + + @classmethod + def FROM_SIZE(cls, size: int) -> IteratorWithWeakReferences: + return cls([IteratorWithWeakReferences._AnObj() for _ in range(size)]) + + def __init__(self, iterable: Iterable): + self._data = deque(element for element in iterable) + self._weakReferences = [weakref.ref(a) for a in self._data] + + def __iter__(self) -> Iterator: + return self + + def __next__(self) -> object: + if len(self._data) == 0: + raise StopIteration + return self._data.popleft() + + def weakReferencesState(self) -> list[bool]: + return [wr() is not None for wr in self._weakReferences] + + +class IsliceExtendedTests(TestCase): + def test_all(self): + iterable = ['0', '1', '2', '3', '4', '5'] + indexes = [*range(-4, 10), None] + steps = [1, 2, 3, 4, -1, -2, -3, -4] + for slice_args in product(indexes, indexes, steps): + with self.subTest(slice_args=slice_args): + actual = list(mi.islice_extended(iterable, *slice_args)) + expected = iterable[slice(*slice_args)] + self.assertEqual(actual, expected, slice_args) + + def test_zero_step(self): + with self.assertRaises(ValueError): + list(mi.islice_extended([1, 2, 3], 0, 1, 0)) + + def test_slicing(self): + iterable = map(str, count()) + first_slice = mi.islice_extended(iterable)[10:] + second_slice = mi.islice_extended(first_slice)[:10] + third_slice = mi.islice_extended(second_slice)[::2] + self.assertEqual(list(third_slice), ['10', '12', '14', '16', '18']) + + def test_slicing_extensive(self): + iterable = range(10) + options = (None, 1, 2, 7, -1) + for start, stop, step in product(options, options, options): + with self.subTest(slice_args=(start, stop, step)): + sliced_tuple_0 = tuple( + mi.islice_extended(iterable)[start:stop:step] + ) + sliced_tuple_1 = tuple( + mi.islice_extended(iterable, start, stop, step) + ) + sliced_range = tuple(iterable[start:stop:step]) + self.assertEqual(sliced_tuple_0, sliced_range) + self.assertEqual(sliced_tuple_1, sliced_range) + + def test_invalid_slice(self): + with self.assertRaises(TypeError): + mi.islice_extended(count())[13] + + def test_elements_lifecycle(self): + # CPython does reference counting. + # GC is not required when ref counting is supported. + refCountSupported = platform.python_implementation() == 'CPython' + + class TestCase(NamedTuple): + initialSize: int + slice: int + # list of expected intermediate elements states (alive or not) + # during a complete iteration + expectedAliveStates: list[list[int]] + + # fmt: off + testCases = [ + # testcases for: start>0, stop>0, step>0 + TestCase(initialSize=3, slice=(None, None, 1), expectedAliveStates=[ # noqa: E501 + [1, 1, 1], [0, 1, 1], [0, 0, 1], [0, 0, 0], [0, 0, 0]]), + TestCase(initialSize=3, slice=(0, None, 1), expectedAliveStates=[ + [1, 1, 1], [0, 1, 1], [0, 0, 1], [0, 0, 0], [0, 0, 0]]), + TestCase(initialSize=3, slice=(1, 2, 1), expectedAliveStates=[ + [1, 1, 1], [0, 0, 1], [0, 0, 1]]), + TestCase(initialSize=4, slice=(0, None, 2), expectedAliveStates=[ + [1, 1, 1, 1], [0, 1, 1, 1], [0, 0, 0, 1], [0, 0, 0, 0]]), + TestCase(initialSize=5, slice=(1, 4, 2), expectedAliveStates=[ + [1, 1, 1, 1, 1], [0, 0, 1, 1, 1], [0, 0, 0, 0, 1], [0, 0, 0, 0, 1]]), # noqa: E501 + TestCase(initialSize=5, slice=(4, 1, 1), expectedAliveStates=[ + [1, 1, 1, 1, 1], [0, 0, 0, 0, 1]]), + + # FYI: to process a negative start/stop index, we need to iterate + # on the whole iterator. All the elements will be consumed + # and will ALWAYS be released on full iteration completion. + + # testcases for: start<0, stop>0, step>0 + TestCase(initialSize=3, slice=(-3, None, 1), expectedAliveStates=[ + [1, 1, 1], [0, 1, 1], [0, 0, 1], [0, 0, 0], [0, 0, 0]]), + TestCase(initialSize=3, slice=(-2, 2, 1), expectedAliveStates=[ + [1, 1, 1], [0, 0, 1], [0, 0, 0]]), + TestCase(initialSize=4, slice=(-4, None, 2), expectedAliveStates=[ + [1, 1, 1, 1], [0, 1, 1, 1], [0, 0, 0, 1], [0, 0, 0, 0]]), + TestCase(initialSize=5, slice=(-4, 4, 2), expectedAliveStates=[ + [1, 1, 1, 1, 1], [0, 0, 1, 1, 1], [0, 0, 0, 0, 1], [0, 0, 0, 0, 0]]), # noqa: E501 + TestCase(initialSize=3, slice=(-2, 0, 1), expectedAliveStates=[ + [1, 1, 1], [0, 0, 0]]), + + # testcases for: start>0, stop<0, step>0 + TestCase(initialSize=3, slice=(None, -1, 1), expectedAliveStates=[ + [1, 1, 1], [0, 1, 1], [0, 0, 1], [0, 0, 0]]), + TestCase(initialSize=4, slice=(1, -1, 1), expectedAliveStates=[ + [1, 1, 1, 1], [0, 0, 1, 1], [0, 0, 0, 1], [0, 0, 0, 0]]), + TestCase(initialSize=5, slice=(None, -2, 2), expectedAliveStates=[ + [1, 1, 1, 1, 1], [0, 1, 1, 1, 1], [0, 0, 0, 1, 1], [0, 0, 0, 0, 0]]), # noqa: E501 + TestCase(initialSize=5, slice=(1, -1, 2), expectedAliveStates=[ + [1, 1, 1, 1, 1], [0, 0, 1, 1, 1], [0, 0, 0, 0, 1], [0, 0, 0, 0, 0]]), # noqa: E501 + TestCase(initialSize=5, slice=(4, -5, 2), expectedAliveStates=[ + [1, 1, 1, 1, 1], [0, 0, 0, 0, 0]]), + + # testcases for: start>0, stop>0, step<0 + TestCase(initialSize=3, slice=(None, None, -1), expectedAliveStates=[ # noqa: E501 + # ⚠️could be improved, elements are only released on final step + [1, 1, 1], [1, 1, 1], [1, 1, 1], [1, 1, 1], [0, 0, 0]]), + TestCase(initialSize=3, slice=(2, None, -1), expectedAliveStates=[ + # ⚠️could be improved, elements are only released on final step + [1, 1, 1], [1, 1, 1], [1, 1, 1], [1, 1, 1], [0, 0, 0]]), + TestCase(initialSize=3, slice=(None, 0, -1), expectedAliveStates=[ + # ⚠️could be improved, elements are only released on final step + [1, 1, 1], [0, 1, 1], [0, 1, 1], [0, 0, 0]]), + TestCase(initialSize=6, slice=(3, 1, -1), expectedAliveStates=[ + # ⚠️could be improved, elements are only released on final step + [1, 1, 1, 1, 1, 1], [0, 0, 1, 1, 1, 1], [0, 0, 1, 1, 1, 1], [0, 0, 0, 0, 1, 1]]), # noqa: E501 + TestCase(initialSize=5, slice=(1, 3, -1), expectedAliveStates=[ + # ⚠️could be improved. Final state could be [0, 0, 1, 1, 1] + [1, 1, 1, 1, 1], [0, 0, 0, 0, 1]]), + + # testcases for: start<0, stop>0, step<0 + TestCase(initialSize=3, slice=(-1, None, -1), expectedAliveStates=[ + # ⚠️could be improved, elements are only released on final step + [1, 1, 1], [1, 1, 1], [1, 1, 1], [1, 1, 1], [0, 0, 0]]), + TestCase(initialSize=3, slice=(-1, 0, -1), expectedAliveStates=[ + # ⚠️could be improved, elements are only released on final step + [1, 1, 1], [0, 1, 1], [0, 1, 1], [0, 0, 0]]), + TestCase(initialSize=6, slice=(-2, None, -2), expectedAliveStates=[ + # ⚠️could be improved, elements are only released on final step + [1, 1, 1, 1, 1, 1], [1, 1, 1, 1, 1, 1], [1, 1, 1, 1, 1, 1], [1, 1, 1, 1, 1, 1], [0, 0, 0, 0, 0, 0]]), # noqa: E501 + TestCase(initialSize=6, slice=(-2, 1, -2), expectedAliveStates=[ + # ⚠️could be improved, elements are only released on final step + [1, 1, 1, 1, 1, 1], [0, 0, 1, 1, 1, 1], [0, 0, 1, 1, 1, 1], [0, 0, 0, 0, 0, 0]]), # noqa: E501 + TestCase(initialSize=6, slice=(-4, 4, -2), expectedAliveStates=[ + [1, 1, 1, 1, 1, 1], [0, 0, 0, 0, 0, 0]]), + + # testcases for: start>0, stop<0, step<0 + TestCase(initialSize=3, slice=(None, -3, -1), expectedAliveStates=[ + # ⚠️could be improved, elements are only released on final step + [1, 1, 1], [0, 1, 1], [0, 1, 1], [0, 0, 0]]), + TestCase(initialSize=3, slice=(None, -4, -1), expectedAliveStates=[ + # ⚠️could be improved, elements are only released on final step + [1, 1, 1], [1, 1, 1], [1, 1, 1], [1, 1, 1], [0, 0, 0]]), + TestCase(initialSize=5, slice=(3, -4, -1), expectedAliveStates=[ + # ⚠️could be improved, elements are only released on final step + [1, 1, 1, 1, 1], [0, 0, 1, 1, 1], [0, 0, 1, 1, 1], [0, 0, 0, 0, 0]]), # noqa: E501 + TestCase(initialSize=5, slice=(1, -1, -1), expectedAliveStates=[ + [1, 1, 1, 1, 1], [0, 0, 0, 0, 0]]), + ] + # fmt: on + + for index, testCase in enumerate(testCases): + with self.subTest(f"{index:02d}", testCase=testCase): + iterator = IteratorWithWeakReferences.FROM_SIZE( + testCase.initialSize + ) + islice_iterator = mi.islice_extended(iterator, *testCase.slice) + + aliveStates = [] + refCountSupported or gc.collect() + # initial alive states + aliveStates.append(iterator.weakReferencesState()) + while True: + try: + next(islice_iterator) + refCountSupported or gc.collect() + # intermediate alive states + aliveStates.append(iterator.weakReferencesState()) + except StopIteration: + refCountSupported or gc.collect() + # final alive states + aliveStates.append(iterator.weakReferencesState()) + break + self.assertEqual(aliveStates, testCase.expectedAliveStates) + + +class ConsecutiveGroupsTest(TestCase): + def test_numbers(self): + iterable = [-10, -8, -7, -6, 1, 2, 4, 5, -1, 7] + actual = [list(g) for g in mi.consecutive_groups(iterable)] + expected = [[-10], [-8, -7, -6], [1, 2], [4, 5], [-1], [7]] + self.assertEqual(actual, expected) + + def test_custom_ordering(self): + iterable = ['1', '10', '11', '20', '21', '22', '30', '31'] + ordering = lambda x: int(x) + actual = [list(g) for g in mi.consecutive_groups(iterable, ordering)] + expected = [['1'], ['10', '11'], ['20', '21', '22'], ['30', '31']] + self.assertEqual(actual, expected) + + def test_exotic_ordering(self): + iterable = [ + ('a', 'b', 'c', 'd'), + ('a', 'c', 'b', 'd'), + ('a', 'c', 'd', 'b'), + ('a', 'd', 'b', 'c'), + ('d', 'b', 'c', 'a'), + ('d', 'c', 'a', 'b'), + ] + ordering = list(permutations('abcd')).index + actual = [list(g) for g in mi.consecutive_groups(iterable, ordering)] + expected = [ + [('a', 'b', 'c', 'd')], + [('a', 'c', 'b', 'd'), ('a', 'c', 'd', 'b'), ('a', 'd', 'b', 'c')], + [('d', 'b', 'c', 'a'), ('d', 'c', 'a', 'b')], + ] + self.assertEqual(actual, expected) + + +class DifferenceTest(TestCase): + def test_normal(self): + iterable = [10, 20, 30, 40, 50] + actual = list(mi.difference(iterable)) + expected = [10, 10, 10, 10, 10] + self.assertEqual(actual, expected) + + def test_custom(self): + iterable = [10, 20, 30, 40, 50] + actual = list(mi.difference(iterable, add)) + expected = [10, 30, 50, 70, 90] + self.assertEqual(actual, expected) + + def test_roundtrip(self): + original = list(range(100)) + accumulated = accumulate(original) + actual = list(mi.difference(accumulated)) + self.assertEqual(actual, original) + + def test_one(self): + self.assertEqual(list(mi.difference([0])), [0]) + + def test_empty(self): + self.assertEqual(list(mi.difference([])), []) + + def test_initial(self): + original = list(range(100)) + accumulated = accumulate(original, initial=100) + actual = list(mi.difference(accumulated, initial=100)) + self.assertEqual(actual, original) + + +class SeekableTest(PeekableMixinTests, TestCase): + cls = mi.seekable + + def test_exhaustion_reset(self): + iterable = [str(n) for n in range(10)] + + s = mi.seekable(iterable) + self.assertEqual(list(s), iterable) # Normal iteration + self.assertEqual(list(s), []) # Iterable is exhausted + + s.seek(0) + self.assertEqual(list(s), iterable) # Back in action + + def test_partial_reset(self): + iterable = [str(n) for n in range(10)] + + s = mi.seekable(iterable) + self.assertEqual(mi.take(5, s), iterable[:5]) # Normal iteration + + s.seek(1) + self.assertEqual(list(s), iterable[1:]) # Get the rest of the iterable + + def test_forward(self): + iterable = [str(n) for n in range(10)] + + s = mi.seekable(iterable) + self.assertEqual(mi.take(1, s), iterable[:1]) # Normal iteration + + s.seek(3) # Skip over index 2 + self.assertEqual(list(s), iterable[3:]) # Result is similar to slicing + + s.seek(0) # Back to 0 + self.assertEqual(list(s), iterable) # No difference in result + + def test_past_end(self): + iterable = [str(n) for n in range(10)] + + s = mi.seekable(iterable) + self.assertEqual(mi.take(1, s), iterable[:1]) # Normal iteration + + s.seek(20) + self.assertEqual(list(s), []) # Iterable is exhausted + + s.seek(0) # Back to 0 + self.assertEqual(list(s), iterable) # No difference in result + + def test_elements(self): + iterable = map(str, count()) + + s = mi.seekable(iterable) + mi.take(10, s) + + elements = s.elements() + self.assertEqual( + [elements[i] for i in range(10)], [str(n) for n in range(10)] + ) + self.assertEqual(len(elements), 10) + + mi.take(10, s) + self.assertEqual(list(elements), [str(n) for n in range(20)]) + + def test_maxlen(self): + iterable = map(str, count()) + + s = mi.seekable(iterable, maxlen=4) + self.assertEqual(mi.take(10, s), [str(n) for n in range(10)]) + self.assertEqual(list(s.elements()), ['6', '7', '8', '9']) + + s.seek(0) + self.assertEqual(mi.take(14, s), [str(n) for n in range(6, 20)]) + self.assertEqual(list(s.elements()), ['16', '17', '18', '19']) + + def test_maxlen_zero(self): + iterable = [str(x) for x in range(5)] + s = mi.seekable(iterable, maxlen=0) + self.assertEqual(list(s), iterable) + self.assertEqual(list(s.elements()), []) + + def test_maxlen_zero_peek_does_not_drop_items(self): + s = mi.seekable([10, 20, 30], maxlen=0) + self.assertEqual(s.peek(), 10) + self.assertEqual(s.peek(), 10) + self.assertEqual(list(s), [10, 20, 30]) + + def test_maxlen_zero_bool_does_not_drop_items(self): + s = mi.seekable([10, 20, 30], maxlen=0) + self.assertTrue(s) + self.assertEqual(list(s), [10, 20, 30]) + + s = mi.seekable([], maxlen=0) + self.assertFalse(s) + + @skipUnless( + platform.python_implementation() == 'CPython', + 'retained allocation check uses CPython tracemalloc', + ) + def test_repeated_zero_cache_peek_has_bounded_memory(self): + import tracemalloc + + if tracemalloc.is_tracing(): + self.skipTest('do not disturb an existing allocation trace') + iterator = mi.seekable(range(3), maxlen=0) + iterator.peek() + tracemalloc.start() + try: + for _ in range(10000): + iterator.peek() + bool(iterator) + retained = tracemalloc.get_traced_memory()[0] + finally: + tracemalloc.stop() + # A lookahead needs one item, not one iterator per peek/truth test. + self.assertLess(retained, 256 * 1024) + self.assertEqual(list(iterator.elements()), []) + self.assertEqual(list(iterator), [0, 1, 2]) + + def test_relative_seek(self): + iterable = [str(x) for x in range(5)] + s = mi.seekable(iterable) + s.relative_seek(2) + self.assertEqual(next(s), '2') + s.relative_seek(-2) + self.assertEqual(next(s), '1') + s.relative_seek(-2) + self.assertEqual( + next(s), '0' + ) # Seek relative to current position within the cache + s.relative_seek(-10) # Lower bound + self.assertEqual(next(s), '0') + s.relative_seek(10) # Lower bound + self.assertEqual(list(s.elements()), [str(x) for x in range(5)]) + + def test_getitem(self): + s = mi.seekable(str(n) for n in range(10)) + with self.assertRaises(IndexError): + s[0] + mi.take(3, s) + self.assertEqual(s[-1], '2') + self.assertEqual(s[0], '0') + self.assertEqual(s[2], '2') + with self.assertRaises(IndexError): + s[3] + + def test_getitem_maxlen(self): + s = mi.seekable((str(n) for n in range(10)), maxlen=2) + mi.take(5, s) + self.assertEqual(s[-1], '4') + self.assertEqual(s[0], '3') + + +class SequenceViewTests(TestCase): + def test_init(self): + view = mi.SequenceView((1, 2, 3)) + self.assertEqual(repr(view), "SequenceView((1, 2, 3))") + self.assertRaises(TypeError, lambda: mi.SequenceView({})) + + def test_update(self): + seq = [1, 2, 3] + view = mi.SequenceView(seq) + self.assertEqual(len(view), 3) + self.assertEqual(repr(view), "SequenceView([1, 2, 3])") + + seq.pop() + self.assertEqual(len(view), 2) + self.assertEqual(repr(view), "SequenceView([1, 2])") + + def test_indexing(self): + seq = ('a', 'b', 'c', 'd', 'e', 'f') + view = mi.SequenceView(seq) + for i in range(-len(seq), len(seq)): + self.assertEqual(view[i], seq[i]) + + def test_slicing(self): + seq = ('a', 'b', 'c', 'd', 'e', 'f') + view = mi.SequenceView(seq) + n = len(seq) + indexes = list(range(-n - 1, n + 1)) + [None] + steps = list(range(-n, n + 1)) + steps.remove(0) + for slice_args in product(indexes, indexes, steps): + i = slice(*slice_args) + self.assertEqual(view[i], seq[i]) + + def test_abc_methods(self): + # collections.Sequence should provide all of this functionality + seq = ('a', 'b', 'c', 'd', 'e', 'f', 'f') + view = mi.SequenceView(seq) + + # __contains__ + self.assertIn('b', view) + self.assertNotIn('g', view) + + # __iter__ + self.assertEqual(list(iter(view)), list(seq)) + + # __reversed__ + self.assertEqual(list(reversed(view)), list(reversed(seq))) + + # index + self.assertEqual(view.index('b'), 1) + + # count + self.assertEqual(seq.count('f'), 2) + + +class RunLengthTest(TestCase): + def test_encode(self): + iterable = (int(str(n)[0]) for n in count(800)) + actual = mi.take(4, mi.run_length.encode(iterable)) + expected = [(8, 100), (9, 100), (1, 1000), (2, 1000)] + self.assertEqual(actual, expected) + + def test_decode(self): + iterable = [('d', 4), ('c', 3), ('b', 2), ('a', 1)] + actual = ''.join(mi.run_length.decode(iterable)) + expected = 'ddddcccbba' + self.assertEqual(actual, expected) + + +class ExactlyNTests(TestCase): + """Tests for ``exactly_n()``""" + + def test_true(self): + """Iterable has ``n`` ``True`` elements""" + self.assertTrue(mi.exactly_n([True, False, True], 2)) + self.assertTrue(mi.exactly_n([1, 1, 1, 0], 3)) + self.assertTrue(mi.exactly_n([False, False], 0)) + self.assertTrue(mi.exactly_n(range(100), 10, lambda x: x < 10)) + self.assertTrue(mi.exactly_n(repeat(True, 100), 100)) + self.assertTrue(mi.exactly_n(repeat(False, 100), 100, predicate=not_)) + + def test_false(self): + """Iterable does not have ``n`` ``True`` elements""" + self.assertFalse(mi.exactly_n([True, False, False], 2)) + self.assertFalse(mi.exactly_n([True, True, False], 1)) + self.assertFalse(mi.exactly_n([False], 1)) + self.assertFalse(mi.exactly_n([True], -1)) + self.assertFalse(mi.exactly_n([True], -10)) + self.assertFalse(mi.exactly_n([], -1)) + self.assertFalse(mi.exactly_n([], -10)) + self.assertFalse(mi.exactly_n([True], 0)) + self.assertFalse(mi.exactly_n(repeat(True), 100)) + + def test_empty(self): + """Return ``True`` if the iterable is empty and ``n`` is 0""" + self.assertTrue(mi.exactly_n([], 0)) + self.assertFalse(mi.exactly_n([], 1)) + + +class AlwaysReversibleTests(TestCase): + """Tests for ``always_reversible()``""" + + def test_regular_reversed(self): + self.assertEqual( + list(reversed(range(10))), list(mi.always_reversible(range(10))) + ) + self.assertEqual( + list(reversed([1, 2, 3])), list(mi.always_reversible([1, 2, 3])) + ) + self.assertEqual( + reversed([1, 2, 3]).__class__, + mi.always_reversible([1, 2, 3]).__class__, + ) + + def test_nonseq_reversed(self): + # Create a non-reversible generator from a sequence + with self.assertRaises(TypeError): + reversed(x for x in range(10)) + + self.assertEqual( + list(reversed(range(10))), + list(mi.always_reversible(x for x in range(10))), + ) + self.assertEqual( + list(reversed([1, 2, 3])), + list(mi.always_reversible(x for x in [1, 2, 3])), + ) + self.assertNotEqual( + reversed((1, 2)).__class__, + mi.always_reversible(x for x in (1, 2)).__class__, + ) + + +class CircularShiftsTests(TestCase): + def test_empty(self): + # empty iterable -> empty list + self.assertEqual(list(mi.circular_shifts([])), []) + + def test_simple_circular_shifts(self): + # test the a simple iterator case + self.assertEqual( + list(mi.circular_shifts(range(4))), + [(0, 1, 2, 3), (1, 2, 3, 0), (2, 3, 0, 1), (3, 0, 1, 2)], + ) + + def test_duplicates(self): + # test non-distinct entries + self.assertEqual( + list(mi.circular_shifts([0, 1, 0, 1])), + [(0, 1, 0, 1), (1, 0, 1, 0), (0, 1, 0, 1), (1, 0, 1, 0)], + ) + + def test_steps_positive(self): + actual = list(mi.circular_shifts(range(5), steps=2)) + expected = [ + (0, 1, 2, 3, 4), + (2, 3, 4, 0, 1), + (4, 0, 1, 2, 3), + (1, 2, 3, 4, 0), + (3, 4, 0, 1, 2), + ] + self.assertEqual(actual, expected) + + def test_steps_negative(self): + actual = list(mi.circular_shifts(range(5), steps=-2)) + expected = [ + (0, 1, 2, 3, 4), + (3, 4, 0, 1, 2), + (1, 2, 3, 4, 0), + (4, 0, 1, 2, 3), + (2, 3, 4, 0, 1), + ] + self.assertEqual(actual, expected) + + def test_steps_zero(self): + with self.assertRaises(ValueError): + list(mi.circular_shifts(range(5), steps=0)) + + +class MakeDecoratorTests(TestCase): + def test_basic(self): + slicer = mi.make_decorator(islice) + + @slicer(1, 10, 2) + def user_function(arg_1, arg_2, kwarg_1=None): + self.assertEqual(arg_1, 'arg_1') + self.assertEqual(arg_2, 'arg_2') + self.assertEqual(kwarg_1, 'kwarg_1') + return map(str, count()) + + it = user_function('arg_1', 'arg_2', kwarg_1='kwarg_1') + actual = list(it) + expected = ['1', '3', '5', '7', '9'] + self.assertEqual(actual, expected) + + def test_result_index(self): + def stringify(*args, **kwargs): + self.assertEqual(args[0], 'arg_0') + iterable = args[1] + self.assertEqual(args[2], 'arg_2') + self.assertEqual(kwargs['kwarg_1'], 'kwarg_1') + return map(str, iterable) + + stringifier = mi.make_decorator(stringify, result_index=1) + + @stringifier('arg_0', 'arg_2', kwarg_1='kwarg_1') + def user_function(n): + return count(n) + + it = user_function(1) + actual = mi.take(5, it) + expected = ['1', '2', '3', '4', '5'] + self.assertEqual(actual, expected) + + def test_wrap_class(self): + seeker = mi.make_decorator(mi.seekable) + + @seeker() + def user_function(n): + return map(str, range(n)) + + it = user_function(5) + self.assertEqual(list(it), ['0', '1', '2', '3', '4']) + + it.seek(0) + self.assertEqual(list(it), ['0', '1', '2', '3', '4']) + + +class MapReduceTests(TestCase): + def test_default(self): + iterable = (str(x) for x in range(5)) + keyfunc = lambda x: int(x) // 2 + actual = sorted(mi.map_reduce(iterable, keyfunc).items()) + expected = [(0, ['0', '1']), (1, ['2', '3']), (2, ['4'])] + self.assertEqual(actual, expected) + + def test_valuefunc(self): + iterable = (str(x) for x in range(5)) + keyfunc = lambda x: int(x) // 2 + valuefunc = int + actual = sorted(mi.map_reduce(iterable, keyfunc, valuefunc).items()) + expected = [(0, [0, 1]), (1, [2, 3]), (2, [4])] + self.assertEqual(actual, expected) + + def test_reducefunc(self): + iterable = (str(x) for x in range(5)) + keyfunc = lambda x: int(x) // 2 + valuefunc = int + reducefunc = lambda value_list: reduce(mul, value_list, 1) + actual = sorted( + mi.map_reduce(iterable, keyfunc, valuefunc, reducefunc).items() + ) + expected = [(0, 0), (1, 6), (2, 4)] + self.assertEqual(actual, expected) + + def test_ret(self): + d = mi.map_reduce([1, 0, 2, 0, 1, 0], bool) + self.assertEqual(d, {False: [0, 0, 0], True: [1, 2, 1]}) + self.assertRaises(KeyError, lambda: d[None].append(1)) + + +class RlocateTests(TestCase): + def test_default_pred(self): + iterable = [0, 1, 1, 0, 1, 0, 0] + for it in (iterable[:], iter(iterable)): + actual = list(mi.rlocate(it)) + expected = [4, 2, 1] + self.assertEqual(actual, expected) + + def test_no_matches(self): + iterable = [0, 0, 0] + for it in (iterable[:], iter(iterable)): + actual = list(mi.rlocate(it)) + expected = [] + self.assertEqual(actual, expected) + + def test_custom_pred(self): + iterable = ['0', 1, 1, '0', 1, '0', '0'] + pred = lambda x: x == '0' + for it in (iterable[:], iter(iterable)): + actual = list(mi.rlocate(it, pred)) + expected = [6, 5, 3, 0] + self.assertEqual(actual, expected) + + def test_efficient_reversal(self): + iterable = range(9**9) # Is efficiently reversible + target = 9**9 - 2 + pred = lambda x: x == target # Find-able from the right + actual = next(mi.rlocate(iterable, pred)) + self.assertEqual(actual, target) + + def test_window_size(self): + iterable = ['0', 1, 1, '0', 1, '0', '0'] + pred = lambda *args: args == ('0', 1) + for it in (iterable, iter(iterable)): + actual = list(mi.rlocate(it, pred, window_size=2)) + expected = [3, 0] + self.assertEqual(actual, expected) + + def test_window_size_large(self): + iterable = [1, 2, 3, 4] + pred = lambda *args: True + for it in (iterable, iter(iterable)): + actual = list(mi.rlocate(iterable, pred, window_size=5)) + expected = [0] + self.assertEqual(actual, expected) + + def test_window_size_zero(self): + iterable = [1, 2, 3, 4] + pred = lambda: True + for it in (iterable, iter(iterable)): + with self.assertRaises(ValueError): + list(mi.locate(iterable, pred, window_size=0)) + + +class ReplaceTests(TestCase): + def test_basic(self): + iterable = range(10) + pred = lambda x: x % 2 == 0 + substitutes = [] + actual = list(mi.replace(iterable, pred, substitutes)) + expected = [1, 3, 5, 7, 9] + self.assertEqual(actual, expected) + + def test_count(self): + iterable = range(10) + pred = lambda x: x % 2 == 0 + substitutes = [] + actual = list(mi.replace(iterable, pred, substitutes, count=4)) + expected = [1, 3, 5, 7, 8, 9] + self.assertEqual(actual, expected) + + def test_window_size(self): + iterable = range(10) + pred = lambda *args: args == (0, 1, 2) + substitutes = [] + actual = list(mi.replace(iterable, pred, substitutes, window_size=3)) + expected = [3, 4, 5, 6, 7, 8, 9] + self.assertEqual(actual, expected) + + def test_window_size_end(self): + iterable = range(10) + pred = lambda *args: args == (7, 8, 9) + substitutes = [] + actual = list(mi.replace(iterable, pred, substitutes, window_size=3)) + expected = [0, 1, 2, 3, 4, 5, 6] + self.assertEqual(actual, expected) + + def test_window_size_count(self): + iterable = range(10) + pred = lambda *args: (args == (0, 1, 2)) or (args == (7, 8, 9)) + substitutes = [] + actual = list( + mi.replace(iterable, pred, substitutes, count=1, window_size=3) + ) + expected = [3, 4, 5, 6, 7, 8, 9] + self.assertEqual(actual, expected) + + def test_window_size_large(self): + iterable = range(4) + pred = lambda *args: True + substitutes = [5, 6, 7] + actual = list(mi.replace(iterable, pred, substitutes, window_size=5)) + expected = [5, 6, 7] + self.assertEqual(actual, expected) + + def test_window_size_zero(self): + iterable = range(10) + pred = lambda *args: True + substitutes = [] + with self.assertRaises(ValueError): + list(mi.replace(iterable, pred, substitutes, window_size=0)) + + def test_iterable_substitutes(self): + iterable = range(5) + pred = lambda x: x % 2 == 0 + substitutes = iter('__') + actual = list(mi.replace(iterable, pred, substitutes)) + expected = ['_', '_', 1, '_', '_', 3, '_', '_'] + self.assertEqual(actual, expected) + + +class PartitionsTest(TestCase): + def test_types(self): + for iterable in ['abcd', ['a', 'b', 'c', 'd'], ('a', 'b', 'c', 'd')]: + with self.subTest(iterable=iterable): + actual = list(mi.partitions(iterable)) + expected = [ + [['a', 'b', 'c', 'd']], + [['a'], ['b', 'c', 'd']], + [['a', 'b'], ['c', 'd']], + [['a', 'b', 'c'], ['d']], + [['a'], ['b'], ['c', 'd']], + [['a'], ['b', 'c'], ['d']], + [['a', 'b'], ['c'], ['d']], + [['a'], ['b'], ['c'], ['d']], + ] + self.assertEqual(actual, expected) + + def test_empty(self): + iterable = [] + actual = list(mi.partitions(iterable)) + expected = [[[]]] + self.assertEqual(actual, expected) + + def test_order(self): + iterable = iter([3, 2, 1]) + actual = list(mi.partitions(iterable)) + expected = [[[3, 2, 1]], [[3], [2, 1]], [[3, 2], [1]], [[3], [2], [1]]] + self.assertEqual(actual, expected) + + def test_duplicates(self): + iterable = [1, 1, 1] + actual = list(mi.partitions(iterable)) + expected = [[[1, 1, 1]], [[1], [1, 1]], [[1, 1], [1]], [[1], [1], [1]]] + self.assertEqual(actual, expected) + + +class _FrozenMultiset(Set): + """ + A helper class, useful to compare two lists without reference to the order + of elements. + + FrozenMultiset represents a hashable set that allows duplicate elements. + """ + + def __init__(self, iterable): + self._collection = frozenset(Counter(iterable).items()) + + def __contains__(self, y): + """ + >>> (0, 1) in _FrozenMultiset([(0, 1), (2,), (0, 1)]) + True + """ + return any(y == x for x, _ in self._collection) + + def __iter__(self): + """ + >>> sorted(_FrozenMultiset([(0, 1), (2,), (0, 1)])) + [(0, 1), (0, 1), (2,)] + """ + return (x for x, c in self._collection for _ in range(c)) + + def __len__(self): + """ + >>> len(_FrozenMultiset([(0, 1), (2,), (0, 1)])) + 3 + """ + return sum(c for x, c in self._collection) + + def has_duplicates(self): + """ + >>> _FrozenMultiset([(0, 1), (2,), (0, 1)]).has_duplicates() + True + """ + return any(c != 1 for _, c in self._collection) + + def __hash__(self): + return hash(self._collection) + + def __repr__(self): + return f'FrozenSet([{", ".join(repr(x) for x in iter(self))}]' + + +class SetPartitionsTests(TestCase): + @staticmethod + def _normalize_partition(p): + """ + Return a normalized, hashable, version of a partition using + _FrozenMultiset + """ + return _FrozenMultiset(_FrozenMultiset(g) for g in p) + + @staticmethod + def _normalize_partitions(ps): + """ + Return a normalized set of all normalized partitions using + _FrozenMultiset + """ + return _FrozenMultiset( + SetPartitionsTests._normalize_partition(p) for p in ps + ) + + def test_repeated(self): + it = 'aaa' + actual = mi.set_partitions(it, 2) + expected = [['a', 'aa'], ['a', 'aa'], ['a', 'aa']] + self.assertEqual( + self._normalize_partitions(expected), + self._normalize_partitions(actual), + ) + + def test_each_correct(self): + a = set(range(6)) + for p in mi.set_partitions(a): + total = {e for g in p for e in g} + self.assertEqual(a, total) + + def test_duplicates(self): + a = set(range(6)) + for p in mi.set_partitions(a): + self.assertFalse(self._normalize_partition(p).has_duplicates()) + + def test_found_all(self): + """small example, hand-checked""" + expected = [ + [[0], [1], [2, 3, 4]], + [[0], [1, 2], [3, 4]], + [[0], [2], [1, 3, 4]], + [[0], [3], [1, 2, 4]], + [[0], [4], [1, 2, 3]], + [[0], [1, 3], [2, 4]], + [[0], [1, 4], [2, 3]], + [[1], [2], [0, 3, 4]], + [[1], [3], [0, 2, 4]], + [[1], [4], [0, 2, 3]], + [[1], [0, 2], [3, 4]], + [[1], [0, 3], [2, 4]], + [[1], [0, 4], [2, 3]], + [[2], [3], [0, 1, 4]], + [[2], [4], [0, 1, 3]], + [[2], [0, 1], [3, 4]], + [[2], [0, 3], [1, 4]], + [[2], [0, 4], [1, 3]], + [[3], [4], [0, 1, 2]], + [[3], [0, 1], [2, 4]], + [[3], [0, 2], [1, 4]], + [[3], [0, 4], [1, 2]], + [[4], [0, 1], [2, 3]], + [[4], [0, 2], [1, 3]], + [[4], [0, 3], [1, 2]], + ] + actual = mi.set_partitions(range(5), 3) + self.assertEqual( + self._normalize_partitions(expected), + self._normalize_partitions(actual), + ) + + def test_stirling_numbers(self): + """Check against https://en.wikipedia.org/wiki/ + Stirling_numbers_of_the_second_kind#Table_of_values""" + cardinality_by_k_by_n = [ + [1], + [1, 1], + [1, 3, 1], + [1, 7, 6, 1], + [1, 15, 25, 10, 1], + [1, 31, 90, 65, 15, 1], + ] + for n, cardinality_by_k in enumerate(cardinality_by_k_by_n, 1): + for k, cardinality in enumerate(cardinality_by_k, 1): + self.assertEqual( + cardinality, len(list(mi.set_partitions(range(n), k))) + ) + + def test_no_group(self): + def helper(): + list(mi.set_partitions(range(4), -1)) + + self.assertRaises(ValueError, helper) + + def test_to_many_groups(self): + self.assertEqual([], list(mi.set_partitions(range(4), 5))) + + def test_min_size(self): + it = 'abc' + actual = mi.set_partitions(it, min_size=2) + expected = [['abc']] + self.assertEqual( + self._normalize_partitions(expected), + self._normalize_partitions(actual), + ) + + def test_max_size(self): + it = 'abc' + actual = mi.set_partitions(it, max_size=2) + expected = [['a', 'bc'], ['ab', 'c'], ['b', 'ac'], ['a', 'b', 'c']] + self.assertEqual( + self._normalize_partitions(expected), + self._normalize_partitions(actual), + ) + + def test_min_max(self): + it = 'abcdefg' + self.assertEqual( + list(mi.set_partitions(it, min_size=4, max_size=3)), [] + ) + + +class TimeLimitedTests(TestCase): + def test_basic(self): + def generator(): + yield 1 + yield 2 + sleep(0.2) + yield 3 + + iterable = mi.time_limited(0.1, generator()) + actual = list(iterable) + expected = [1, 2] + self.assertEqual(actual, expected) + self.assertTrue(iterable.timed_out) + + def test_complete(self): + iterable = mi.time_limited(2, iter(range(10))) + actual = list(iterable) + expected = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9] + self.assertEqual(actual, expected) + self.assertFalse(iterable.timed_out) + + def test_zero_limit(self): + iterable = mi.time_limited(0, count()) + actual = list(iterable) + expected = [] + self.assertEqual(actual, expected) + self.assertTrue(iterable.timed_out) + + def test_invalid_limit(self): + with self.assertRaises(ValueError): + list(mi.time_limited(-0.1, count())) + + +class OnlyTests(TestCase): + def test_defaults(self): + self.assertEqual(mi.only([]), None) + self.assertEqual(mi.only([1]), 1) + self.assertRaises(ValueError, lambda: mi.only([1, 2])) + + def test_custom_value(self): + self.assertEqual(mi.only([], default='!'), '!') + self.assertEqual(mi.only([1], default='!'), 1) + self.assertRaises(ValueError, lambda: mi.only([1, 2], default='!')) + + def test_custom_exception(self): + self.assertEqual(mi.only([], too_long=RuntimeError), None) + self.assertEqual(mi.only([1], too_long=RuntimeError), 1) + self.assertRaises( + RuntimeError, lambda: mi.only([1, 2], too_long=RuntimeError) + ) + + def test_default_exception_message(self): + self.assertRaisesRegex( + ValueError, + "Expected exactly one item in iterable, " + "but got 'foo', 'bar', and perhaps more", + lambda: mi.only(['foo', 'bar', 'baz']), + ) + + def test_falsy_custom_exception(self): + class FalsyError(Exception): + def __bool__(self): + return False + + too_long = FalsyError('too many') + self.assertRaises( + FalsyError, lambda: mi.only([1, 2], too_long=too_long) + ) + + def test_too_long_does_not_evaluate_repr(self): + class NoRepr: + def __repr__(self): + raise RuntimeError('repr should not be called') + + self.assertRaises( + OverflowError, + lambda: mi.only( + [NoRepr(), NoRepr(), NoRepr()], too_long=OverflowError + ), + ) + + +class IchunkedTests(TestCase): + def test_even(self): + iterable = (str(x) for x in range(10)) + actual = [''.join(c) for c in mi.ichunked(iterable, 5)] + expected = ['01234', '56789'] + self.assertEqual(actual, expected) + + def test_odd(self): + iterable = (str(x) for x in range(10)) + actual = [''.join(c) for c in mi.ichunked(iterable, 4)] + expected = ['0123', '4567', '89'] + self.assertEqual(actual, expected) + + def test_zero(self): + iterable = [] + actual = [list(c) for c in mi.ichunked(iterable, 0)] + expected = [] + self.assertEqual(actual, expected) + + def test_zero_nonempty(self): + # chunked(..., 0) yields nothing; ichunked should match, including + # when the source still has items (the empty-iterable case hid this). + iterable = [1, 2, 3] + actual = [list(c) for c in mi.ichunked(iterable, 0)] + self.assertEqual(actual, []) + + source = iter([1, 2, 3]) + self.assertEqual([list(c) for c in mi.ichunked(source, 0)], []) + self.assertEqual(list(source), [1, 2, 3]) + + def test_negative(self): + iterable = count() + with self.assertRaises(ValueError) as cm: + [list(c) for c in mi.ichunked(iterable, -1)] + self.assertIn('n must be at least 0', str(cm.exception)) + + def test_out_of_order(self): + iterable = map(str, count()) + it = mi.ichunked(iterable, 4) + chunk_1 = next(it) + chunk_2 = next(it) + self.assertEqual(''.join(chunk_2), '4567') + self.assertEqual(''.join(chunk_1), '0123') + + def test_laziness(self): + def gen(): + yield 0 + raise RuntimeError + yield from count(1) + + it = mi.ichunked(gen(), 4) + chunk = next(it) + self.assertEqual(next(chunk), 0) + self.assertRaises(RuntimeError, next, it) + + def test_memory_in_order(self): + gen_numbers = [] + + def gen(): + for gen_number in count(): + gen_numbers.append(gen_number) + yield gen_number + + # No items should be kept in memory when a ichunked is first called + all_chunks = mi.ichunked(gen(), 4) + self.assertEqual(gen_numbers, []) + + # The first item of each chunk should be generated on chunk generation + first_chunk = next(all_chunks) + self.assertEqual(gen_numbers, [0]) + + # If we don't read a chunk before getting its successor, its contents + # will be cached + second_chunk = next(all_chunks) + self.assertEqual(gen_numbers, [0, 1, 2, 3, 4]) + + # Check if we can read in cached values + self.assertEqual(list(first_chunk), [0, 1, 2, 3]) + self.assertEqual(list(second_chunk), [4, 5, 6, 7]) + + # Again only the most recent chunk should have an item cached + third_chunk = next(all_chunks) + self.assertEqual(len(gen_numbers), 9) + + # No new item should be cached when reading past the first number + next(third_chunk) + self.assertEqual(len(gen_numbers), 9) + + # we should not be able to read spent chunks + self.assertEqual(list(first_chunk), []) + self.assertEqual(list(second_chunk), []) + + +class DistinctCombinationsTests(TestCase): + def test_basic(self): + for iterable in [ + (1, 2, 2, 3, 3, 3), # In order + range(6), # All distinct + 'abbccc', # Not numbers + 'cccbba', # Backward + 'mississippi', # No particular order + ]: + for r in range(len(iterable)): + with self.subTest(iterable=iterable, r=r): + actual = list(mi.distinct_combinations(iterable, r)) + expected = list( + mi.unique_everseen(combinations(iterable, r)) + ) + self.assertEqual(actual, expected) + + def test_negative(self): + with self.assertRaises(ValueError): + list(mi.distinct_combinations([], -1)) + + def test_empty(self): + self.assertEqual(list(mi.distinct_combinations([], 2)), []) + + +class FilterExceptTests(TestCase): + def test_no_exceptions_pass(self): + iterable = '0123' + actual = list(mi.filter_except(int, iterable)) + expected = ['0', '1', '2', '3'] + self.assertEqual(actual, expected) + + def test_no_exceptions_raise(self): + iterable = ['0', '1', 'two', '3'] + with self.assertRaises(ValueError): + list(mi.filter_except(int, iterable)) + + def test_raise(self): + iterable = ['0', '12', 'three', None] + with self.assertRaises(TypeError): + list(mi.filter_except(int, iterable, ValueError)) + + def test_false(self): + # Even if the validator returns false, we pass through + validator = lambda x: False + iterable = ['0', '1', '2', 'three', None] + actual = list(mi.filter_except(validator, iterable, Exception)) + expected = ['0', '1', '2', 'three', None] + self.assertEqual(actual, expected) + + def test_multiple(self): + iterable = ['0', '1', '2', 'three', None, '4'] + actual = list(mi.filter_except(int, iterable, ValueError, TypeError)) + expected = ['0', '1', '2', '4'] + self.assertEqual(actual, expected) + + +class MapExceptTests(TestCase): + def test_no_exceptions_pass(self): + iterable = '0123' + actual = list(mi.map_except(int, iterable)) + expected = [0, 1, 2, 3] + self.assertEqual(actual, expected) + + def test_no_exceptions_raise(self): + iterable = ['0', '1', 'two', '3'] + with self.assertRaises(ValueError): + list(mi.map_except(int, iterable)) + + def test_raise(self): + iterable = ['0', '12', 'three', None] + with self.assertRaises(TypeError): + list(mi.map_except(int, iterable, ValueError)) + + def test_multiple(self): + iterable = ['0', '1', '2', 'three', None, '4'] + actual = list(mi.map_except(int, iterable, ValueError, TypeError)) + expected = [0, 1, 2, 4] + self.assertEqual(actual, expected) + + +class MapIfTests(TestCase): + def test_without_func_else(self): + iterable = list(range(-5, 5)) + actual = list(mi.map_if(iterable, lambda x: x > 3, lambda x: 'toobig')) + expected = [-5, -4, -3, -2, -1, 0, 1, 2, 3, 'toobig'] + self.assertEqual(actual, expected) + + def test_with_func_else(self): + iterable = list(range(-5, 5)) + actual = list( + mi.map_if( + iterable, lambda x: x >= 0, lambda x: 'notneg', lambda x: 'neg' + ) + ) + expected = ['neg'] * 5 + ['notneg'] * 5 + self.assertEqual(actual, expected) + + def test_empty(self): + actual = list(mi.map_if([], lambda x: len(x) > 5, lambda x: None)) + expected = [] + self.assertEqual(actual, expected) + + +class SampleTests(TestCase): + def test_specific_sample(self): + """Verify reproducibility.""" + + # Note, this test is surprisingly robust. Although it depends on the quality + # of the underlying libmath implementations for log, exp, and log1p, the + # number of samples and population size are small enough that small errors + # in those underlying functions won't affect the sample. + + seed(8675309) + self.assertEqual( + list(mi.sample(range(10**5), k=5)), + [16845, 79805, 76057, 58302, 40472], + ) + + seed(8675309) + self.assertEqual( + list(mi.sample(range(10**5), counts=[1, 2] * (10**5 // 2), k=5)), + [87899, 53203, 38868, 11230, 50705], + ) + + seed(8675309) + self.assertEqual( + list(mi.sample(range(10**5), weights=range(1, 10**5 + 1), k=5)), + [50915, 33816, 32250, 98284, 43517], + ) + + def test_unit_case(self): + """Test against a fixed case by seeding the random module.""" + # Beware that this test really just verifies random.random() behavior. + # If the algorithm is changed (e.g. to a more naive implementation) + # this test will fail, but the algorithm might be correct. + # Also, this test can pass and the algorithm can be completely wrong. + data = "abcdef" + weights = list(range(1, len(data) + 1)) + seed(123) + actual = mi.sample(data, k=2, weights=weights) + expected = ['f', 'e'] + self.assertEqual(actual, expected) + + def test_negative(self): + data = [1, 2, 3, 4, 5] + with self.assertRaises(ValueError): + mi.sample(data, k=-1) + + def test_length(self): + """Check that *k* elements are sampled.""" + data = [1, 2, 3, 4, 5] + for k in [0, 3, 5, 7]: + sampled = mi.sample(data, k=k) + actual = len(sampled) + expected = min(k, len(data)) + self.assertEqual(actual, expected) + + def test_strict(self): + data = ['1', '2', '3', '4', '5'] + self.assertEqual(set(mi.sample(data, 6, strict=False)), set(data)) + with self.assertRaises(ValueError): + mi.sample(data, 6, strict=True) + + def test_counts(self): + # Test with counts + seed(0) + iterable = ['red', 'blue'] + counts = [4, 2] + k = 5 + actual = list(mi.sample(iterable, counts=counts, k=k)) + + # Test without counts + seed(0) + decoded_iterable = (['red'] * 4) + (['blue'] * 2) + expected = list(mi.sample(decoded_iterable, k=k)) + + self.assertEqual(actual, expected) + + def test_counts_all(self): + actual = Counter(mi.sample('uwxyz', 35, counts=(1, 0, 4, 10, 20))) + expected = Counter({'u': 1, 'x': 4, 'y': 10, 'z': 20}) + self.assertEqual(actual, expected) + + def test_sampling_entire_iterable(self): + """If k=len(iterable), the sample contains the original elements.""" + data = ["a", 2, "a", 4, (1, 2, 3)] + actual = set(mi.sample(data, k=len(data))) + expected = set(data) + self.assertEqual(actual, expected) + + def test_scale_invariance_of_weights(self): + """The probability of choosing element a_i is w_i / sum(weights). + Scaling weights should not change the probability or outcome.""" + data = "abcdef" + + weights = list(range(1, len(data) + 1)) + seed(123) + first_sample = mi.sample(data, k=2, weights=weights) + + # Scale the weights and sample again + weights_scaled = [w / 1e10 for w in weights] + seed(123) + second_sample = mi.sample(data, k=2, weights=weights_scaled) + + self.assertEqual(first_sample, second_sample) + + def test_invariance_under_permutations_unweighted(self): + """The order of the data should not matter. This is a stochastic test, + but it will fail in less than 1 / 10_000 cases.""" + + # Create a data set and a reversed data set + data = list(range(100)) + data_rev = list(reversed(data)) + + # Sample each data set 10 times + data_means = [mean(mi.sample(data, k=50)) for _ in range(10)] + data_rev_means = [mean(mi.sample(data_rev, k=50)) for _ in range(10)] + + # The difference in the means should be low, i.e. little bias + difference_in_means = abs(mean(data_means) - mean(data_rev_means)) + + # The observed largest difference in 10,000 simulations was 5.09599 + self.assertTrue(difference_in_means < 5.1) + + def test_invariance_under_permutations_weighted(self): + """The order of the data should not matter. This is a stochastic test, + but it will fail in less than 1 / 10_000 cases.""" + + # Create a data set and a reversed data set + data = list(range(1, 101)) + data_rev = list(reversed(data)) + + # Sample each data set 10 times + data_means = [ + mean(mi.sample(data, k=50, weights=data)) for _ in range(10) + ] + data_rev_means = [ + mean(mi.sample(data_rev, k=50, weights=data_rev)) + for _ in range(10) + ] + + # The difference in the means should be low, i.e. little bias + difference_in_means = abs(mean(data_means) - mean(data_rev_means)) + + # The observed largest difference in 10,000 simulations was 4.337999 + self.assertTrue(difference_in_means < 4.4) + + def test_error_cases(self): + # weights and counts are mutually exclusive + with self.assertRaises(TypeError): + mi.sample( + 'abcde', 3, weights=[1, 2, 3, 4, 5], counts=[1, 2, 3, 4, 5] + ) + + # Weighted sample larger than population + with self.assertRaises(ValueError): + mi.sample('abcde', 10, weights=[1, 2, 3, 4, 5], strict=True) + + # Counted sample larger than population + with self.assertRaises(ValueError): + mi.sample('abcde', 10, counts=[1, 1, 1, 1, 1], strict=True) + + +class BarelySortable: + def __init__(self, value): + self.value = value + + def __lt__(self, other): + return self.value < other.value + + def __int__(self): + return int(self.value) + + +class IsSortedTests(TestCase): + def test_basic(self): + for iterable, kwargs, expected in [ + ([], {}, True), + ([1], {}, True), + ([1, 2, 3], {}, True), + ([1, 1, 2, 3], {}, True), + ([1, 10, 2, 3], {}, False), + (['1', '10', '2', '3'], {}, True), + (['1', '10', '2', '3'], {'key': int}, False), + ([1, 2, 3], {'reverse': True}, False), + ([1, 1, 2, 3], {'reverse': True}, False), + ([1, 10, 2, 3], {'reverse': True}, False), + (['3', '2', '10', '1'], {'reverse': True}, True), + (['3', '2', '10', '1'], {'key': int, 'reverse': True}, False), + # strict + ([], {'strict': True}, True), + ([1], {'strict': True}, True), + ([1, 1], {'strict': True}, False), + ([1, 2, 3], {'strict': True}, True), + ([1, 1, 2, 3], {'strict': True}, False), + ([1, 10, 2, 3], {'strict': True}, False), + (['1', '10', '2', '3'], {'strict': True}, True), + (['1', '10', '2', '3', '3'], {'strict': True}, False), + (['1', '10', '2', '3'], {'strict': True, 'key': int}, False), + ([1, 2, 3], {'strict': True, 'reverse': True}, False), + ([1, 1, 2, 3], {'strict': True, 'reverse': True}, False), + ([1, 10, 2, 3], {'strict': True, 'reverse': True}, False), + (['3', '2', '10', '1'], {'strict': True, 'reverse': True}, True), + ( + ['3', '2', '10', '10', '1'], + {'strict': True, 'reverse': True}, + False, + ), + ( + ['3', '2', '10', '1'], + {'strict': True, 'key': int, 'reverse': True}, + False, + ), + ]: + key = kwargs.get('key', None) + reverse = kwargs.get('reverse', False) + strict = kwargs.get('strict', False) + + with self.subTest( + iterable=iterable, key=key, reverse=reverse, strict=strict + ): + mi_result = mi.is_sorted( + map(BarelySortable, iterable), + key=key, + reverse=reverse, + strict=strict, + ) + + sorted_iterable = sorted(iterable, key=key, reverse=reverse) + if strict: + sorted_iterable = list(mi.unique_justseen(sorted_iterable)) + + py_result = iterable == sorted_iterable + + self.assertEqual(mi_result, expected) + self.assertEqual(mi_result, py_result) + + +class CallbackIterTests(TestCase): + def _target(self, cb=None, exc=None, wait=0): + total = 0 + for i, c in enumerate('abc', 1): + total += i + if wait: + sleep(wait) + if cb: + cb(i, c, intermediate_total=total) + if exc: + raise exc('error in target') + + return total + + def test_basic(self): + func = lambda callback=None: self._target(cb=callback, wait=0.02) + with mi.callback_iter(func, wait_seconds=0.01) as it: + # Execution doesn't start until we begin iterating + self.assertFalse(it.done) + + # Consume everything + self.assertEqual( + list(it), + [ + ((1, 'a'), {'intermediate_total': 1}), + ((2, 'b'), {'intermediate_total': 3}), + ((3, 'c'), {'intermediate_total': 6}), + ], + ) + + # After consuming everything the future is done and the + # result is available. + self.assertTrue(it.done) + self.assertEqual(it.result, 6) + + # This examines the internal state of the ThreadPoolExecutor. This + # isn't documented, so may break in future Python versions. + self.assertTrue(it._executor._shutdown) + + def test_callback_kwd(self): + with mi.callback_iter(self._target, callback_kwd='cb') as it: + self.assertEqual( + list(it), + [ + ((1, 'a'), {'intermediate_total': 1}), + ((2, 'b'), {'intermediate_total': 3}), + ((3, 'c'), {'intermediate_total': 6}), + ], + ) + + def test_partial_consumption(self): + func = lambda callback=None: self._target(cb=callback) + with mi.callback_iter(func) as it: + self.assertEqual(next(it), ((1, 'a'), {'intermediate_total': 1})) + + self.assertTrue(it._executor._shutdown) + + def test_abort(self): + func = lambda callback=None: self._target(cb=callback, wait=0.1) + with mi.callback_iter(func) as it: + self.assertEqual(next(it), ((1, 'a'), {'intermediate_total': 1})) + + with self.assertRaises(mi.AbortThread): + it.result + + def test_no_result(self): + func = lambda callback=None: self._target(cb=callback) + with mi.callback_iter(func) as it: + with self.assertRaises(RuntimeError): + it.result + + def test_no_result_while_running(self): + # Unlike test_no_result, here the function has started but hasn't + # returned, so there is a pending future to interrogate. + gate = Event() + + def func(callback=None): + callback(1, 'a', intermediate_total=1) + gate.wait() + return 6 + + with mi.callback_iter(func, wait_seconds=0.01) as it: + self.assertEqual(next(it), ((1, 'a'), {'intermediate_total': 1})) + self.assertFalse(it.done) + with self.assertRaises(RuntimeError): + it.result + + gate.set() + + def test_exception(self): + func = lambda callback=None: self._target(cb=callback, exc=ValueError) + with mi.callback_iter(func) as it: + self.assertEqual( + next(it), + ((1, 'a'), {'intermediate_total': 1}), + ) + + with self.assertRaises(ValueError): + it.result + + +class WindowedCompleteTests(TestCase): + """Tests for ``windowed_complete()``""" + + def test_basic(self): + actual = list(mi.windowed_complete([1, 2, 3, 4, 5], 3)) + expected = [ + ((), (1, 2, 3), (4, 5)), + ((1,), (2, 3, 4), (5,)), + ((1, 2), (3, 4, 5), ()), + ] + self.assertEqual(actual, expected) + + def test_zero_length(self): + actual = list(mi.windowed_complete([1, 2, 3], 0)) + expected = [ + ((), (), (1, 2, 3)), + ((1,), (), (2, 3)), + ((1, 2), (), (3,)), + ((1, 2, 3), (), ()), + ] + self.assertEqual(actual, expected) + + def test_wrong_length(self): + seq = [1, 2, 3, 4, 5] + for n in (-10, -1, len(seq) + 1, len(seq) + 10): + with self.subTest(n=n): + with self.assertRaises(ValueError): + list(mi.windowed_complete(seq, n)) + + def test_every_partition(self): + every_partition = lambda seq: chain( + *map(partial(mi.windowed_complete, seq), range(len(seq))) + ) + + seq = 'ABC' + actual = list(every_partition(seq)) + expected = [ + ((), (), ('A', 'B', 'C')), + (('A',), (), ('B', 'C')), + (('A', 'B'), (), ('C',)), + (('A', 'B', 'C'), (), ()), + ((), ('A',), ('B', 'C')), + (('A',), ('B',), ('C',)), + (('A', 'B'), ('C',), ()), + ((), ('A', 'B'), ('C',)), + (('A',), ('B', 'C'), ()), + ] + self.assertEqual(actual, expected) + + +class AllUniqueTests(TestCase): + def test_basic(self): + for iterable, expected in [ + ([], True), + ([1, 2, 3], True), + ([1, 1], False), + ([1, 2, 3, 1], False), + ([1, 2, 3, '1'], True), + ]: + with self.subTest(args=(iterable,)): + self.assertEqual(mi.all_unique(iterable), expected) + + def test_non_hashable(self): + self.assertEqual(mi.all_unique([[1, 2], [3, 4]]), True) + self.assertEqual(mi.all_unique([[1, 2], [3, 4], [1, 2]]), False) + + def test_partially_hashable(self): + self.assertEqual(mi.all_unique([[1, 2], [3, 4], (5, 6)]), True) + self.assertEqual( + mi.all_unique([[1, 2], [3, 4], (5, 6), [1, 2]]), False + ) + self.assertEqual( + mi.all_unique([[1, 2], [3, 4], (5, 6), (5, 6)]), False + ) + + def test_key(self): + iterable = ['A', 'B', 'C', 'b'] + self.assertEqual(mi.all_unique(iterable, lambda x: x), True) + self.assertEqual(mi.all_unique(iterable, str.lower), False) + + def test_infinite(self): + self.assertEqual(mi.all_unique(mi.prepend(3, count())), False) + + +class NthProductTests(TestCase): + def test_basic(self): + iterables = ['ab', 'cdef', 'ghi'] + for index, expected in enumerate(product(*iterables)): + actual = mi.nth_product(index, *iterables) + self.assertEqual(actual, expected) + + def test_long(self): + actual = mi.nth_product(1337, range(101), range(22), range(53)) + expected = (1, 3, 12) + self.assertEqual(actual, expected) + + def test_negative(self): + iterables = ['abc', 'de', 'fghi'] + for index, expected in enumerate(product(*iterables)): + actual = mi.nth_product(index - 24, *iterables) + self.assertEqual(actual, expected) + + def test_invalid_index(self): + with self.assertRaises(IndexError): + mi.nth_product(24, 'ab', 'cde', 'fghi') + + def test_repeat(self): + self.assertEqual( + mi.nth_product(1234, 'abcde', repeat=5), + mi.nth_product(1234, 'abcde', 'abcde', 'abcde', 'abcde', 'abcde'), + ) + self.assertEqual( + mi.nth_product(123, 'AB', 'CD', 'EFG', repeat=2), + mi.nth_product(123, 'AB', 'CD', 'EFG', 'AB', 'CD', 'EFG'), + ) + self.assertEqual( + mi.nth_product(123, iter('AB'), iter('CD'), iter('EFG'), repeat=2), + mi.nth_product(123, 'AB', 'CD', 'EFG', 'AB', 'CD', 'EFG'), + ) + + +class NthCombinationWithReplacementTests(TestCase): + def test_basic(self): + iterable = 'abcdefg' + r = 4 + for index, expected in enumerate( + combinations_with_replacement(iterable, r) + ): + actual = mi.nth_combination_with_replacement(iterable, r, index) + self.assertEqual(actual, expected) + self.assertEqual( + mi.nth_combination_with_replacement('abcde', 7, 320), + ('c', 'd', 'e', 'e', 'e', 'e', 'e'), + ) + + def test_long(self): + actual = mi.nth_combination_with_replacement(range(90), 4, 2000000) + expected = (22, 65, 68, 81) + self.assertEqual(actual, expected) + + def test_invalid_r(self): + with self.assertRaises(ValueError): + mi.nth_combination_with_replacement([], -1, 0) + + def test_invalid_index(self): + with self.assertRaises(IndexError): + mi.nth_combination_with_replacement('abcdefg', 3, -85) + with self.assertRaises(IndexError): + mi.nth_combination_with_replacement('abcde', 7, 400) + + +class ValueChainTests(TestCase): + def test_empty(self): + actual = list(mi.value_chain()) + expected = [] + self.assertEqual(actual, expected) + + def test_simple(self): + actual = list(mi.value_chain(1, 2.71828, False, 'foo')) + expected = [1, 2.71828, False, 'foo'] + self.assertEqual(actual, expected) + + def test_more(self): + actual = list(mi.value_chain(b'bar', [1, 2, 3], 4, {'key': 1})) + expected = [b'bar', 1, 2, 3, 4, 'key'] + self.assertEqual(actual, expected) + + def test_empty_lists(self): + actual = list(mi.value_chain(1, 2, [], [3, 4])) + expected = [1, 2, 3, 4] + self.assertEqual(actual, expected) + + def test_complex(self): + obj = object() + actual = list( + mi.value_chain( + (1, (2, (3,))), + ['foo', ['bar', ['baz']], 'tic'], + {'key': {'foo': 1}}, + obj, + ) + ) + expected = [1, (2, (3,)), 'foo', ['bar', ['baz']], 'tic', 'key', obj] + self.assertEqual(actual, expected) + + def test_type_error_while_iterating(self): + # A TypeError raised while consuming an argument means the argument + # is broken, not that it is a scalar - it should propagate. + def gen(): + yield 1 + yield 2 + raise TypeError('failure inside the iterable') + + it = mi.value_chain(gen(), 3) + self.assertEqual([next(it), next(it)], [1, 2]) + with self.assertRaises(TypeError): + next(it) + + def test_type_error_from_iter(self): + # An object that cannot be iterated at all is still emitted as-is. + class NotIterable: + def __iter__(self): + raise TypeError('not iterable') + + obj = NotIterable() + self.assertEqual(list(mi.value_chain(1, obj, [2, 3])), [1, obj, 2, 3]) + + +class ProductIndexTests(TestCase): + def test_basic(self): + iterables = ['ab', 'cdef', 'ghi'] + first_index = {} + for index, element in enumerate(product(*iterables)): + actual = mi.product_index(element, *iterables) + expected = first_index.setdefault(element, index) + self.assertEqual(actual, expected) + + def test_multiplicity(self): + iterables = ['ab', 'bab', 'cab'] + first_index = {} + for index, element in enumerate(product(*iterables)): + actual = mi.product_index(element, *iterables) + expected = first_index.setdefault(element, index) + self.assertEqual(actual, expected) + + def test_long(self): + actual = mi.product_index((1, 3, 12), range(101), range(22), range(53)) + expected = 1337 + self.assertEqual(actual, expected) + + def test_invalid_empty(self): + with self.assertRaises(ValueError): + mi.product_index('', 'ab', 'cde', 'fghi') + + def test_invalid_small(self): + with self.assertRaises(ValueError): + mi.product_index('ac', 'ab', 'cde', 'fghi') + + def test_invalid_large(self): + with self.assertRaises(ValueError): + mi.product_index('achi', 'ab', 'cde', 'fghi') + + def test_invalid_match(self): + with self.assertRaises(ValueError): + mi.product_index('axf', 'ab', 'cde', 'fghi') + + def test_iterator_input(self): + self.assertEqual( + mi.product_index(iter(['i', 'a']), iter('snicker'), iter('snack')), + 12, + ) + + def test_repeat(self): + self.assertEqual( + mi.product_index([1, 2, 3, 4], range(10), repeat=4), + mi.product_index( + [1, 2, 3, 4], range(10), range(10), range(10), range(10) + ), + ) + target = ['B', 'D', 'E', 'A', 'C', 'G'] + self.assertEqual( + mi.product_index(target, 'AB', 'CD', 'EFG', repeat=2), + mi.product_index(target, 'AB', 'CD', 'EFG', 'AB', 'CD', 'EFG'), + ) + self.assertEqual( + mi.product_index( + iter(target), iter('AB'), iter('CD'), iter('EFG'), repeat=2 + ), + mi.product_index(target, 'AB', 'CD', 'EFG', 'AB', 'CD', 'EFG'), + ) + + +class CombinationIndexTests(TestCase): + def test_r_less_than_n(self): + iterable = 'abcdefg' + r = 4 + first_index = {} + for index, element in enumerate(combinations(iterable, r)): + actual = mi.combination_index(element, iterable) + expected = first_index.setdefault(element, index) + self.assertEqual(actual, expected) + + def test_r_equal_to_n(self): + iterable = 'abcd' + r = len(iterable) + first_index = {} + for index, element in enumerate(combinations(iterable, r=r)): + actual = mi.combination_index(element, iterable) + expected = first_index.setdefault(element, index) + self.assertEqual(actual, expected) + + def test_multiplicity(self): + iterable = 'abacba' + r = 3 + first_index = {} + for index, element in enumerate(combinations(iterable, r)): + actual = mi.combination_index(element, iterable) + expected = first_index.setdefault(element, index) + self.assertEqual(actual, expected) + + def test_null(self): + actual = mi.combination_index(tuple(), []) + expected = 0 + self.assertEqual(actual, expected) + + def test_long(self): + actual = mi.combination_index((2, 12, 35, 126), range(180)) + expected = 2000000 + self.assertEqual(actual, expected) + + def test_invalid_order(self): + with self.assertRaises(ValueError): + mi.combination_index(tuple('acb'), 'abcde') + + def test_invalid_large(self): + with self.assertRaises(ValueError): + mi.combination_index(tuple('abcdefg'), 'abcdef') + + def test_invalid_match(self): + with self.assertRaises(ValueError): + mi.combination_index(tuple('axe'), 'abcde') + + +class CombinationWithReplacementIndexTests(TestCase): + def test_none_values(self): + for iterable in ( + [None, 1, 2], + [1, None, 2], + [1, 2, None], + [None, 1, None], + ): + for r in range(4): + first_index = {} + for index, element in enumerate( + combinations_with_replacement(iterable, r) + ): + with self.subTest(iterable=iterable, element=element): + actual = mi.combination_with_replacement_index( + iter(element), iter(iterable) + ) + expected = first_index.setdefault(element, index) + self.assertEqual(actual, expected) + + def test_invalid_none_values(self): + for element, iterable in ( + ((None,), [1, 2]), + ((1, None), [1, 2]), + ((1, None), [None, 1]), + ((None, 2), [1, None]), + ): + with self.subTest(element=element, iterable=iterable): + with self.assertRaises(ValueError): + mi.combination_with_replacement_index(element, iterable) + + def test_r_less_than_n(self): + iterable = 'abcdefg' + r = 4 + first_index = {} + for index, element in enumerate( + combinations_with_replacement(iterable, r) + ): + actual = mi.combination_with_replacement_index(element, iterable) + expected = first_index.setdefault(element, index) + self.assertEqual(actual, expected) + + def test_r_equal_to_n(self): + iterable = 'abcd' + r = len(iterable) + first_index = {} + for index, element in enumerate( + combinations_with_replacement(iterable, r=r) + ): + actual = mi.combination_with_replacement_index(element, iterable) + expected = first_index.setdefault(element, index) + self.assertEqual(actual, expected) + + def test_multiplicity(self): + iterable = 'abacba' + r = 3 + first_index = {} + for index, element in enumerate( + combinations_with_replacement(iterable, r) + ): + actual = mi.combination_with_replacement_index(element, iterable) + expected = first_index.setdefault(element, index) + self.assertEqual(actual, expected) + + def test_null(self): + actual = mi.combination_with_replacement_index(tuple(), []) + expected = 0 + self.assertEqual(actual, expected) + + def test_long(self): + actual = mi.combination_with_replacement_index( + (22, 65, 68, 81), range(90) + ) + expected = 2000000 + self.assertEqual(actual, expected) + + def test_invalid_order(self): + with self.assertRaises(ValueError): + mi.combination_with_replacement_index(tuple('acb'), 'abcde') + + def test_invalid_large(self): + with self.assertRaises(ValueError): + mi.combination_with_replacement_index(tuple('abcdefg'), 'abcdef') + + def test_invalid_match(self): + with self.assertRaises(ValueError): + mi.combination_with_replacement_index(tuple('axe'), 'abcde') + + +class PermutationIndexTests(TestCase): + def test_r_less_than_n(self): + iterable = 'abcdefg' + r = 4 + first_index = {} + for index, element in enumerate(permutations(iterable, r)): + actual = mi.permutation_index(element, iterable) + expected = first_index.setdefault(element, index) + self.assertEqual(actual, expected) + + def test_r_equal_to_n(self): + iterable = 'abcd' + first_index = {} + for index, element in enumerate(permutations(iterable)): + actual = mi.permutation_index(element, iterable) + expected = first_index.setdefault(element, index) + self.assertEqual(actual, expected) + + def test_multiplicity(self): + iterable = 'abacba' + r = 3 + first_index = {} + for index, element in enumerate(permutations(iterable, r)): + actual = mi.permutation_index(element, iterable) + expected = first_index.setdefault(element, index) + self.assertEqual(actual, expected) + + def test_null(self): + actual = mi.permutation_index(tuple(), []) + expected = 0 + self.assertEqual(actual, expected) + + def test_long(self): + actual = mi.permutation_index((2, 12, 35, 126), range(180)) + expected = 11631678 + self.assertEqual(actual, expected) + + def test_invalid_large(self): + with self.assertRaises(ValueError): + mi.permutation_index(tuple('abcdefg'), 'abcdef') + + def test_invalid_match(self): + with self.assertRaises(ValueError): + mi.permutation_index(tuple('axe'), 'abcde') + + +class CountableTests(TestCase): + def test_empty(self): + iterable = [] + it = mi.countable(iterable) + self.assertEqual(it.items_seen, 0) + self.assertEqual(list(it), []) + + def test_basic(self): + iterable = '0123456789' + it = mi.countable(iterable) + self.assertEqual(it.items_seen, 0) + self.assertEqual(next(it), '0') + self.assertEqual(it.items_seen, 1) + self.assertEqual(''.join(it), '123456789') + self.assertEqual(it.items_seen, 10) + + +class ChunkedEvenTests(TestCase): + """Tests for ``chunked_even()``""" + + def test_0(self): + self._test_finite('', 3, []) + + def test_1(self): + self._test_finite('A', 1, [['A']]) + + def test_4(self): + self._test_finite('ABCD', 3, [['A', 'B'], ['C', 'D']]) + + def test_5(self): + self._test_finite('ABCDE', 3, [['A', 'B', 'C'], ['D', 'E']]) + + def test_6(self): + self._test_finite('ABCDEF', 3, [['A', 'B', 'C'], ['D', 'E', 'F']]) + + def test_7(self): + self._test_finite( + 'ABCDEFG', 3, [['A', 'B', 'C'], ['D', 'E'], ['F', 'G']] + ) + + def _test_finite(self, seq, n, expected): + # Check with and without `len()` + self.assertEqual(list(mi.chunked_even(seq, n)), expected) + self.assertEqual(list(mi.chunked_even(iter(seq), n)), expected) + + def test_infinite(self): + for n in range(1, 5): + k = 0 + + def count_with_assert(): + for i in count(): + # Look-ahead should be less than n^2 + self.assertLessEqual(i, n * k + n * n) + yield i + + ls = mi.chunked_even(count_with_assert(), n) + while k < 2: + self.assertEqual(next(ls), list(range(k * n, (k + 1) * n))) + k += 1 + + def test_evenness(self): + for N in range(1, 50): + for n in range(1, N + 2): + lengths = [] + items = [] + for l in mi.chunked_even(range(N), n): + L = len(l) + self.assertLessEqual(L, n) + self.assertGreaterEqual(L, 1) + lengths.append(L) + items.extend(l) + self.assertEqual(items, list(range(N))) + self.assertLessEqual(max(lengths) - min(lengths), 1) + + +class ZipBroadcastTests(TestCase): + def test_single_use_iterable(self): + class SingleUseIterable: + def __init__(self, values): + self.values = values + self.opened = False + + def __iter__(self): + if self.opened: + raise RuntimeError('input already opened') + self.opened = True + return iter(self.values) + + for values in ([], [1, 2]): + for strict in (False, True): + with self.subTest(values=values, strict=strict): + source = SingleUseIterable(values) + self.assertEqual( + list(mi.zip_broadcast('label', source, strict=strict)), + [('label', value) for value in values], + ) + + def test_zip(self): + for objects, zipped, strict_ok in [ + # Empty + ([], [], True), + # One argument + ([1], [(1,)], True), + ([[1]], [(1,)], True), + ([[1, 2]], [(1,), (2,)], True), + # All scalars + ([1, 2], [(1, 2)], True), + ([1, 2, 3], [(1, 2, 3)], True), + # Iterables with length = 0 + ([[], 1], [], True), + ([1, []], [], True), + ([[], []], [], True), + ([[], 1, 2], [], True), + ([[], 1, []], [], True), + ([1, [], 2], [], True), + ([1, [], []], [], True), + ([[], [], 1], [], True), + ([[], [], []], [], True), + # Iterables with length = 1 + ([1, [2]], [(1, 2)], True), + ([[1], 2], [(1, 2)], True), + ([[1], [2]], [(1, 2)], True), + ([1, [2], 3], [(1, 2, 3)], True), + ([1, [2], [3]], [(1, 2, 3)], True), + ([[1], 2, 3], [(1, 2, 3)], True), + ([[1], 2, [3]], [(1, 2, 3)], True), + ([[1], [2], 3], [(1, 2, 3)], True), + ([[1], [2], [3]], [(1, 2, 3)], True), + # Iterables with length > 1 + ([1, [2, 3]], [(1, 2), (1, 3)], True), + ([[1, 2], 3], [(1, 3), (2, 3)], True), + ([[1, 2], [3, 4]], [(1, 3), (2, 4)], True), + ([1, [2, 3], 4], [(1, 2, 4), (1, 3, 4)], True), + ([1, [2, 3], [4, 5]], [(1, 2, 4), (1, 3, 5)], True), + ([[1, 2], 3, 4], [(1, 3, 4), (2, 3, 4)], True), + ([[1, 2], 3, [4, 5]], [(1, 3, 4), (2, 3, 5)], True), + ([[1, 2], [3, 4], 5], [(1, 3, 5), (2, 4, 5)], True), + ([[1, 2], [3, 4], [5, 6]], [(1, 3, 5), (2, 4, 6)], True), + # Iterables with different lengths + ([[], [1]], [], False), + ([[1], []], [], False), + ([[1], [2, 3]], [(1, 2)], False), + ([[1, 2], [3]], [(1, 3)], False), + ([[1, 2], [3], [4]], [(1, 3, 4)], False), + ([[1], [2, 3], [4]], [(1, 2, 4)], False), + ([[1], [2], [3, 4]], [(1, 2, 3)], False), + ([[1], [2, 3], [4, 5]], [(1, 2, 4)], False), + ([[1, 2], [3], [4, 5]], [(1, 3, 4)], False), + ([[1, 2], [3, 4], [5]], [(1, 3, 5)], False), + ([1, [2, 3], [4, 5, 6]], [(1, 2, 4), (1, 3, 5)], False), + ([[1, 2], 3, [4, 5, 6]], [(1, 3, 4), (2, 3, 5)], False), + ([1, [2, 3, 4], [5, 6]], [(1, 2, 5), (1, 3, 6)], False), + ([[1, 2, 3], 4, [5, 6]], [(1, 4, 5), (2, 4, 6)], False), + ([[1, 2], [3, 4, 5], 6], [(1, 3, 6), (2, 4, 6)], False), + ([[1, 2, 3], [4, 5], 6], [(1, 4, 6), (2, 5, 6)], False), + # Infinite + ([count(), 1, [2]], [(0, 1, 2)], False), + ([count(), 1, [2, 3]], [(0, 1, 2), (1, 1, 3)], False), + # Miscellaneous + (['a', [1, 2], [3, 4, 5]], [('a', 1, 3), ('a', 2, 4)], False), + ]: + # Truncate by default + with self.subTest(objects=objects, strict=False, zipped=zipped): + self.assertEqual(list(mi.zip_broadcast(*objects)), zipped) + + # Raise an exception for strict=True + with self.subTest(objects=objects, strict=True, zipped=zipped): + if strict_ok: + self.assertEqual( + list(mi.zip_broadcast(*objects, strict=True)), + zipped, + ) + else: + with self.assertRaises(ValueError): + list(mi.zip_broadcast(*objects, strict=True)) + + def test_scalar_types(self): + # Default: str and bytes are treated as scalar + self.assertEqual( + list(mi.zip_broadcast('ab', [1, 2, 3])), + [('ab', 1), ('ab', 2), ('ab', 3)], + ) + self.assertEqual( + list(mi.zip_broadcast(b'ab', [1, 2, 3])), + [(b'ab', 1), (b'ab', 2), (b'ab', 3)], + ) + # scalar_types=None allows str and bytes to be treated as iterable + self.assertEqual( + list(mi.zip_broadcast('abc', [1, 2, 3], scalar_types=None)), + [('a', 1), ('b', 2), ('c', 3)], + ) + # Use a custom type + self.assertEqual( + list(mi.zip_broadcast({'a': 'b'}, [1, 2, 3], scalar_types=dict)), + [({'a': 'b'}, 1), ({'a': 'b'}, 2), ({'a': 'b'}, 3)], + ) + + +class UniqueInWindowTests(TestCase): + def test_invalid_n(self): + with self.assertRaises(ValueError): + list(mi.unique_in_window([], 0)) + + def test_basic(self): + for iterable, n, expected in [ + (range(9), 10, list(range(9))), + (range(20), 10, list(range(20))), + ([1, 2, 3, 4, 4, 4], 1, [1, 2, 3, 4, 4, 4]), + ([1, 2, 3, 4, 4, 4], 2, [1, 2, 3, 4]), + ([1, 2, 3, 4, 4, 4], 3, [1, 2, 3, 4]), + ([1, 2, 3, 4, 4, 4], 4, [1, 2, 3, 4]), + ([1, 2, 3, 4, 4, 4], 5, [1, 2, 3, 4]), + ( + [0, 1, 1, 1, 1, 1, 1, 1, 1, 0, 2, 3, 4, 2], + 2, + [0, 1, 0, 2, 3, 4, 2], + ), + ]: + with self.subTest(expected=expected): + actual = list(mi.unique_in_window(iterable, n)) + self.assertEqual(actual, expected) + + def test_key(self): + iterable = [0, 1, 3, 4, 5, 6, 7, 8, 9] + n = 3 + key = lambda x: x // 3 + actual = list(mi.unique_in_window(iterable, n, key=key)) + expected = [0, 3, 6, 9] + self.assertEqual(actual, expected) + + +class StrictlyNTests(TestCase): + def test_basic(self): + iterable = ['a', 'b', 'c', 'd'] + n = 4 + actual = list(mi.strictly_n(iter(iterable), n)) + expected = iterable + self.assertEqual(actual, expected) + + def test_too_short_default(self): + iterable = ['a', 'b', 'c', 'd'] + n = 5 + with self.assertRaises(ValueError) as exc: + list(mi.strictly_n(iter(iterable), n)) + + self.assertEqual( + 'Too few items in iterable (got 4)', exc.exception.args[0] + ) + + def test_too_long_default(self): + iterable = ['a', 'b', 'c', 'd'] + n = 3 + with self.assertRaises(ValueError) as cm: + list(mi.strictly_n(iter(iterable), n)) + + self.assertEqual( + 'Too many items in iterable (got at least 4)', + cm.exception.args[0], + ) + + def test_too_short_custom(self): + call_count = 0 + + def too_short(item_count): + nonlocal call_count + call_count += 1 + + iterable = ['a', 'b', 'c', 'd'] + n = 6 + actual = [] + for item in mi.strictly_n(iter(iterable), n, too_short=too_short): + actual.append(item) + expected = ['a', 'b', 'c', 'd'] + self.assertEqual(actual, expected) + self.assertEqual(call_count, 1) + + def test_too_long_custom(self): + import logging + + iterable = ['a', 'b', 'c', 'd'] + n = 2 + too_long = lambda item_count: logging.warning( + 'Picked the first %s items', n + ) + + with self.assertLogs(level='WARNING') as cm: + actual = list(mi.strictly_n(iter(iterable), n, too_long=too_long)) + + self.assertEqual(actual, ['a', 'b']) + self.assertIn('Picked the first 2 items', cm.output[0]) + + +class DuplicatesEverSeenTests(TestCase): + def test_basic(self): + for iterable, expected in [ + ([], []), + ([1, 2, 3], []), + ([1, 1], [1]), + ([1, 2, 1, 2], [1, 2]), + ([1, 2, 3, '1'], []), + ]: + with self.subTest(args=(iterable,)): + self.assertEqual( + list(mi.duplicates_everseen(iterable)), expected + ) + + def test_non_hashable(self): + self.assertEqual(list(mi.duplicates_everseen([[1, 2], [3, 4]])), []) + self.assertEqual( + list(mi.duplicates_everseen([[1, 2], [3, 4], [1, 2]])), [[1, 2]] + ) + + def test_partially_hashable(self): + self.assertEqual( + list(mi.duplicates_everseen([[1, 2], [3, 4], (5, 6)])), [] + ) + self.assertEqual( + list(mi.duplicates_everseen([[1, 2], [3, 4], (5, 6), [1, 2]])), + [[1, 2]], + ) + self.assertEqual( + list(mi.duplicates_everseen([[1, 2], [3, 4], (5, 6), (5, 6)])), + [(5, 6)], + ) + + def test_key_hashable(self): + iterable = 'HEheHEhe' + self.assertEqual(list(mi.duplicates_everseen(iterable)), list('HEhe')) + self.assertEqual( + list(mi.duplicates_everseen(iterable, str.lower)), + list('heHEhe'), + ) + + def test_key_non_hashable(self): + iterable = [[1, 2], [3, 0], [5, -2], [5, 6]] + self.assertEqual( + list(mi.duplicates_everseen(iterable, lambda x: x)), [] + ) + self.assertEqual( + list(mi.duplicates_everseen(iterable, sum)), [[3, 0], [5, -2]] + ) + + def test_key_partially_hashable(self): + iterable = [[1, 2], (1, 2), [1, 2], [5, 6]] + self.assertEqual( + list(mi.duplicates_everseen(iterable, lambda x: x)), [[1, 2]] + ) + self.assertEqual( + list(mi.duplicates_everseen(iterable, list)), [(1, 2), [1, 2]] + ) + + +class DuplicatesJustSeenTests(TestCase): + def test_basic(self): + for iterable, expected in [ + ([], []), + ([1, 2, 3, 3, 2, 2], [3, 2]), + ([1, 1], [1]), + ([1, 2, 1, 2], []), + ([1, 2, 3, '1'], []), + ]: + with self.subTest(args=(iterable,)): + self.assertEqual( + list(mi.duplicates_justseen(iterable)), expected + ) + + def test_non_hashable(self): + self.assertEqual(list(mi.duplicates_justseen([[1, 2], [3, 4]])), []) + self.assertEqual( + list( + mi.duplicates_justseen( + [[1, 2], [3, 4], [3, 4], [3, 4], [1, 2]] + ) + ), + [[3, 4], [3, 4]], + ) + + def test_partially_hashable(self): + self.assertEqual( + list(mi.duplicates_justseen([[1, 2], [3, 4], (5, 6)])), [] + ) + self.assertEqual( + list( + mi.duplicates_justseen( + [[1, 2], [3, 4], (5, 6), [1, 2], [1, 2]] + ) + ), + [[1, 2]], + ) + self.assertEqual( + list( + mi.duplicates_justseen( + [[1, 2], [3, 4], (5, 6), (5, 6), (5, 6)] + ) + ), + [(5, 6), (5, 6)], + ) + + def test_key_hashable(self): + iterable = 'HEheHHHhEheeEe' + self.assertEqual(list(mi.duplicates_justseen(iterable)), list('HHe')) + self.assertEqual( + list(mi.duplicates_justseen(iterable, str.lower)), + list('HHheEe'), + ) + + def test_key_non_hashable(self): + iterable = [[1, 2], [3, 0], [5, -2], [5, 6], [1, 2]] + self.assertEqual( + list(mi.duplicates_justseen(iterable, lambda x: x)), [] + ) + self.assertEqual( + list(mi.duplicates_justseen(iterable, sum)), [[3, 0], [5, -2]] + ) + + def test_key_partially_hashable(self): + iterable = [[1, 2], (1, 2), [1, 2], [5, 6], [1, 2]] + self.assertEqual( + list(mi.duplicates_justseen(iterable, lambda x: x)), [] + ) + self.assertEqual( + list(mi.duplicates_justseen(iterable, list)), [(1, 2), [1, 2]] + ) + + def test_nested(self): + iterable = [[[1, 2], [1, 2]], [5, 6], [5, 6]] + self.assertEqual(list(mi.duplicates_justseen(iterable)), [[5, 6]]) + + +class ClassifyUniqueTests(TestCase): + def test_basic(self): + self.assertEqual( + list(mi.classify_unique('mississippi')), + [ + ('m', True, True), + ('i', True, True), + ('s', True, True), + ('s', False, False), + ('i', True, False), + ('s', True, False), + ('s', False, False), + ('i', True, False), + ('p', True, True), + ('p', False, False), + ('i', True, False), + ], + ) + + def test_non_hashable(self): + self.assertEqual( + list(mi.classify_unique([[1, 2], [3, 4], [3, 4], [1, 2]])), + [ + ([1, 2], True, True), + ([3, 4], True, True), + ([3, 4], False, False), + ([1, 2], True, False), + ], + ) + + def test_partially_hashable(self): + self.assertEqual( + list( + mi.classify_unique( + [[1, 2], [3, 4], (5, 6), (5, 6), (3, 4), [1, 2]] + ) + ), + [ + ([1, 2], True, True), + ([3, 4], True, True), + ((5, 6), True, True), + ((5, 6), False, False), + ((3, 4), True, True), + ([1, 2], True, False), + ], + ) + + def test_key_hashable(self): + iterable = 'HEheHHHhEheeEe' + self.assertEqual( + list(mi.classify_unique(iterable)), + [ + ('H', True, True), + ('E', True, True), + ('h', True, True), + ('e', True, True), + ('H', True, False), + ('H', False, False), + ('H', False, False), + ('h', True, False), + ('E', True, False), + ('h', True, False), + ('e', True, False), + ('e', False, False), + ('E', True, False), + ('e', True, False), + ], + ) + self.assertEqual( + list(mi.classify_unique(iterable, str.lower)), + [ + ('H', True, True), + ('E', True, True), + ('h', True, False), + ('e', True, False), + ('H', True, False), + ('H', False, False), + ('H', False, False), + ('h', False, False), + ('E', True, False), + ('h', True, False), + ('e', True, False), + ('e', False, False), + ('E', False, False), + ('e', False, False), + ], + ) + + def test_key_non_hashable(self): + iterable = [[1, 2], [3, 0], [5, -2], [5, 6], [1, 2]] + self.assertEqual( + list(mi.classify_unique(iterable, lambda x: x)), + [ + ([1, 2], True, True), + ([3, 0], True, True), + ([5, -2], True, True), + ([5, 6], True, True), + ([1, 2], True, False), + ], + ) + self.assertEqual( + list(mi.classify_unique(iterable, sum)), + [ + ([1, 2], True, True), + ([3, 0], False, False), + ([5, -2], False, False), + ([5, 6], True, True), + ([1, 2], True, False), + ], + ) + + def test_key_partially_hashable(self): + iterable = [[1, 2], (1, 2), [1, 2], [5, 6], [1, 2]] + self.assertEqual( + list(mi.classify_unique(iterable, lambda x: x)), + [ + ([1, 2], True, True), + ((1, 2), True, True), + ([1, 2], True, False), + ([5, 6], True, True), + ([1, 2], True, False), + ], + ) + self.assertEqual( + list(mi.classify_unique(iterable, list)), + [ + ([1, 2], True, True), + ((1, 2), False, False), + ([1, 2], False, False), + ([5, 6], True, True), + ([1, 2], True, False), + ], + ) + + def test_vs_unique_everseen(self): + input = 'AAAABBBBCCDAABBB' + output = [e for e, j, u in mi.classify_unique(input) if u] + self.assertEqual(output, ['A', 'B', 'C', 'D']) + self.assertEqual(list(mi.unique_everseen(input)), output) + + def test_vs_unique_everseen_key(self): + input = 'aAbACCc' + output = [e for e, j, u in mi.classify_unique(input, str.lower) if u] + self.assertEqual(output, list('abC')) + self.assertEqual(list(mi.unique_everseen(input, str.lower)), output) + + def test_vs_unique_justseen(self): + input = 'AAAABBBCCDABB' + output = [e for e, j, u in mi.classify_unique(input) if j] + self.assertEqual(output, list('ABCDAB')) + self.assertEqual(list(mi.unique_justseen(input)), output) + + def test_vs_unique_justseen_key(self): + input = 'AABCcAD' + output = [e for e, j, u in mi.classify_unique(input, str.lower) if j] + self.assertEqual(output, list('ABCAD')) + self.assertEqual(list(mi.unique_justseen(input, str.lower)), output) + + def test_vs_duplicates_everseen(self): + input = [1, 2, 1, 2] + output = [e for e, j, u in mi.classify_unique(input) if not u] + self.assertEqual(output, [1, 2]) + self.assertEqual(list(mi.duplicates_everseen(input)), output) + + def test_vs_duplicates_everseen_key(self): + input = 'HEheHEhe' + output = [ + e for e, j, u in mi.classify_unique(input, str.lower) if not u + ] + self.assertEqual(output, list('heHEhe')) + self.assertEqual( + list(mi.duplicates_everseen(input, str.lower)), output + ) + + def test_vs_duplicates_justseen(self): + input = [1, 2, 3, 3, 2, 2] + output = [e for e, j, u in mi.classify_unique(input) if not j] + self.assertEqual(output, [3, 2]) + self.assertEqual(list(mi.duplicates_justseen(input)), output) + + def test_vs_duplicates_justseen_key(self): + input = 'HEheHHHhEheeEe' + output = [ + e for e, j, u in mi.classify_unique(input, str.lower) if not j + ] + self.assertEqual(output, list('HHheEe')) + self.assertEqual( + list(mi.duplicates_justseen(input, str.lower)), output + ) + + +class LongestCommonPrefixTests(TestCase): + def test_basic(self): + iterables = [[1, 2], [1, 2, 3], [1, 2, 4]] + self.assertEqual(list(mi.longest_common_prefix(iterables)), [1, 2]) + + def test_iterators(self): + iterables = iter([iter([1, 2]), iter([1, 2, 3]), iter([1, 2, 4])]) + self.assertEqual(list(mi.longest_common_prefix(iterables)), [1, 2]) + + def test_no_iterables(self): + iterables = [] + self.assertEqual(list(mi.longest_common_prefix(iterables)), []) + + def test_empty_iterables_only(self): + iterables = [[], [], []] + self.assertEqual(list(mi.longest_common_prefix(iterables)), []) + + def test_includes_empty_iterables(self): + iterables = [[1, 2], [1, 2, 3], [1, 2, 4], []] + self.assertEqual(list(mi.longest_common_prefix(iterables)), []) + + def test_non_hashable(self): + # See https://github.com/more-itertools/more-itertools/issues/603 + iterables = [[[1], [2]], [[1], [2], [3]], [[1], [2], [4]]] + self.assertEqual(list(mi.longest_common_prefix(iterables)), [[1], [2]]) + + def test_prefix_contains_elements_of_the_first_iterable(self): + iterables = [[[1], [2]], [[1], [2], [3]], [[1], [2], [4]]] + prefix = list(mi.longest_common_prefix(iterables)) + self.assertIs(prefix[0], iterables[0][0]) + self.assertIs(prefix[1], iterables[0][1]) + self.assertIsNot(prefix[0], iterables[1][0]) + self.assertIsNot(prefix[1], iterables[1][1]) + self.assertIsNot(prefix[0], iterables[2][0]) + self.assertIsNot(prefix[1], iterables[2][1]) + + def test_infinite_iterables(self): + prefix = mi.longest_common_prefix([count(), count()]) + self.assertEqual(next(prefix), 0) + self.assertEqual(next(prefix), 1) + self.assertEqual(next(prefix), 2) + + def test_contains_infinite_iterables(self): + iterables = [[0, 1, 2], count()] + self.assertEqual(list(mi.longest_common_prefix(iterables)), [0, 1, 2]) + + +class IequalsTests(TestCase): + def test_basic(self): + self.assertTrue(mi.iequals("abc", iter("abc"))) + self.assertTrue(mi.iequals(range(3), [0, 1, 2])) + self.assertFalse(mi.iequals("abc", [0, 1, 2])) + + def test_no_iterables(self): + self.assertTrue(mi.iequals()) + + def test_one_iterable(self): + self.assertTrue(mi.iequals("abc")) + + def test_more_than_two_iterable(self): + self.assertTrue(mi.iequals("abc", iter("abc"), ['a', 'b', 'c'])) + self.assertFalse(mi.iequals("abc", iter("abc"), ['a', 'b', 'd'])) + + def test_order_matters(self): + self.assertFalse(mi.iequals("abc", "acb")) + + def test_not_equal_lengths(self): + self.assertFalse(mi.iequals("abc", "ab")) + self.assertFalse(mi.iequals("abc", "bc")) + self.assertFalse(mi.iequals("aaa", "aaaa")) + + def test_empty_iterables(self): + self.assertTrue(mi.iequals([], "")) + + def test_none_is_not_a_sentinel(self): + # See https://stackoverflow.com/a/900444 + self.assertFalse(mi.iequals([1, 2], [1, 2, None])) + self.assertFalse(mi.iequals([1, 2], [None, 1, 2])) + + def test_not_identical_but_equal(self): + self.assertTrue([1, True], [1.0, complex(1, 0)]) + + def test_fillvalue_not_fakeable(self): + # See https://github.com/more-itertools/more-itertools/issues/900 + self.assertFalse(mi.iequals([], [mock.ANY])) + + +class ConstrainedBatchesTests(TestCase): + def test_basic(self): + zen = [ + 'Beautiful is better than ugly', + 'Explicit is better than implicit', + 'Simple is better than complex', + 'Complex is better than complicated', + 'Flat is better than nested', + 'Sparse is better than dense', + 'Readability counts', + ] + for size, expected in ( + ( + 34, + [ + (zen[0],), + (zen[1],), + (zen[2],), + (zen[3],), + (zen[4],), + (zen[5],), + (zen[6],), + ], + ), + ( + 61, + [ + (zen[0], zen[1]), + (zen[2],), + (zen[3], zen[4]), + (zen[5], zen[6]), + ], + ), + ( + 90, + [ + (zen[0], zen[1], zen[2]), + (zen[3], zen[4], zen[5]), + (zen[6],), + ], + ), + ( + 124, + [(zen[0], zen[1], zen[2], zen[3]), (zen[4], zen[5], zen[6])], + ), + ( + 150, + [(zen[0], zen[1], zen[2], zen[3], zen[4]), (zen[5], zen[6])], + ), + ( + 177, + [(zen[0], zen[1], zen[2], zen[3], zen[4], zen[5]), (zen[6],)], + ), + ): + with self.subTest(size=size): + actual = list(mi.constrained_batches(iter(zen), size)) + self.assertEqual(actual, expected) + + def test_max_count(self): + iterable = ['1', '1', '12345678', '12345', '12345'] + max_size = 10 + max_count = 2 + actual = list(mi.constrained_batches(iterable, max_size, max_count)) + expected = [('1', '1'), ('12345678',), ('12345', '12345')] + self.assertEqual(actual, expected) + + def test_max_count_one(self): + self.assertEqual( + list(mi.constrained_batches(['', 'a', ''], 10, max_count=1)), + [('',), ('a',), ('',)], + ) + + def test_nonpositive_max_count(self): + for max_count in (0, -1): + for items in ([], ['a', 'b']): + with self.subTest(max_count=max_count, items=items): + source = iter(items) + with self.assertRaisesRegex( + ValueError, 'maximum count must be greater than zero' + ): + list(mi.constrained_batches(source, 10, max_count)) + self.assertEqual(list(source), items) + + def test_strict(self): + iterable = ['1', '123456789', '1'] + size = 8 + with self.assertRaises(ValueError): + list(mi.constrained_batches(iterable, size)) + + actual = list(mi.constrained_batches(iterable, size, strict=False)) + expected = [('1',), ('123456789',), ('1',)] + self.assertEqual(actual, expected) + + def test_get_len(self): + class Record(tuple): + def total_size(self): + return sum(len(x) for x in self) + + record_3 = Record(('1', '23')) + record_5 = Record(('1234', '1')) + record_10 = Record(('1', '12345678', '1')) + record_2 = Record(('1', '1')) + iterable = [record_3, record_5, record_10, record_2] + + self.assertEqual( + list( + mi.constrained_batches( + iterable, 10, get_len=lambda x: x.total_size() + ) + ), + [(record_3, record_5), (record_10,), (record_2,)], + ) + + def test_bad_max(self): + with self.assertRaises(ValueError): + list(mi.constrained_batches([], 0)) + + +class GrayProductTests(TestCase): + def test_basic(self): + self.assertEqual( + tuple(mi.gray_product(('a', 'b', 'c'), range(1, 3))), + (("a", 1), ("b", 1), ("c", 1), ("c", 2), ("b", 2), ("a", 2)), + ) + out = mi.gray_product(('foo', 'bar'), (3, 4, 5, 6), ['quz', 'baz']) + self.assertEqual(next(out), ('foo', 3, 'quz')) + self.assertEqual( + list(out), + [ + ('bar', 3, 'quz'), + ('bar', 4, 'quz'), + ('foo', 4, 'quz'), + ('foo', 5, 'quz'), + ('bar', 5, 'quz'), + ('bar', 6, 'quz'), + ('foo', 6, 'quz'), + ('foo', 6, 'baz'), + ('bar', 6, 'baz'), + ('bar', 5, 'baz'), + ('foo', 5, 'baz'), + ('foo', 4, 'baz'), + ('bar', 4, 'baz'), + ('bar', 3, 'baz'), + ('foo', 3, 'baz'), + ], + ) + self.assertEqual(tuple(mi.gray_product()), ((),)) + self.assertEqual(tuple(mi.gray_product((1, 2))), ((1,), (2,))) + + def test_errors(self): + with self.assertRaises(ValueError): + list(mi.gray_product((1, 2), ())) + with self.assertRaises(ValueError): + list(mi.gray_product((1, 2), (2,))) + + def test_vs_product(self): + iters = ( + ("a", "b"), + range(3, 6), + [None, None], + {"i", "j", "k", "l"}, + "XYZ", + ) + self.assertEqual( + sorted(product(*iters)), sorted(mi.gray_product(*iters)) + ) + + def test_repeat(self): + self.assertEqual( + list(mi.gray_product('ABC', repeat=5)), + list(mi.gray_product('ABC', 'ABC', 'ABC', 'ABC', 'ABC')), + ) + self.assertEqual( + list(mi.gray_product('ABC', 'DE', repeat=5)), + list( + mi.gray_product( + 'ABC', + 'DE', + 'ABC', + 'DE', + 'ABC', + 'DE', + 'ABC', + 'DE', + 'ABC', + 'DE', + ) + ), + ) + + def test_repeat_with_iterators(self): + self.assertEqual( + list(mi.gray_product(iter('abc'), iter('def'), repeat=2)), + list(mi.gray_product('abc', 'def', repeat=2)), + ) + + +class PartialProductTests(TestCase): + def test_no_iterables(self): + self.assertEqual(tuple(mi.partial_product()), ((),)) + + def test_empty_iterable(self): + self.assertEqual(tuple(mi.partial_product('AB', '', 'CD')), ()) + + def test_one_iterable(self): + # a single iterable should pass through + self.assertEqual( + tuple(mi.partial_product('ABCD')), + ( + ('A',), + ('B',), + ('C',), + ('D',), + ), + ) + + def test_two_iterables(self): + self.assertEqual( + list(mi.partial_product('ABCD', [1])), + [('A', 1), ('B', 1), ('C', 1), ('D', 1)], + ) + expected = [ + ('A', 1), + ('B', 1), + ('C', 1), + ('D', 1), + ('D', 2), + ('D', 3), + ('D', 4), + ] + self.assertEqual( + list(mi.partial_product('ABCD', [1, 2, 3, 4])), expected + ) + + def test_basic(self): + ones = [1, 2, 3] + tens = [10, 20, 30, 40, 50] + hundreds = [100, 200] + + expected = [ + (1, 10, 100), + (2, 10, 100), + (3, 10, 100), + (3, 20, 100), + (3, 30, 100), + (3, 40, 100), + (3, 50, 100), + (3, 50, 200), + ] + + actual = list(mi.partial_product(ones, tens, hundreds)) + self.assertEqual(actual, expected) + + def test_uneven_length_iterables(self): + # this is also the docstring example + expected = [ + ('A', 'C', 'D'), + ('B', 'C', 'D'), + ('B', 'C', 'E'), + ('B', 'C', 'F'), + ] + + self.assertEqual(list(mi.partial_product('AB', 'C', 'DEF')), expected) + + def test_repeat(self): + self.assertEqual( + list(mi.partial_product('ABC', repeat=5)), + list(mi.partial_product('ABC', 'ABC', 'ABC', 'ABC', 'ABC')), + ) + self.assertEqual( + list(mi.partial_product('ABC', 'DE', repeat=5)), + list( + mi.partial_product( + 'ABC', + 'DE', + 'ABC', + 'DE', + 'ABC', + 'DE', + 'ABC', + 'DE', + 'ABC', + 'DE', + ) + ), + ) + + def test_repeat_with_iterators(self): + self.assertEqual( + list(mi.partial_product(iter('abc'), iter('def'), repeat=2)), + list(mi.partial_product('abc', 'def', repeat=2)), + ) + + +class IterateTests(TestCase): + def test_basic(self) -> None: + result = list(islice(mi.iterate(lambda x: 2 * x, start=1), 10)) + expected = [1, 2, 4, 8, 16, 32, 64, 128, 256, 512] + self.assertEqual(result, expected) + + def test_func_controls_iteration_stop(self) -> None: + def func(num): + if num > 100: + raise StopIteration + return num * 2 + + result = list(islice(mi.iterate(func, start=1), 10)) + expected = [1, 2, 4, 8, 16, 32, 64, 128] + self.assertEqual(result, expected) + + +class TakewhileInclusiveTests(TestCase): + def test_basic(self) -> None: + result = list(mi.takewhile_inclusive(lambda x: x < 5, [1, 4, 6, 4, 1])) + expected = [1, 4, 6] + self.assertEqual(result, expected) + + def test_empty_iterator(self) -> None: + result = list(mi.takewhile_inclusive(lambda x: True, [])) + expected = [] + self.assertEqual(result, expected) + + def test_collatz_sequence(self) -> None: + is_even = lambda n: n % 2 == 0 + start = 11 + result = list( + mi.takewhile_inclusive( + lambda n: n != 1, + mi.iterate( + lambda n: n // 2 if is_even(n) else 3 * n + 1, start + ), + ) + ) + expected = [11, 34, 17, 52, 26, 13, 40, 20, 10, 5, 16, 8, 4, 2, 1] + self.assertEqual(result, expected) + + +class OuterProductTests(TestCase): + def test_basic(self) -> None: + greetings = ['Hello', 'Goodbye'] + names = ['Alice', 'Bob', 'Carol'] + greet = lambda greeting, name: f'{greeting}, {name}!' + result = list(mi.outer_product(greet, greetings, names)) + expected = [ + ('Hello, Alice!', 'Hello, Bob!', 'Hello, Carol!'), + ('Goodbye, Alice!', 'Goodbye, Bob!', 'Goodbye, Carol!'), + ] + self.assertEqual(result, expected) + + +class IterSuppressTests(TestCase): + class Producer: + def __init__(self, exc, die_early=False): + self.exc = exc + self.pos = 0 + self.die_early = die_early + + def __iter__(self): + if self.die_early: + raise self.exc + + return self + + def __next__(self): + ret = self.pos + if self.pos >= 5: + raise self.exc + self.pos += 1 + return ret + + def test_no_error(self): + iterator = range(5) + actual = list(mi.iter_suppress(iterator, RuntimeError)) + expected = [0, 1, 2, 3, 4] + self.assertEqual(actual, expected) + + def test_raises_error(self): + iterator = self.Producer(ValueError) + with self.assertRaises(ValueError): + list(mi.iter_suppress(iterator, RuntimeError)) + + def test_suppression(self): + iterator = self.Producer(ValueError) + actual = list(mi.iter_suppress(iterator, RuntimeError, ValueError)) + expected = [0, 1, 2, 3, 4] + self.assertEqual(actual, expected) + + def test_early_suppression(self): + iterator = self.Producer(ValueError, die_early=True) + actual = list(mi.iter_suppress(iterator, RuntimeError, ValueError)) + expected = [] + self.assertEqual(actual, expected) + + +class FilterMapTests(TestCase): + def test_no_iterables(self): + actual = list(mi.filter_map(lambda _: None, [])) + expected = [] + self.assertEqual(actual, expected) + + def test_filter(self): + actual = list(mi.filter_map(lambda _: None, [1, 2, 3])) + expected = [] + self.assertEqual(actual, expected) + + def test_map(self): + actual = list(mi.filter_map(lambda x: x + 1, [1, 2, 3])) + expected = [2, 3, 4] + self.assertEqual(actual, expected) + + def test_filter_map(self): + actual = list( + mi.filter_map( + lambda x: int(x) if x.isnumeric() else None, + ['1', 'a', '2', 'b', '3'], + ) + ) + expected = [1, 2, 3] + self.assertEqual(actual, expected) + + +class PowersetOfSetsTests(TestCase): + def test_simple(self): + iterable = [0, 1, 2] + actual = list(mi.powerset_of_sets(iterable)) + expected = [set(), {0}, {1}, {2}, {0, 1}, {0, 2}, {1, 2}, {0, 1, 2}] + self.assertEqual(actual, expected) + + def test_hash_count(self): + hash_count = 0 + + class Str(str): + def __hash__(self): + nonlocal hash_count + hash_count += 1 + return super().__hash__() + + # Six distinct input values should have 2 ** 6 subsets. + iterable = map(Str, 'ABBBCDDEEEF') + self.assertEqual(mi.ilen(mi.powerset_of_sets(iterable)), 2**6) + + # Eleven input values should have 11 hash calls. + # However, PyPy's version of dict.fromkeys calls hash a + # second time for set inputs. This doubles the 11 to 22. + self.assertIn(hash_count, {11, 22}) + + def test_baseset(self): + iterable = [0, 1, 2] + for kind in (set, frozenset): + ps = list(mi.powerset_of_sets(iterable, baseset=kind)) + self.assertEqual(set(map(type, ps)), {kind}) + + # Verify that an actual set can be formed. + ps = set(mi.powerset_of_sets('abc', baseset=frozenset)) + self.assertIn({'a', 'b'}, ps) + + +class JoinMappingTests(TestCase): + def test_basic(self): + salary_map = {'e1': 12, 'e2': 23, 'e3': 34} + dept_map = {'e1': 'eng', 'e2': 'sales', 'e3': 'eng'} + service_map = {'e1': 5, 'e2': 9, 'e3': 2} + field_to_map = { + 'salary': salary_map, + 'dept': dept_map, + 'service': service_map, + } + expected = { + 'e1': {'salary': 12, 'dept': 'eng', 'service': 5}, + 'e2': {'salary': 23, 'dept': 'sales', 'service': 9}, + 'e3': {'salary': 34, 'dept': 'eng', 'service': 2}, + } + self.assertEqual(dict(mi.join_mappings(**field_to_map)), expected) + + def test_empty(self): + self.assertEqual(dict(mi.join_mappings()), {}) + + +class DiscreteFourierTransformTests(TestCase): + def test_basic(self): + # Example calculation from: + # https://en.wikipedia.org/wiki/Discrete_Fourier_transform#Example + xarr = [1, 2 - 1j, -1j, -1 + 2j] + Xarr = [2, -2 - 2j, -2j, 4 + 4j] + self.assertTrue(all(map(cmath.isclose, mi.dft(xarr), Xarr))) + self.assertTrue(all(map(cmath.isclose, mi.idft(Xarr), xarr))) + + def test_roundtrip(self): + for _ in range(1_000): + N = randrange(35) + xarr = [complex(random(), random()) for i in range(N)] + Xarr = list(mi.dft(xarr)) + assert all(map(cmath.isclose, mi.idft(Xarr), xarr)) + + +class DoubleStarMapTests(TestCase): + def test_construction(self): + iterable = [{'price': 1.23}, {'price': 42}, {'price': 0.1}] + actual = list(mi.doublestarmap('{price:.2f}'.format, iterable)) + expected = ['1.23', '42.00', '0.10'] + self.assertEqual(actual, expected) + + def test_identity(self): + iterable = [{'x': 1}, {'x': 2}, {'x': 3}] + actual = list(mi.doublestarmap(lambda x: x, iterable)) + expected = [1, 2, 3] + self.assertEqual(actual, expected) + + def test_adding(self): + iterable = [{'a': 1, 'b': 2}, {'a': 3, 'b': 4}] + actual = list(mi.doublestarmap(lambda a, b: a + b, iterable)) + expected = [3, 7] + self.assertEqual(actual, expected) + + def test_mismatch_function_smaller(self): + iterable = [{'a': 1, 'b': 2}, {'a': 3, 'b': 4}] + with self.assertRaises(TypeError): + list(mi.doublestarmap(lambda a: a, iterable)) + + def test_mismatch_function_different(self): + iterable = [{'a': 1}, {'a': 2}] + with self.assertRaises(TypeError): + list(mi.doublestarmap(lambda x: x, iterable)) + + def test_mismatch_function_larger(self): + iterable = [{'a': 1}, {'a': 2}] + with self.assertRaises(TypeError): + list(mi.doublestarmap(lambda a, b: a + b, iterable)) + + def test_no_mapping(self): + iterable = [1, 2, 3, 4] + with self.assertRaises(TypeError): + list(mi.doublestarmap(lambda x: x, iterable)) + + def test_empty(self): + actual = list(mi.doublestarmap(lambda x: x, [])) + expected = [] + self.assertEqual(actual, expected) + + +class ArgMinArgMaxTests(TestCase): + def test_basic(self): + for i, iterable, expected_min, expected_max in ( + (1, [10, 2, 20, 5, 17, 4], 1, 2), + (2, [10, -2, -20, 5, 17, 4], 2, 4), + (3, [10, 10, 20, 10], 0, 2), + (4, [30, 30, 20, 30], 2, 0), + ): + with self.subTest(i=i): + self.assertEqual(mi.argmin(iterable), expected_min) + self.assertEqual(mi.argmax(iterable), expected_max) + + def test_key(self): + for i, iterable, key, expected_min, expected_max in ( + (1, [10, -2, -20, 5, 17, 4], abs, 1, 2), + ( + 2, + [[0] * 10, [0] * 5, [0] * 3, [0] * 12, [0] * 2, [0] * 3], + len, + 4, + 3, + ), + ): + with self.subTest(i=i): + self.assertEqual(mi.argmin(iterable, key=key), expected_min) + self.assertEqual(mi.argmax(iterable, key=key), expected_max) + + +class ExtractTests(TestCase): + def test_basics(self): + extract = mi.extract + data = 'abcdefghijklmnopqrstuvwxyz' + + # Test iterator inputs, increasing and decreasing indices, and repeats. + self.assertEqual( + list(extract(iter(data), iter([7, 4, 11, 11, 14]))), + ['h', 'e', 'l', 'l', 'o'], + ) + + # Empty indices + self.assertEqual(list(extract(iter(data), iter([]))), []) + + # Result is an iterator + iterator = extract('abc', [0, 1, 2]) + self.assertTrue(hasattr(iterator, '__next__')) + + # Error cases + + with self.assertRaises(TypeError): + list(extract(None, [])) # Non-iterable data source + with self.assertRaises(TypeError): + list(extract(data, None)) # Non-iterable indices + with self.assertRaises(ValueError): + list(extract(data, [0.0, 1.0, 2.0])) # Non-integer indices + with self.assertRaises(ValueError): + list(extract(data, [1, 2, -3])) # Negative indices + with self.assertRaises(IndexError): + list(extract(data, [1, 2, len(data)])) # Indices out of range + + def test_negative_one_bug(self): + # When the lowest index was exactly -1, it matched the initial + # iterator_position of -1 giving a zero advance step. + extract = mi.extract + + with self.assertRaises(ValueError): + list(extract('abcdefg', [1, 2, -1])) + + def test_none_value_bug(self): + # The buffer used to be a list with unused slots marked with None. + # The mark got conflated with None values in the data stream. + extract = mi.extract + data = ['a', 'b', 'None', 'c', 'd'] + self.assertEqual(list(extract(data, range(5))), data) + + def test_all_orderings(self): + # Thorough test for all cases of five indices to detect + # obscure corner case bugs. + extract = mi.extract + + data = 'abcdefg' + for indices in product(range(6), repeat=5): + with self.subTest(indices=indices): + actual = tuple(extract(data, indices)) + expected = itemgetter(*indices)(data) + self.assertEqual(actual, expected) + + def test_early_free(self): + # No references are held for skipped values or for previously + # emitted values regardless of how long they were in the buffer. + + extract = mi.extract + + class TrackDels(str): + def __del__(self): + dead.add(str(self)) + + dead = set() + iterator = extract(map(TrackDels, 'ABCDEF'), [3, 2, 4, 5]) + + value = next(iterator) + gc.collect() # Force collection on PyPy. + self.assertEqual(value, 'D') # Returns D. Buffered C is alive. + self.assertEqual(dead, {'A', 'B'}) # A and B are dead. + + value = next(iterator) + gc.collect() # Force collection on PyPy + self.assertEqual(value, 'C') # Returns C. + + value = next(iterator) + gc.collect() # Force collection on PyPy + self.assertEqual(value, 'E') # Returns E. + self.assertEqual(dead, {'A', 'B', 'D', 'C'}) # D and C are now dead. + + def test_monotonic(self): + collatz = mi.iterate(lambda x: 3 * x + 1 if x % 2 == 1 else x // 2, 42) + indices = count(0, 2) + self.assertEqual( + mi.take(3, mi.extract(collatz, indices, monotonic=True)), + [42, 64, 16], + ) + self.assertEqual(next(collatz), 8) + self.assertEqual(next(indices), 6) + + # Finite Inputs + self.assertEqual( + list(mi.extract('abcdefgh', [0, 2, 4], monotonic=True)), + ['a', 'c', 'e'], + ) + with self.assertRaises(IndexError): + list(mi.extract('abcdefgh', [0, 2, 40], monotonic=True)) + + # Error cases + with self.assertRaises(ValueError): + list( + mi.extract('abcdefg', [2, 4, 3], monotonic=True) + ) # decreasing index + with self.assertRaises(ValueError): + list( + mi.extract('abcdefg', [-1, 0, 1], monotonic=True) + ) # negative index + + def test_lazy_consumption(self): + extract = mi.extract + + input_stream = mi.peekable(iter('ABCDEFGHIJKLM')) + iterator = extract(input_stream, [4, 2, 10]) + + self.assertEqual(next(iterator), 'E') # C is still buffered + self.assertEqual(input_stream.peek(), 'F') + + self.assertEqual(next(iterator), 'C') + self.assertEqual(input_stream.peek(), 'F') + + # Infinite input + self.assertEqual( + list(extract(count(), [5, 7, 3, 9, 4])), [5, 7, 3, 9, 4] + ) + + +class TestSerialize(TestCase): + def test_concurrent_calls(self): + result = 0 + result_lock = Lock() + + def producer(limit): + 'Non-concurrent producer. A generator version of range(limit).' + for x in range(limit): + yield x + + def consumer(counter): + 'Concurrent data consumer' + nonlocal result + total = 0 + for x in counter: + total += x + with result_lock: + result += total + + limit = 10**6 + counter = mi.serialize(producer(limit)) + workers = [Thread(target=consumer, args=[counter]) for _ in range(10)] + for worker in workers: + worker.start() + for worker in workers: + worker.join() + + self.assertEqual(result, limit * (limit - 1) // 2) + + def test_serialize_generator_methods(self): + # A generator that yields and receives + def echo(): + try: + while True: + val = yield "ready" + yield f"received {val}" + except ValueError: + yield "caught" + + it = mi.serialize(echo()) + + # Test __next__ + self.assertEqual(next(it), "ready") + + # Test send() + self.assertEqual(it.send("hello"), "received hello") + self.assertEqual(next(it), "ready") + + # Test throw() + self.assertEqual(it.throw(ValueError), "caught") + + # Test close() + it.close() + with self.assertRaises(StopIteration): + next(it) + + def test_serialize_methods_attribute_error(self): + # A standard iterator that does not have send/throw/close + # should raise AttributeError when called. + standard_it = mi.serialize([1, 2, 3]) + + with self.assertRaises(AttributeError): + standard_it.send("foo") + + with self.assertRaises(AttributeError): + standard_it.throw(ValueError) + + with self.assertRaises(AttributeError): + standard_it.close() + + def test_serialize_generator_methods_locking(self): + # Verifies that generator methods also acquire the lock. + # We can test this by checking if the lock is held during the call. + + class LockCheckingGenerator: + def __init__(self, lock): + self.lock = lock + + def __iter__(self): + return self + + def send(self, value): + if not self.lock.locked(): + raise RuntimeError("Lock not held during send()") + return value + + def throw(self, *args): + if not self.lock.locked(): + raise RuntimeError("Lock not held during throw()") + + def close(self): + if not self.lock.locked(): + raise RuntimeError("Lock not held during close()") + + # Manually create the serialize object to inspect the lock + it = mi.serialize([]) + mock_gen = LockCheckingGenerator(it._lock) + it._iterator = mock_gen + + # These should not raise RuntimeError + it.send(1) + it.throw(ValueError) + it.close() + + +class TestSynchronized(TestCase): + def test_concurrent_calls(self): + unique = 10 # Number of distinct counters + repetitions = 5 # Number of times each counter is used + limit = 100 # Calls per counter per repetition + + @mi.synchronized + def atomic_counter(): + # This is a generator so that non-concurrent calls are detectable. + # To make calls while running more likely, this code uses random + # time delays. + i = 0 + while True: + yield i + next_i = i + 1 + sleep(random() / 1000) + i = next_i + + def consumer(counter): + for i in range(limit): + next(counter) + + unique_counters = [atomic_counter() for _ in range(unique)] + counters = unique_counters * repetitions + workers = [ + Thread(target=consumer, args=[counter]) for counter in counters + ] + for worker in workers: + worker.start() + for worker in workers: + worker.join() + self.assertEqual( + {next(counter) for counter in unique_counters}, + {limit * repetitions}, + ) + + +class TestConcurrentTee(TestCase): + def test_concurrent_consumers(self): + result = 0 + result_lock = Lock() + + def producer(limit): + 'Non-concurrent producer. A generator version of range(limit).' + for x in range(limit): + yield x + + def consumer(iterator): + 'Concurrent data consumer' + nonlocal result + reconstructed = [x for x in iterator] + if reconstructed == list(range(limit)): + with result_lock: + result += 1 + + limit = 10**5 + num_threads = 100 + non_concurrent_source = producer(limit) + tees = mi.concurrent_tee(non_concurrent_source, n=num_threads) + + # Verify that locks are shared + self.assertEqual(len({id(t_obj.lock) for t_obj in tees}), 1) + + # Run the consumers + workers = [Thread(target=consumer, args=[t_obj]) for t_obj in tees] + for worker in workers: + worker.start() + for worker in workers: + worker.join() + + # Verify that every consumer received 100% of the data (no dups or drops). + self.assertEqual(result, len(tees)) + + # Corner case + non_concurrent_source = producer(limit) + tees = mi.concurrent_tee(non_concurrent_source, n=0) # Zero n + self.assertEqual(tees, ()) + + # Error cases + with self.assertRaises(ValueError): + non_concurrent_source = producer(limit) + mi.concurrent_tee(non_concurrent_source, n=-1) # Negative n + + +def grow_to_window(data, maxlen): + "Return growing window views upto maxlen." + for j in range(1, len(data) + 1): + i = max(j - maxlen, 0) + yield data[i:j] + + +class TestRunningMin(TestCase): + def test_basic(self): + for i, (iterable, expected) in enumerate( + [ + ([], []), + ([1], [1]), + ([1, 2], [1, 1]), + ( + [Fraction(1, 1), Fraction(2, 1)], + [Fraction(1, 1), Fraction(1, 1)], + ), + ( + [Decimal('1.0'), Decimal('2.0')], + [Decimal('1.0'), Decimal('1.0')], + ), + ([8.5, 9.5, 7.5, 6.5], [8.5, 8.5, 7.5, 6.5]), + ] + ): + with self.subTest(i=i, iterable=iterable, expected=expected): + actual = list(mi.running_min(iterable)) + self.assertEqual(actual, expected) + + def test_maxlen(self): + data = choices(range(20), k=1000) + + # Window size must be positive + with self.assertRaises(ValueError): + list(mi.running_min(iter(data), maxlen=0)) + + # Window size of 1 should return the original dataset unchanged + self.assertEqual(list(mi.running_min(iter(data), maxlen=1)), data) + + # Window size normal cases + for maxlen in range(2, 6): + with self.subTest(maxlen=maxlen): + actual = list(mi.running_min(iter(data), maxlen=maxlen)) + expected = list(map(min, grow_to_window(data, maxlen))) + self.assertEqual(actual, expected) + + # Window size larger than the data same as the unbounded case + self.assertEqual( + list(mi.running_min(iter(data), maxlen=len(data) * 2)), + list(mi.running_min(iter(data))), + ) + + def test_stability(self): + # min(x, y) returns x when x == y + data = [0, 0.0, Fraction(0)] + self.assertEqual( + list(map(type, mi.running_min(data, maxlen=2))), + [type(min(data[0:1])), type(min(data[0:2])), type(min(data[1:3]))], + ) + + +class TestRunningMax(TestCase): + def test_basic(self): + for i, (iterable, expected) in enumerate( + [ + ([], []), + ([1], [1]), + ([1, 2], [1, 2]), + ( + [Fraction(1, 1), Fraction(2, 1)], + [Fraction(1, 1), Fraction(2, 1)], + ), + ( + [Decimal('1.0'), Decimal('2.0')], + [Decimal('1.0'), Decimal('2.0')], + ), + ([8.5, 9.5, 7.5, 6.5], [8.5, 9.5, 9.5, 9.5]), + ] + ): + with self.subTest(i=i, iterable=iterable, expected=expected): + actual = list(mi.running_max(iterable)) + self.assertEqual(actual, expected) + + def test_maxlen(self): + data = choices(range(20), k=1000) + + # Window size must be positive + with self.assertRaises(ValueError): + list(mi.running_max(iter(data), maxlen=0)) + + # Window size of 1 should return the original dataset unchanged + self.assertEqual(list(mi.running_max(iter(data), maxlen=1)), data) + + # Window size normal cases + for maxlen in range(2, 6): + with self.subTest(maxlen=maxlen): + actual = list(mi.running_max(iter(data), maxlen=maxlen)) + expected = list(map(max, grow_to_window(data, maxlen))) + self.assertEqual(actual, expected) + + # Window size larger than the data same as the unbounded case + self.assertEqual( + list(mi.running_max(iter(data), maxlen=len(data) * 2)), + list(mi.running_max(iter(data))), + ) + + def test_stability(self): + # max(x, y) returns x when x == y + data = [0, 0.0, Fraction(0)] + self.assertEqual( + list(map(type, mi.running_max(data, maxlen=2))), + [type(max(data[0:1])), type(max(data[0:2])), type(max(data[1:3]))], + ) + + +class TestRunningStats(TestCase): + # Assumes that the component functions are already tested, + # so we only need to test whether they are properly integrated + # (i.e correctly pass *maxlen* to each and correctly send + # each statistic to the correct slot in Stats). + + def test_single_example(self): + data = [4, 3, 7, 0, 8, 1, 6, 2, 9, 5] + + it = mi.running_statistics(iter(data)) + self.assertEqual( + next(it), + mi.Stats(size=1, minimum=4, median=4, maximum=4, mean=4.0), + ) + self.assertEqual( + next(it), + mi.Stats(size=2, minimum=3, median=3.5, maximum=4, mean=7 / 2), + ) + self.assertEqual( + next(it), + mi.Stats(size=3, minimum=3, median=4, maximum=7, mean=14 / 3), + ) + self.assertEqual( + next(it), + mi.Stats(size=4, minimum=0, median=3.5, maximum=7, mean=14 / 4), + ) + self.assertEqual( + next(it), + mi.Stats(size=5, minimum=0, median=4, maximum=8, mean=22 / 5), + ) + w = [4, 3, 7, 0, 8, 1] # Unbounded window + self.assertEqual( + next(it), + mi.Stats( + size=len(w), + minimum=min(w), + median=median(w), + maximum=max(w), + mean=mean(w), + ), + ) + w = [4, 3, 7, 0, 8, 1, 6] # Unbounded window + self.assertEqual( + next(it), + mi.Stats( + size=len(w), + minimum=min(w), + median=median(w), + maximum=max(w), + mean=mean(w), + ), + ) + + it = mi.running_statistics(iter(data), maxlen=3) + self.assertEqual( + next(it), + mi.Stats(size=1, minimum=4, median=4, maximum=4, mean=4.0), + ) + self.assertEqual( + next(it), + mi.Stats(size=2, minimum=3, median=3.5, maximum=4, mean=7 / 2), + ) + self.assertEqual( + next(it), + mi.Stats(size=3, minimum=3, median=4, maximum=7, mean=14 / 3), + ) + self.assertEqual( + next(it), + mi.Stats(size=3, minimum=0, median=3, maximum=7, mean=10 / 3), + ) + self.assertEqual( + next(it), + mi.Stats(size=3, minimum=0, median=7, maximum=8, mean=15 / 3), + ) + w = [0, 8, 1] # Window with maxlen=3 + self.assertEqual( + next(it), + mi.Stats( + size=len(w), + minimum=min(w), + median=median(w), + maximum=max(w), + mean=mean(w), + ), + ) + w = [8, 1, 6] # Window with maxlen=3 + self.assertEqual( + next(it), + mi.Stats( + size=len(w), + minimum=min(w), + median=median(w), + maximum=max(w), + mean=mean(w), + ), + ) + + def test_early_error_detection(self): + with self.assertRaises(TypeError): + mi.running_statistics(None) # Non-iterable input + + with self.assertRaises(ValueError): + mi.running_statistics([10, 5, 20], maxlen=0) # Non-positive maxlen + + def test_stat_properties(self): + st = mi.Stats(size=5, minimum=0, median=4, maximum=8, mean=22 / 5) + + # Comparable + self.assertEqual( + st, mi.Stats(size=5, minimum=0, median=4, maximum=8, mean=22 / 5) + ) + + # Hashable + hash(st) + + # Slots + with self.assertRaises(TypeError): + vars(st) + + # Immutable + with self.assertRaises(AttributeError): + st.size = 10 + + # Not sizeable, indexable, or iterable + # (so we can add new fields in the future) + with self.assertRaises(TypeError): + len(st) + with self.assertRaises(TypeError): + st[0] + with self.assertRaises(TypeError): + iter(st) + + def test_datatypes(self): + data = [4, 3, 7, 0, 8, 1, 6, 2, 9, 5] + for kind in (int, float, Fraction, Decimal): + list(mi.running_statistics(map(kind, data))) + list(mi.running_statistics(map(kind, data), maxlen=3)) + with self.assertRaises(TypeError): + list(mi.running_statistics(map(complex, data))) + with self.assertRaises(TypeError): + list(mi.running_statistics(map(complex, data), maxlen=3)) + + +class TestSubfactorial(TestCase): + def test_oeis_baseline(self): + # https://oeis.org/A000166 + reference_cases = """ + 1, 0, 1, 2, 9, 44, 265, 1854, 14833, 133496, 1334961, 14684570, + 176214841, 2290792932, 32071101049, 481066515734, 7697064251745, + 130850092279664, 2355301661033953, 44750731559645106, + 895014631192902121, 18795307255050944540, 413496759611120779881, + 9510425471055777937262 + """.replace(',', ' ').split() + for n, target_str in enumerate(reference_cases): + with self.subTest(n=n): + self.assertEqual(mi.subfactorial(n), int(target_str)) + + def test_vs_derangements(self): + for n in range(10): + with self.subTest(n=n): + self.assertEqual( + mi.subfactorial(n), mi.ilen(mi.derangements('.' * n)) + ) + + def test_error_cases(self): + with self.assertRaises(TypeError): + mi.subfactorial(5.0) # float input + with self.assertRaises(TypeError): + mi.subfactorial('5') # string input + with self.assertRaises(ValueError): + mi.subfactorial(-1) # negative input + + +class TestRandomOrderedRange(TestCase): + test_inputs = [ + (-1,), + (0,), + (1,), + (2,), + (3,), + (4,), + (10,), + (50,), + (10_000,), + (100, 200), + (1000, 2000, 10), + (2000, 1000, -10), + ] + + def test_vs_ordered_baseline(self): + # Verify that the randomized values match the ordered values. + for args in self.test_inputs: + with self.subTest(args=args): + randomized = sorted(mi.random_ordered_range(*args)) + ordered = sorted(range(*args)) + self.assertEqual(randomized, ordered) + + def test_lcg_mitigation(self): + # Without mitigation the LCG doesn't produce many distinct permutations + perms = set(tuple(mi.random_ordered_range(6)) for _ in range(10**5)) + self.assertEqual(len(perms), factorial(6))