github.com/google/grumpy@v0.0.0-20171122020858-3ec87959189c/lib/itertools_test.py (about) 1 # Copyright 2016 Google Inc. All Rights Reserved. 2 # 3 # Licensed under the Apache License, Version 2.0 (the "License"); 4 # you may not use this file except in compliance with the License. 5 # You may obtain a copy of the License at 6 # 7 # http://www.apache.org/licenses/LICENSE-2.0 8 # 9 # Unless required by applicable law or agreed to in writing, software 10 # distributed under the License is distributed on an "AS IS" BASIS, 11 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 # See the License for the specific language governing permissions and 13 # limitations under the License. 14 15 import itertools 16 17 import weetest 18 19 def TestCycle(): 20 want = [] 21 got = [] 22 for x in itertools.cycle(()): 23 got.append(x) 24 assert got == want, 'empty cycle yields no elements' 25 26 arg = (0, 1, 2) 27 want = (0, 1, 2) * 10 28 got = [] 29 limit = 10 * len(arg) 30 counter = 0 31 for x in itertools.cycle((0, 1, 2)): 32 got.append(x) 33 counter += 1 34 if counter == limit: 35 break 36 assert tuple(got) == want, 'tuple(cycle%s) == %s, want %s' % (arg, tuple(got), want) 37 38 39 def TestDropwhile(): 40 r = range(10) 41 cases = [ 42 ((lambda x: x < 5, r), (5, 6, 7, 8, 9)), 43 ((lambda x: True, r), ()), 44 ((lambda x: False, r), tuple(r)), 45 ] 46 for args, want in cases: 47 got = tuple(itertools.dropwhile(*args)) 48 assert got == want, 'tuple(dropwhile%s) == %s, want %s' % (args, got, want) 49 50 51 def TestChain(): 52 r = range(10) 53 cases = [ 54 ([r], tuple(r)), 55 ([r, r], tuple(r) + tuple(r)), 56 ([], ()) 57 ] 58 for args, want in cases: 59 got = tuple(itertools.chain(*args)) 60 assert got == want, 'tuple(chain%s) == %s, want %s' % (args, got, want) 61 62 63 def TestFromIterable(): 64 r = range(10) 65 cases = [ 66 ([r], tuple(r)), 67 ([r, r], tuple(r) + tuple(r)), 68 ([], ()) 69 ] 70 for args, want in cases: 71 got = tuple(itertools.chain.from_iterable(args)) 72 assert got == want, 'tuple(from_iterable%s) == %s, want %s' % (args, got, want) 73 74 75 def TestIFilter(): 76 r = range(10) 77 cases = [ 78 ((lambda x: x < 5, r), (0, 1, 2, 3, 4)), 79 ((lambda x: False, r), ()), 80 ((lambda x: True, r), tuple(r)), 81 ((None, r), (1, 2, 3, 4, 5, 6, 7, 8, 9)) 82 ] 83 for args, want in cases: 84 got = tuple(itertools.ifilter(*args)) 85 assert got == want, 'tuple(ifilter%s) == %s, want %s' % (args, got, want) 86 87 88 def TestIFilterFalse(): 89 r = range(10) 90 cases = [ 91 ((lambda x: x < 5, r), (5, 6, 7, 8, 9)), 92 ((lambda x: False, r), tuple(r)), 93 ((lambda x: True, r), ()), 94 ((None, r), (0,)) 95 ] 96 for args, want in cases: 97 got = tuple(itertools.ifilterfalse(*args)) 98 assert got == want, 'tuple(ifilterfalse%s) == %s, want %s' % (args, got, want) 99 100 101 def TestISlice(): 102 r = range(10) 103 cases = [ 104 ((r, 5), (0, 1, 2, 3, 4)), 105 ((r, 25, 30), ()), 106 ((r, 1, None, 3), (1, 4, 7)), 107 ] 108 for args, want in cases: 109 got = tuple(itertools.islice(*args)) 110 assert got == want, 'tuple(islice%s) == %s, want %s' % (args, got, want) 111 112 113 def TestIZipLongest(): 114 cases = [ 115 (('abc', range(6)), (('a', 0), ('b', 1), ('c', 2), (None, 3), (None, 4), (None, 5))), 116 ((range(6), 'abc'), ((0, 'a'), (1, 'b'), (2, 'c'), (3, None), (4, None), (5, None))), 117 (([1, None, 3], 'ab', range(1)), ((1, 'a', 0), (None, 'b', None), (3, None, None))), 118 ] 119 for args, want in cases: 120 got = tuple(itertools.izip_longest(*args)) 121 assert got == want, 'tuple(izip_longest%s) == %s, want %s' % (args, got, want) 122 123 124 def TestProduct(): 125 cases = [ 126 (([1, 2], ['a', 'b']), ((1, 'a'), (1, 'b'), (2, 'a'), (2, 'b'))), 127 (([1], ['a', 'b']), ((1, 'a'), (1, 'b'))), 128 (([],), ()), 129 ] 130 for args, want in cases: 131 got = tuple(itertools.product(*args)) 132 assert got == want, 'tuple(product%s) == %s, want %s' % (args, got, want) 133 134 135 def TestPermutations(): 136 cases = [ 137 (('AB',), (('A', 'B'), ('B', 'A'))), 138 (('ABC', 2), (('A', 'B'), ('A', 'C'), ('B', 'A'), ('B', 'C'), ('C', 'A'), ('C', 'B'))), 139 ((range(3),), ((0, 1, 2), (0, 2, 1), (1, 0, 2), (1, 2, 0), (2, 0, 1), (2, 1, 0))), 140 (([],), ((),)), 141 (([], 0), ((),)), 142 ((range(3), 4), ()), 143 ] 144 for args, want in cases: 145 got = tuple(itertools.permutations(*args)) 146 assert got == want, 'tuple(permutations%s) == %s, want %s' % (args, got, want) 147 148 149 def TestCombinations(): 150 cases = [ 151 ((range(4), 3), ((0, 1, 2), (0, 1, 3), (0, 2, 3), (1, 2, 3))), 152 ] 153 for args, want in cases: 154 got = tuple(itertools.combinations(*args)) 155 assert got == want, 'tuple(combinations%s) == %s, want %s' % (args, got, want) 156 157 158 def TestCombinationsWithReplacement(): 159 cases = [ 160 (([-12], 2), (((-12, -12),))), 161 (('AB', 3), (('A', 'A', 'A'), ('A', 'A', 'B'), ('A', 'B', 'B'), ('B', 'B', 'B'))), 162 (([], 2), ()), 163 (([], 0), ((),)) 164 ] 165 for args, want in cases: 166 got = tuple(itertools.combinations_with_replacement(*args)) 167 assert got == want, 'tuple(combinations_with_replacement%s) == %s, want %s' % (args, got, want) 168 169 170 def TestGroupBy(): 171 cases = [ 172 (([1, 2, 2, 3, 3, 3, 4, 4, 4, 4],), [(1, [1]), (2, [2, 2]), (3, [3, 3, 3]), (4, [4, 4, 4, 4])]), 173 ((['aa', 'ab', 'abc', 'bcd', 'abcde'], len), [(2, ['aa', 'ab']), (3, ['abc', 'bcd']), (5, ['abcde'])]), 174 ] 175 for args, want in cases: 176 got = [(k, list(v)) for k, v in itertools.groupby(*args)] 177 assert got == want, 'groupby %s == %s, want %s' % (args, got, want) 178 179 180 def TestTakewhile(): 181 r = range(10) 182 cases = [ 183 ((lambda x: x % 2 == 0, r), (0,)), 184 ((lambda x: True, r), tuple(r)), 185 ((lambda x: False, r), ()) 186 ] 187 for args, want in cases: 188 got = tuple(itertools.takewhile(*args)) 189 assert got == want, 'tuple(takewhile%s) == %s, want %s' % (args, got, want) 190 191 192 if __name__ == '__main__': 193 weetest.RunTests()