-
Notifications
You must be signed in to change notification settings - Fork 0
/
interleave.py
57 lines (45 loc) · 1.55 KB
/
interleave.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
from itertools import zip_longest
def interleave2(*iterables):
iters = [iter(x) for x in iterables]
while iters:
for it in list(iters):
try:
yield next(it)
except StopIteration:
iters.remove(it)
def interleave(*iterables):
dummy = object()
for items in zip_longest(*iterables, fillvalue=dummy):
for item in items:
if item is not dummy:
yield item
def test_numbers_basic():
numbers = [1, 2, 3, 4]
res = interleave(numbers, range(5, 9))
assert list(res) == [1, 5, 2, 6, 3, 7, 4, 8]
def test_numbers_basic2():
numbers = [1, 2, 3, 4]
res = interleave(numbers, (n**2 for n in numbers))
assert list(res) == [1, 1, 2, 4, 3, 9, 4, 16]
def test_numbers_iterator():
i = interleave([1,2,3,4], [5,6,7,8])
assert next(i) == 1
assert list(i) == [5, 2, 6, 3, 7, 4, 8]
def test_three_args():
expected = [1, 4, 7, 2, 5, 8, 3, 6, 9]
actual = interleave([1,2,3], [4,5,6], [7,8,9])
assert list(actual) == expected
def test_unequal_lengths():
actual = interleave([1, 2, 3], [4, 5, 6, 7, 8])
expected = [1, 4, 2, 5, 3, 6, 7, 8]
assert list(actual) == expected
def test_different_length_lists():
in1 = [1, 2, 3]
in2 = [4, 5, 6, 7, 8]
in3 = [9]
out1 = [1, 4, 9, 2, 5, 3, 6, 7, 8]
assert list(interleave(in1, in2, in3)) == out1
def test_different_length_lists2():
actual = interleave([1, 2], [3], [4, 5, 6], [7, 8], [9])
expected = [1, 3, 4, 7, 9, 2, 5, 8, 6]
assert list(actual) == expected