Files
RustPython/extra_tests/snippets/stdlib_itertools.py
Jeong, YunWon d5d19e0611 Fix 33 reproduced fuzzing and static-review defects (#8514)
* mmap: guard move() against dest/src past the mapping size

An offset larger than the mapping made `size - dest` underflow, producing
an out-of-range slice index panic. Check the high side first, as write() does.

Assisted-by: Claude

* exceptions: keep ImportError name/path/name_from out of an empty instance dict

ImportError.__reduce__ unwrapped get_arg(0), which is None when the exception
was constructed with no positional argument, aborting on pickle.dumps(ImportError()).
Fall back to the full args tuple there, and expose name/path/name_from as class
attributes defaulting to None so a bare ImportError() reduces to (cls, ()).

Assisted-by: Claude

* asyncio: stop unwrapping the _current_tasks downcast in current_task()

_asyncio._current_tasks is a reassignable module attribute; the three sibling
functions already degrade gracefully when it is not a dict, current_task() did not.

Assisted-by: Claude

* asyncio: raise TypeError when FutureIter.throw()'s exception class returns a non-exception

exc_type.call() goes through a Python-controlled __new__, so its result can be
any object; the downcast was unwrapped. Report it as a TypeError instead.

Assisted-by: Claude

* sequence: use a fallible reservation for sequence repetition

The guard only rejects a repeat whose per-element size crosses MAX_MEMORY_SIZE,
so (1,) * (10**12) still reached Vec::with_capacity and aborted in the allocator.
Reserve fallibly and surface a MemoryError.

Assisted-by: Claude

* collections: reject an oversized deque repetition with MemoryError

deque * sys.maxsize reached the allocator and aborted with a capacity overflow.
Apply the MAX_MEMORY_SIZE guard the other sequences already carry.

Assisted-by: Claude

* itertools: raise OverflowError for an out-of-range r argument

combinations/combinations_with_replacement/permutations narrowed r with
to_usize().unwrap(), so r=2**64 panicked instead of raising.

Assisted-by: Claude

* math: stream the generic sumprod path instead of collecting both iterables

The big-int path collected each argument into a Vec before multiplying, so an
unbounded iterable exhausted memory. Advance both iterators in lockstep and
accumulate, reporting a length mismatch when only one is exhausted.

Assisted-by: Claude

* _imp: raise TypeError for a second positional argument to find_frozen

withdata is keyword-only and unimplemented; passing it positionally hit an
unimplemented!() and aborted.

Assisted-by: Claude

* _typing: check _idfunc arity before indexing args

_typing._idfunc() with no argument indexed args[0] out of bounds.

Assisted-by: Claude

* exceptions: require a sequence for the ExceptionGroup excs argument

The argument was collected before being validated, so an unbounded iterable
such as itertools.count() exhausted memory. Reject a non-sequence up front.

Assisted-by: Claude

* mmap: treat an inverted find/rfind range as empty

find(b"x", 5, 2) built a slice whose start exceeded its end and panicked.

Assisted-by: Claude

* collections: give deque and defaultdict a GC traverse

Neither type opted into traversal, so a reference cycle through a deque or a
defaultdict was never collected.

Assisted-by: Claude

* itertools: opt the iterator types into GC traverse

cycle and its siblings hold Python references but declared no traverse, so a
cycle built through one of them leaked.

Assisted-by: Claude

* _ctypes: mask an out-of-range int instead of panicking

c_char_p(2**64) and pointer item assignment called .expect() on the narrowing
conversion. Wrap the value to the target width, which is what the C
implementation stores.

Assisted-by: Claude

* hashlib: import _hashlib when a hash module is loaded

The hash object type is a static type owned by _hashlib; calling _md5.md5()
without _hashlib imported hit an uninitialized static type and panicked.

Assisted-by: Claude

* _csv: fall back to the built-in dialect defaults when none is registered

The dialect table is empty until csv.py registers 'excel', so _csv.reader([])
unwrapped a missing entry and panicked.

Assisted-by: Claude

* builtins: reject surrogates in compile()/eval() source instead of panicking

expect_str() panics on a str containing surrogates; convert with try_as_utf8
so eval(chr(0xd800)) raises.

Assisted-by: Claude

* _suggestions: require a list for _generate_suggestions candidates

The argument was collected before validation, so an unbounded iterable
exhausted memory.

Assisted-by: Claude

* lzma: size the filter chain before consuming it

filters= was collected into a Vec before the length check, so an unbounded
iterable exhausted memory. Take the length through the sequence protocol first.

Assisted-by: Claude

* classmethod: opt into GC traverse

staticmethod already declares traverse; classmethod did not, so a cycle through
the wrapped callable leaked.

Assisted-by: Claude

* posix: reject unbounded iterables in posix_spawn and setgroups

argv, setsigdef, setsigmask and setgroups bound an ArgIterable and collected it
before validating, so an infinite generator exhausted memory. Require a
list/tuple for argv, validate signals while streaming, and take setgroups
through the sequence protocol.

Assisted-by: Claude

* _ctypes: size Array slice assignment and _argtypes_ before collecting

Both eagerly collected their argument, so an unbounded iterable exhausted
memory before the length check ran.

Assisted-by: Claude

* sys: propagate the breakpointhook warning failure

warn() was unwrapped, so an unimportable $PYTHONBREAKPOINT under
-W error panicked instead of raising.

Assisted-by: Claude

* structseq: require the sequence argument when constructing a struct sequence

A no-argument construction produced an empty backing tuple, and reading any
named field then indexed out of bounds.

Assisted-by: Claude

* Guard the hash slot dispatch against unbounded recursion

`PyObject::hash` invoked the type's hash slot directly, so an element-wise
`__hash__` following a deeply nested object graph recursed one native frame
per level and overflowed the stack. Wrap the dispatch in `with_recursion`,
matching the repr and rich-compare dispatches in the same file, so
`hash(x)` on a nested tuple/GenericAlias/slice raises `RecursionError`.

Assisted-by: Claude

* genericalias: guard the __parameters__ walk against unbounded recursion

`make_parameters_from_slice` recursed into every list/tuple argument with
nothing counting the frames, so `list[L]` for a self-referential or deeply
nested `L` overflowed the native stack at subscript time. Wrap the descent
in `with_recursion`, which makes the walk fallible; `PyGenericAlias::new`,
`from_args` and `make_parameters` now return `PyResult` and every caller
propagates.

Assisted-by: Claude

* Fix the type confusion in PyAtomicRef's Debug impl

`PyAtomicRef<T>` stores a pointer to a `Py<T>`, as `Deref`, `load_raw`,
`swap` and `Drop` all read it, but `Debug` cast it to a bare `T` and
formatted the object header as payload bytes. For `PyFunction`, whose
`code: PyAtomicRef<PyCode>` has a pointer-chasing `Debug`, that
dereferenced header words and segfaulted.

Cast to `PyObject` instead, which is what `Drop` already does and which
also covers the `PyAtomicRef<PyObject>` and `PyAtomicRef<Option<T>>`
instantiations that have no `Py<T>`.

`_asyncio._enter_task` reached this through `{:?}` in its "Cannot enter
into task" message; format the two tasks with their Python repr, which is
what the message is meant to show.

Assisted-by: Claude

* _sre: disallow instantiating Match

`Match` inherited `object.__new__`, so `M.__new__(M)` produced an instance
whose `regs`/`string`/`pattern` were never filled in by a match run; the
mapping subscript path read them and dereferenced garbage. The type has no
public constructor, so mark it `DISALLOW_INSTANTIATION`, which makes
`M.__new__(M)` raise `TypeError: cannot create 're.Match' instances`.

Assisted-by: Claude

* utils: return the empty repr instead of asserting a non-empty collection

`collection_repr` took the first element with an `.expect()` justified by the
caller's preceding non-empty check. Another thread clearing the collection
between that check and the iteration made the iterator yield nothing and
panicked the worker. Fall back to the caller-supplied empty form, which is
the text those callers already produce for an empty collection.

Assisted-by: Claude

* itertools: advance cycle's index atomically

`cycle.__next__` did `fetch_add(1)` and then reset the index to 0 in a
separate store, so two threads replaying the saved items could both read an
index past the end of `saved` and panic on the slice access. Do the advance
and the wrap in one `fetch_update`.

Assisted-by: Claude

* _ctypes: reject a float argument to a foreign call without argtypes

`conv_param`, the conversion used when `argtypes` is not set, converted its
argument with `try_int`, which goes through `__int__` and so accepts a
float. `libc.strlen(1.5)` therefore passed 1 where the callee expects a
`char *` and the callee dereferenced it. Match `ConvParam`, which does a
`PyLong_Check` and converts nothing: take the branch only for an `int` (or
a subclass, so `True` still converts), and let a float fall through to
"Don't know how to convert parameter".

The branch below it converted a float to a C double, but `try_int` claimed
every float before it could run, so it was dead; `CArgValue::Double`
existed only for that branch and both go. Typed doubles are unaffected —
they travel as `CArgValue::Typed` with code 'd'.

Assisted-by: Claude

* Report the iterator itself from PyIter's traverse

`Traverse for PyIter<O>` delegated to the inherent `PyObject::traverse` of
the object it wraps, so it reported that iterator's referents instead of
the iterator. The iterator's own reference to those referents was then
never subtracted during the collector's reference-subtraction pass, the
referents kept a non-zero gc_refs, and every object reachable from them
was classified as a root.

Any cycle running through a type with a `PyIter` field therefore survived
collection: `map`, `filter`, `zip`, `enumerate`, `reversed` and the
`itertools` iterators all leaked, while the same cycle through a `list`,
`tuple` or `list_iterator` collected. Report the wrapped object, as the
`PyObjectRef`, `PyRef<T>` and `PyStackRef` impls do.

`itertools.tee` still leaks: its shared buffer is a `PyRc<PyItertoolsTeeData>`
rather than a Python object, so the collector cannot see through it.

Assisted-by: Claude

* Remove the obsolete expectedFailure on test_code_module.test_unicode_error

Compiling a source string containing a lone surrogate now raises
UnicodeEncodeError, so the test passes.

Assisted-by: Claude

* itertools: reserve the combination indices fallibly

`combinations` and `combinations_with_replacement` built their index vector
with an infallible allocation, so an `r` that passes the ssize_t check but
does not fit in memory aborted the process instead of raising MemoryError.

Assisted-by: Claude

* Apply the struct sequence constructor's dict argument

`structseq(sequence, dict)` discarded its second argument, so the hidden
fields past `n_sequence_fields` — `tm_zone`, `st_atime` and friends — were
always None when constructed directly or restored from a `(sequence, dict)`
pickle, and a non-dict second argument was accepted silently.

Take the dict, require it to be a dict, and fill the hidden slots the
sequence did not cover from it. A key that names a field the sequence
already supplied, or no field at all, is now a
"got duplicate or unexpected field name(s)" TypeError instead of being
dropped. Both arguments are bindable by name, as `sequence` and `dict`.

`os.stat_result` and `os.statvfs_result` did not accept a second argument
at all; they and `time.struct_time` now share the parsing.

Assisted-by: Claude

* _imp: report the argument count in find_frozen's arity error

Assisted-by: Claude

* Add regression tests for the reproduced crashers

One case per catalog entry, each asserting the behavior the fix produces:
recursion guards, the memory-unsafety sites, the overflow and unbounded
allocation guards, the eager-collection rejections, the unwrap sites, and
the cycles the collector now breaks. Every expected value was checked
against CPython 3.14.

Assisted-by: Claude

* Tolerate a changed-size RuntimeError in the set repr stress test

The mutator and reader threads race on purpose; a "changed size during
iteration" RuntimeError is a valid outcome of that race and should not fail
the test. The panic it guards against is not.

Assisted-by: Claude

* Move the crash regression tests into per-module snippets

crash_regressions.py collected every reproduced crasher in one file. Split it
into the snippet for the module each case exercises, and add stdlib_gc.py,
stdlib_asyncio.py, stdlib_lzma.py, stdlib_threading_set_repr.py and
stdlib_threading_itertools_cycle.py for the cases with no existing home.

The suite runs every snippet under the host CPython too, so the checks only
RustPython raises are guarded by sys.implementation.name: the hash and
__parameters__ recursion depth, and the deque repeat overflow. The float
ctypes argument accepts either TypeError or ctypes.ArgumentError.

Assisted-by: Claude
2026-08-14 00:57:10 +09:00

543 lines
12 KiB
Python

import itertools
from testutils import assert_raises
# itertools.chain tests
chain = itertools.chain
# empty
assert list(chain()) == []
assert list(chain([], "", b"", ())) == []
assert list(chain([1, 2, 3, 4])) == [1, 2, 3, 4]
assert list(chain("ab", "cd", (), "e")) == ["a", "b", "c", "d", "e"]
with assert_raises(TypeError):
list(chain(1))
x = chain("ab", 1)
assert next(x) == "a"
assert next(x) == "b"
with assert_raises(TypeError):
next(x)
# empty
with assert_raises(TypeError):
chain.from_iterable()
with assert_raises(TypeError):
chain.from_iterable("abc", "def")
with assert_raises(TypeError):
# iterables are lazily evaluated -- can be constructed but will fail to execute
list(chain.from_iterable([1, 2, 3]))
with assert_raises(TypeError):
list(chain(1))
args = ["abc", "def"]
assert list(chain.from_iterable(args)) == ["a", "b", "c", "d", "e", "f"]
args = [[], "", b"", ()]
assert list(chain.from_iterable(args)) == []
args = ["ab", "cd", (), "e"]
assert list(chain.from_iterable(args)) == ["a", "b", "c", "d", "e"]
x = chain.from_iterable(["ab", 1])
assert next(x) == "a"
assert next(x) == "b"
with assert_raises(TypeError):
next(x)
# iterables are lazily evaluated
x = chain.from_iterable(itertools.repeat(range(2)))
assert next(x) == 0
assert next(x) == 1
assert next(x) == 0
assert next(x) == 1
x = chain(1, [2])
with assert_raises(TypeError):
next(x)
with assert_raises(StopIteration):
next(x)
# itertools.count tests
# default arguments
c = itertools.count()
assert next(c) == 0
assert next(c) == 1
assert next(c) == 2
# positional
c = itertools.count(2, 3)
assert next(c) == 2
assert next(c) == 5
assert next(c) == 8
# backwards
c = itertools.count(1, -10)
assert next(c) == 1
assert next(c) == -9
assert next(c) == -19
# step = 0
c = itertools.count(5, 0)
assert next(c) == 5
assert next(c) == 5
# itertools.count TODOs: kwargs and floats
# step kwarg
# c = itertools.count(step=5)
# assert next(c) == 0
# assert next(c) == 5
# start kwarg
# c = itertools.count(start=10)
# assert next(c) == 10
# float start
# c = itertools.count(0.5)
# assert next(c) == 0.5
# assert next(c) == 1.5
# assert next(c) == 2.5
# float step
# c = itertools.count(1, 0.5)
# assert next(c) == 1
# assert next(c) == 1.5
# assert next(c) == 2
# float start + step
# c = itertools.count(0.5, 0.5)
# assert next(c) == 0.5
# assert next(c) == 1
# assert next(c) == 1.5
# itertools.cycle tests
r = itertools.cycle([1, 2, 3])
assert next(r) == 1
assert next(r) == 2
assert next(r) == 3
assert next(r) == 1
assert next(r) == 2
assert next(r) == 3
assert next(r) == 1
r = itertools.cycle([1])
assert next(r) == 1
assert next(r) == 1
assert next(r) == 1
r = itertools.cycle([])
with assert_raises(StopIteration):
next(r)
with assert_raises(TypeError):
itertools.cycle(None)
with assert_raises(TypeError):
itertools.cycle(10)
# itertools.repeat tests
# no times
r = itertools.repeat(5)
assert next(r) == 5
assert next(r) == 5
assert next(r) == 5
# times
r = itertools.repeat(1, 2)
assert next(r) == 1
assert next(r) == 1
with assert_raises(StopIteration):
next(r)
# times = 0
r = itertools.repeat(1, 0)
with assert_raises(StopIteration):
next(r)
# negative times
r = itertools.repeat(1, -1)
with assert_raises(StopIteration):
next(r)
# itertools.starmap tests
starmap = itertools.starmap
assert list(starmap(pow, zip(range(3), range(1, 7)))) == [0**1, 1**2, 2**3]
assert list(starmap(pow, [])) == []
assert list(starmap(pow, [iter([4, 5])])) == [4**5]
with assert_raises(TypeError):
starmap(pow)
# itertools.takewhile tests
from itertools import takewhile as tw
t = tw(lambda n: n < 5, [1, 2, 5, 1, 3])
assert next(t) == 1
assert next(t) == 2
with assert_raises(StopIteration):
next(t)
# not iterable
with assert_raises(TypeError):
tw(lambda n: n < 1, 1)
# not callable
t = tw(5, [1, 2])
with assert_raises(TypeError):
next(t)
# non-bool predicate
t = tw(lambda n: n, [1, 2, 0])
assert next(t) == 1
assert next(t) == 2
with assert_raises(StopIteration):
next(t)
# bad predicate prototype
t = tw(lambda: True, [1])
with assert_raises(TypeError):
next(t)
# StopIteration before attempting to call (bad) predicate
t = tw(lambda: True, [])
with assert_raises(StopIteration):
next(t)
# doesn't try again after the first predicate failure
t = tw(lambda n: n < 1, [1, 0])
with assert_raises(StopIteration):
next(t)
with assert_raises(StopIteration):
next(t)
# itertools.islice tests
def assert_matches_seq(it, seq):
assert list(it) == list(seq)
i = itertools.islice
it = i([1, 2, 3, 4, 5], 3)
assert_matches_seq(it, [1, 2, 3])
it = i([0.5, 1, 1.5, 2, 2.5, 3, 4, 5], 1, 6, 2)
assert_matches_seq(it, [1, 2, 3])
it = i([1, 2], None)
assert_matches_seq(it, [1, 2])
it = i([1, 2, 3], None, None, None)
assert_matches_seq(it, [1, 2, 3])
it = i([1, 2, 3], 1, None, None)
assert_matches_seq(it, [2, 3])
it = i([1, 2, 3], None, 2, None)
assert_matches_seq(it, [1, 2])
it = i([1, 2, 3], None, None, 3)
assert_matches_seq(it, [1])
# itertools.filterfalse
it = itertools.filterfalse(lambda x: x % 2, range(10))
assert 0 == next(it)
assert 2 == next(it)
assert 4 == next(it)
assert 6 == next(it)
assert 8 == next(it)
with assert_raises(StopIteration):
next(it)
l = [0, 1, None, False, True, [], {}]
it = itertools.filterfalse(None, l)
assert 0 == next(it)
assert None == next(it)
assert False == next(it)
assert [] == next(it)
assert {} == next(it)
# itertools.dropwhile
it = itertools.dropwhile(lambda x: x < 5, [1, 4, 6, 4, 1])
assert 6 == next(it)
assert 4 == next(it)
assert 1 == next(it)
with assert_raises(StopIteration):
next(it)
# itertools.accumulate
it = itertools.accumulate([6, 3, 7, 1, 0, 9, 8, 8])
assert 6 == next(it)
assert 9 == next(it)
assert 16 == next(it)
assert 17 == next(it)
assert 17 == next(it)
assert 26 == next(it)
assert 34 == next(it)
assert 42 == next(it)
with assert_raises(StopIteration):
next(it)
it = itertools.accumulate([3, 2, 4, 1, 0, 5, 8], lambda a, v: a * v)
assert 3 == next(it)
assert 6 == next(it)
assert 24 == next(it)
assert 24 == next(it)
assert 0 == next(it)
assert 0 == next(it)
assert 0 == next(it)
with assert_raises(StopIteration):
next(it)
# itertools.compress
assert list(itertools.compress("ABCDEF", [1, 0, 1, 0, 1, 1])) == list("ACEF")
assert list(itertools.compress("ABCDEF", [0, 0, 0, 0, 0, 0])) == list("")
assert list(itertools.compress("ABCDEF", [1, 1, 1, 1, 1, 1])) == list("ABCDEF")
assert list(itertools.compress("ABCDEF", [1, 0, 1])) == list("AC")
assert list(itertools.compress("ABC", [0, 1, 1, 1, 1, 1])) == list("BC")
assert list(itertools.compress("ABCDEF", [True, False, "t", "", 1, 9])) == list("ACEF")
# itertools.tee
t = itertools.tee([])
assert len(t) == 2
assert t[0] is not t[1]
assert list(t[0]) == list(t[1]) == []
with assert_raises(TypeError):
itertools.tee()
t = itertools.tee(range(1000))
assert list(t[0]) == list(t[1]) == list(range(1000))
t = itertools.tee([1, 22, 333], 3)
assert len(t) == 3
assert 1 == next(t[0])
assert 1 == next(t[1])
assert 22 == next(t[0])
assert 333 == next(t[0])
assert 1 == next(t[2])
assert 22 == next(t[1])
assert 333 == next(t[1])
assert 22 == next(t[2])
with assert_raises(StopIteration):
next(t[0])
assert 333 == next(t[2])
with assert_raises(StopIteration):
next(t[2])
with assert_raises(StopIteration):
next(t[1])
t0, t1 = itertools.tee([1, 2, 3])
tc = t0.__copy__()
assert list(t0) == [1, 2, 3]
assert list(t1) == [1, 2, 3]
assert list(tc) == [1, 2, 3]
t0, t1 = itertools.tee([1, 2, 3])
assert 1 == next(t0) # advance index of t0 by 1 before __copy__()
t0c = t0.__copy__()
t1c = t1.__copy__()
assert list(t0) == [2, 3]
assert list(t0c) == [2, 3]
assert list(t1) == [1, 2, 3]
assert list(t1c) == [1, 2, 3]
t0, t1 = itertools.tee([1, 2, 3])
t2, t3 = itertools.tee(t0)
assert list(t1) == [1, 2, 3]
assert list(t2) == [1, 2, 3]
assert list(t3) == [1, 2, 3]
t = itertools.tee([1, 2, 3])
assert list(t[0]) == [1, 2, 3]
assert list(t[0]) == []
# itertools.product
it = itertools.product([1, 2], [3, 4])
assert (1, 3) == next(it)
assert (1, 4) == next(it)
assert (2, 3) == next(it)
assert (2, 4) == next(it)
with assert_raises(StopIteration):
next(it)
it = itertools.product([1, 2], repeat=2)
assert (1, 1) == next(it)
assert (1, 2) == next(it)
assert (2, 1) == next(it)
assert (2, 2) == next(it)
with assert_raises(StopIteration):
next(it)
with assert_raises(TypeError):
itertools.product(None)
with assert_raises(TypeError):
itertools.product([1, 2], repeat=None)
# itertools.combinations
it = itertools.combinations([1, 2, 3, 4], 2)
assert list(it) == [(1, 2), (1, 3), (1, 4), (2, 3), (2, 4), (3, 4)]
it = itertools.combinations([1, 2, 3], 0)
assert list(it) == [()]
it = itertools.combinations([1, 2, 3], 1)
assert list(it) == [(1,), (2,), (3,)]
it = itertools.combinations([1, 2, 3], 2)
assert next(it) == (1, 2)
assert next(it) == (1, 3)
assert next(it) == (2, 3)
with assert_raises(StopIteration):
next(it)
it = itertools.combinations([1, 2, 3], 4)
with assert_raises(StopIteration):
next(it)
with assert_raises(ValueError):
itertools.combinations([1, 2, 3, 4], -2)
with assert_raises(TypeError):
itertools.combinations([1, 2, 3, 4], None)
# itertools.combinations
it = itertools.combinations_with_replacement([1, 2, 3], 0)
assert list(it) == [()]
it = itertools.combinations_with_replacement([1, 2, 3], 1)
assert list(it) == [(1,), (2,), (3,)]
it = itertools.combinations_with_replacement([1, 2, 3], 2)
assert list(it) == [(1, 1), (1, 2), (1, 3), (2, 2), (2, 3), (3, 3)]
it = itertools.combinations_with_replacement([1, 2], 3)
assert list(it) == [(1, 1, 1), (1, 1, 2), (1, 2, 2), (2, 2, 2)]
with assert_raises(ValueError):
itertools.combinations_with_replacement([1, 2, 3, 4], -2)
with assert_raises(TypeError):
itertools.combinations_with_replacement([1, 2, 3, 4], None)
# itertools.permutations
it = itertools.permutations([1, 2, 3])
assert list(it) == [(1, 2, 3), (1, 3, 2), (2, 1, 3), (2, 3, 1), (3, 1, 2), (3, 2, 1)]
it = itertools.permutations([1, 2, 3], None)
assert list(it) == [(1, 2, 3), (1, 3, 2), (2, 1, 3), (2, 3, 1), (3, 1, 2), (3, 2, 1)]
it = itertools.permutations([1, 2, 3], 0)
assert list(it) == [()]
it = itertools.permutations([1, 2, 3], 1)
assert list(it) == [(1,), (2,), (3,)]
it = itertools.permutations([1, 2, 3], 2)
assert next(it) == (1, 2)
assert next(it) == (1, 3)
assert next(it) == (2, 1)
assert next(it) == (2, 3)
assert next(it) == (3, 1)
assert next(it) == (3, 2)
with assert_raises(StopIteration):
next(it)
it = itertools.permutations([1, 2, 3], 4)
with assert_raises(StopIteration):
next(it)
with assert_raises(ValueError):
itertools.permutations([1, 2, 3, 4], -2)
with assert_raises(TypeError):
itertools.permutations([1, 2, 3, 4], 3.14)
# itertools.zip_longest tests
zl = itertools.zip_longest
assert list(zl(["a", "b", "c"], range(3), [9, 8, 7])) == [
("a", 0, 9),
("b", 1, 8),
("c", 2, 7),
]
assert list(zl(["a", "b", "c"], range(3), [9, 8, 7, 99])) == [
("a", 0, 9),
("b", 1, 8),
("c", 2, 7),
(None, None, 99),
]
assert list(zl(["a", "b", "c"], range(3), [9, 8, 7, 99], fillvalue="d")) == [
("a", 0, 9),
("b", 1, 8),
("c", 2, 7),
("d", "d", 99),
]
assert list(zl(["a", "b", "c"])) == [("a",), ("b",), ("c",)]
assert list(zl()) == []
assert list(zl(*zl(["a", "b", "c"], range(1, 4)))) == [("a", "b", "c"), (1, 2, 3)]
assert list(zl(*zl(["a", "b", "c"], range(1, 5)))) == [
("a", "b", "c", None),
(1, 2, 3, 4),
]
assert list(zl(*zl(["a", "b", "c"], range(1, 5), fillvalue=100))) == [
("a", "b", "c", 100),
(1, 2, 3, 4),
]
# test infinite iterator
class Counter(object):
def __init__(self, counter=0):
self.counter = counter
def __next__(self):
self.counter += 1
return self.counter
def __iter__(self):
return self
it = zl(Counter(), Counter(3))
assert next(it) == (1, 4)
assert next(it) == (2, 5)
it = zl([1, 2], [3])
assert next(it) == (1, 3)
assert next(it) == (2, None)
with assert_raises(StopIteration):
next(it)
# r is an arbitrary Python int: one too large for an index must raise
# OverflowError, and a representable one that cannot be allocated must raise
# MemoryError.
for factory in (
itertools.combinations,
itertools.combinations_with_replacement,
itertools.permutations,
):
with assert_raises(OverflowError):
factory(range(5), 2**64)
with assert_raises(MemoryError):
itertools.combinations(range(5), 2**44)
with assert_raises(MemoryError):
itertools.combinations_with_replacement(range(5), 2**44)