Files
2026-10-23 18:03:11 +09:00

706 lines
17 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
product_source = iter([1])
assert list(itertools.product(product_source, repeat=0)) == [()]
assert next(product_source) == 1
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)
# repeat is an arbitrary Python int: a negative one is refused, and one whose
# pool cannot be allocated must raise rather than take the process down.
with assert_raises(ValueError):
itertools.product([1], repeat=-1)
with assert_raises(OverflowError):
itertools.product([1, 2], repeat=2**60)
# The pools are filled by their own count, so a repeat with nothing to repeat
# answers at once instead of counting up to it.
assert list(itertools.product(repeat=2**62)) == [()]
assert list(itertools.product(repeat=0)) == [()]
# The count is settled before the arguments are read, so a repeat too large to
# serve does not run their code first.
ran = []
class Watched:
def __iter__(self):
ran.append(True)
return iter([1])
with assert_raises(OverflowError):
itertools.product(Watched(), repeat=2**62)
assert ran == []
# compress consumes data, then the selector, then tests the selector's truth.
def check_compress_consumption():
events = []
def tracked(name, values):
for value in values:
events.append(name)
yield value
events.append(name + " exhausted")
class Selector:
def __init__(self, value):
self.value = value
def __bool__(self):
events.append("bool")
return self.value
assert list(
itertools.compress(
tracked("data", [1, 2]),
tracked("selectors", [Selector(False), Selector(True), Selector(True)]),
)
) == [2]
assert events == [
"data",
"selectors",
"bool",
"data",
"selectors",
"bool",
"data exhausted",
], events
events.clear()
data = tracked("data", [1, 2, 3])
assert list(itertools.compress(data, tracked("selectors", [True]))) == [1]
assert events == ["data", "selectors", "data", "selectors exhausted"], events
assert next(data) == 3
events.clear()
assert (
list(itertools.compress(tracked("data", []), tracked("selectors", [True])))
== []
)
assert events == ["data exhausted"], events
def failing(name):
events.append(name)
raise ValueError(name)
yield
events.clear()
with assert_raises(ValueError) as caught:
next(itertools.compress(failing("data"), tracked("selectors", [True])))
assert str(caught.exception) == "data"
assert events == ["data"], events
events.clear()
with assert_raises(ValueError) as caught:
next(itertools.compress(tracked("data", [1]), failing("selectors")))
assert str(caught.exception) == "selectors"
assert events == ["data", "selectors"], events
class BadSelector:
def __bool__(self):
events.append("bool")
raise ValueError("truth")
events.clear()
with assert_raises(ValueError) as caught:
next(itertools.compress(tracked("data", [1]), [BadSelector()]))
assert str(caught.exception) == "truth"
assert events == ["data", "bool"], events
check_compress_consumption()
# These public iterator types also support Python subclasses.
for base, args in (
(itertools.count, ()),
(itertools.groupby, ([1],)),
(itertools.accumulate, ([1],)),
(itertools.product, ([1],)),
(itertools.combinations, ([1], 1)),
(itertools.combinations_with_replacement, ([1], 1)),
(itertools.permutations, ([1],)),
(itertools.zip_longest, ([1],)),
(itertools.pairwise, ([1, 2],)),
):
class Subclass(base):
pass
instance = Subclass(*args)
assert type(instance) is Subclass
next(instance)
class PairwiseWithLabel(itertools.pairwise):
def __init__(self, iterable, *, label):
self.label = label
pairs = PairwiseWithLabel([1, 2], label="pairs")
assert pairs.label == "pairs"
assert next(pairs) == (1, 2)
# Consumed tee values must not remain alive for the lifetime of the iterators.
import weakref
class TeeItem:
pass
tee_left, tee_right = itertools.tee(TeeItem() for _ in range(200))
tee_first = next(tee_left)
tee_ref = weakref.ref(tee_first)
del tee_first
next(tee_right)
for _ in range(199):
next(tee_left)
next(tee_right)
assert tee_ref() is None