micropython-lib/python-stdlib/itertools/test_itertools.py

238 wiersze
6.8 KiB
Python

import itertools
# accumulate
assert list(itertools.accumulate([])) == []
assert list(itertools.accumulate([0])) == [0]
assert list(itertools.accumulate([0, 2, 3])) == [0, 2, 5]
assert list(itertools.accumulate(reversed([0, 2, 3]))) == [3, 5, 5]
assert list(itertools.accumulate([1, 2, 3], lambda x, y: x * y)) == [1, 2, 6]
assert list(itertools.accumulate([1, 2, 3, 4, 5], func=lambda x, y: x - y, initial=10)) == [
10,
9,
7,
4,
0,
-5,
]
# chain
assert list(itertools.chain()) == []
assert list(itertools.chain([], [], [])) == []
assert list(itertools.chain(range(3), [2 * (i + 1) for i in range(4)])) == [0, 1, 2, 2, 4, 6, 8]
assert list(itertools.chain("abcd", [], range(5))) == ["a", "b", "c", "d", 0, 1, 2, 3, 4]
assert list(itertools.chain.from_iterable([])) == []
assert list(itertools.chain.from_iterable(["ABC", "DEF"])) == ["A", "B", "C", "D", "E", "F"]
# combinations
assert list(itertools.combinations("", 1)) == []
assert list(itertools.combinations("ABCD", 0)) == [()]
assert list(itertools.combinations("ABCD", 1)) == [("A",), ("B",), ("C",), ("D",)]
assert list(itertools.combinations("ABCD", 3)) == [
("A", "B", "C"),
("A", "B", "D"),
("A", "C", "D"),
("B", "C", "D"),
]
assert list(itertools.combinations("ABCD", 4)) == [("A", "B", "C", "D")]
assert list(itertools.combinations("ABCD", 5)) == []
assert len(list(itertools.combinations(range(7), 4))) == 35
assert len(set(itertools.combinations(range(7), 4))) == 35
# combinations with replacement
assert list(itertools.combinations_with_replacement("ABCD", 0)) == [()]
assert list(itertools.combinations_with_replacement("ABCD", 1)) == [("A",), ("B",), ("C",), ("D",)]
assert list(itertools.combinations_with_replacement("ABC", 2)) == [
("A", "A"),
("A", "B"),
("A", "C"),
("B", "B"),
("B", "C"),
("C", "C"),
]
assert list(itertools.combinations_with_replacement("ABC", 3)) == [
("A", "A", "A"),
("A", "A", "B"),
("A", "A", "C"),
("A", "B", "B"),
("A", "B", "C"),
("A", "C", "C"),
("B", "B", "B"),
("B", "B", "C"),
("B", "C", "C"),
("C", "C", "C"),
]
# compress
assert tuple(itertools.compress("ABCDEF", (1, 0, 1, 0, 1, 1))) == ("A", "C", "E", "F")
assert tuple(itertools.compress("ABCDEF", (1, 0, 1, 1, 0))) == ("A", "C", "D")
# count
it = itertools.count(4, 3)
for _ in range(200):
n = next(it)
assert list(next(it) for _ in range(5)) == [604, 607, 610, 613, 616]
# cycle
it = itertools.cycle(iter("abcde"))
assert list(next(it) for _ in range(12)) == [
"a",
"b",
"c",
"d",
"e",
"a",
"b",
"c",
"d",
"e",
"a",
"b",
]
it = itertools.cycle([2, 4, "x"])
assert list(next(it) for _ in range(7)) == [2, 4, "x", 2, 4, "x", 2]
# dropwhile
assert list(itertools.dropwhile(lambda x: x < 5, [1, 4, 6, 4, 1])) == [6, 4, 1]
assert list(itertools.dropwhile(lambda x: ord(x) < 118, "")) == []
assert list(itertools.dropwhile(lambda x: ord(x) < 118, "dropwhile")) == ["w", "h", "i", "l", "e"]
# filterfalse
assert list(itertools.filterfalse(lambda x: ord(x) < 110, "dropwhile")) == ["r", "o", "p", "w"]
# groupby
assert list(((k, "".join(g)) for k, g in itertools.groupby("aaaabbbccdaa"))) == [
("a", "aaaa"),
("b", "bbb"),
("c", "cc"),
("d", "d"),
("a", "aa"),
]
# islice
assert "".join(itertools.islice("", 2, 5)) == ""
assert "".join(itertools.islice("abcdefgh", 2, 5)) == "cde"
assert "".join(itertools.islice("abcdefghij", 2, None, 3)) == "cfi"
assert "".join(itertools.islice("abcdefghij", 2, None)) == "cdefghij"
assert "".join(itertools.islice("abcdefghij", 6)) == "abcdef"
assert "".join(itertools.islice("abcdefghij", 6, 6)) == ""
assert list(itertools.islice(range(10), 2, 6, 2)) == [2, 4]
assert list(itertools.islice(itertools.cycle([1, 2, 3]), 10)) == [1, 2, 3, 1, 2, 3, 1, 2, 3, 1]
# pairwise
assert list(itertools.pairwise(range(5))) == [(0, 1), (1, 2), (2, 3), (3, 4)]
assert list(("".join(t) for t in itertools.pairwise("abcdefg"))) == [
"ab",
"bc",
"cd",
"de",
"ef",
"fg",
]
assert list(("".join(t) for t in itertools.pairwise("ab"))) == ["ab"]
assert list(("".join(t) for t in itertools.pairwise("a"))) == []
assert list(("".join(t) for t in itertools.pairwise(""))) == []
# permutations
assert list(itertools.permutations("", 1)) == []
assert list(("".join(t) for t in itertools.permutations("a", 2))) == []
assert list(("".join(t) for t in itertools.permutations("abcd", 0))) == [""]
assert list(("".join(t) for t in itertools.permutations("ab", 2))) == ["ab", "ba"]
assert list(("".join(t) for t in itertools.permutations("abcd", 2))) == [
"ab",
"ac",
"ad",
"ba",
"bc",
"bd",
"ca",
"cb",
"cd",
"da",
"db",
"dc",
]
assert list(itertools.permutations(range(3))) == [
(0, 1, 2),
(0, 2, 1),
(1, 0, 2),
(1, 2, 0),
(2, 0, 1),
(2, 1, 0),
]
# product
assert list(itertools.product()) == [()]
assert list(itertools.product(range(2), repeat=0)) == [()]
assert list(("".join(t) for t in itertools.product("ABC", "xy"))) == [
"Ax",
"Ay",
"Bx",
"By",
"Cx",
"Cy",
]
assert list(("".join(t) for t in itertools.product("A", "xy", repeat=2))) == [
"AxAx",
"AxAy",
"AyAx",
"AyAy",
]
assert list(("".join(map(str, t)) for t in itertools.product(range(2), repeat=3))) == [
"000",
"001",
"010",
"011",
"100",
"101",
"110",
"111",
]
# repeat
assert list(itertools.repeat(10, 0)) == []
assert list(itertools.repeat(10, 1)) == [10]
assert list(itertools.repeat(10, 3)) == [10, 10, 10]
# starmap
assert list(itertools.starmap(pow, [])) == []
assert list(itertools.starmap(pow, [(2, 5), (3, 2), (10, 3)])) == [32, 9, 1000]
assert list(itertools.starmap(lambda x, y: x * y, [[1, 2], [2, 3], [3, 4]])) == [2, 6, 12]
# takewhile
assert list(itertools.takewhile(lambda x: x < 5, [1, 4, 6, 4, 1])) == [1, 4]
assert list(itertools.takewhile(lambda x: ord(x) < 118, "dropwhile")) == ["d", "r", "o", "p"]
# tee
def genx(n):
i = 1
while True:
yield i
i += 1
if i > n:
return
it1, it2, it3 = itertools.tee(genx(1000), 3)
_ = next(it1) # case of iterable that is unique; iterate once
assert [min(it1), max(it2), sum(it3)] == [2, 1000, 500500]
it1, it2, it3 = itertools.tee(range(2, 10), 3)
_ = next(it1) # iterable that is not unique; iterate once
assert [min(it1), max(it2), sum(it3)] == [
3,
9,
44,
] # the min is increased, other iterators remained full
# zip_longest
assert list(itertools.zip_longest("", "")) == []
assert list(itertools.zip_longest("", "", fillvalue="-")) == []
assert list(itertools.zip_longest("", "xy")) == [(None, "x"), (None, "y")]
assert list(itertools.zip_longest("", "xy", fillvalue="-")) == [("-", "x"), ("-", "y")]
assert list(itertools.zip_longest("ABCD", "xy", fillvalue="-")) == [
("A", "x"),
("B", "y"),
("C", "-"),
("D", "-"),
]