bernsteinbear.com
this repo has no description
1#!/usr/bin/env python3
2# Copyright (c) Facebook, Inc. and its affiliates. (http://www.facebook.com)
3import itertools
4import operator
5import unittest
6from unittest.mock import Mock
7
8
9class ChainTests(unittest.TestCase):
10 def test_chain_with_no_iterables_returns_stoped_iterator(self):
11 self.assertTupleEqual(tuple(itertools.chain()), ())
12
13 def test_chain_with_empty_iterables_returns_stoped_iterator(self):
14 self.assertTupleEqual(tuple(itertools.chain([])), ())
15 self.assertTupleEqual(tuple(itertools.chain([], [])), ())
16 self.assertTupleEqual(tuple(itertools.chain([], [], [])), ())
17
18 def test_chain_with_one_iterable(self):
19 self.assertTupleEqual(tuple(itertools.chain("1")), ("1",))
20 self.assertTupleEqual(tuple(itertools.chain("12")), ("1", "2"))
21 self.assertTupleEqual(tuple(itertools.chain("123")), ("1", "2", "3"))
22
23 def test_chain_with_many_iterables(self):
24 self.assertTupleEqual(tuple(itertools.chain("1", "")), ("1",))
25 self.assertTupleEqual(tuple(itertools.chain("1", "2")), ("1", "2"))
26 self.assertTupleEqual(tuple(itertools.chain("12", "")), ("1", "2"))
27 self.assertTupleEqual(tuple(itertools.chain("12", "3")), ("1", "2", "3"))
28 self.assertTupleEqual(tuple(itertools.chain("12", "34")), ("1", "2", "3", "4"))
29 self.assertTupleEqual(tuple(itertools.chain("", "1")), ("1",))
30 self.assertTupleEqual(tuple(itertools.chain("", "1", "")), ("1",))
31 self.assertTupleEqual(tuple(itertools.chain("1", "23")), ("1", "2", "3"))
32 self.assertTupleEqual(tuple(itertools.chain("1", "23", "")), ("1", "2", "3"))
33 self.assertTupleEqual(tuple(itertools.chain("1", "2", "3")), ("1", "2", "3"))
34
35 def test_from_iterable(self):
36 self.assertTupleEqual(tuple(itertools.chain.from_iterable([])), ())
37 self.assertTupleEqual(tuple(itertools.chain.from_iterable([""])), ())
38 self.assertTupleEqual(tuple(itertools.chain.from_iterable(["1"])), ("1",))
39 self.assertTupleEqual(tuple(itertools.chain.from_iterable(["12"])), ("1", "2"))
40 self.assertTupleEqual(
41 tuple(itertools.chain.from_iterable(["12", ""])), ("1", "2")
42 )
43 self.assertTupleEqual(
44 tuple(itertools.chain.from_iterable(["12", "3"])), ("1", "2", "3")
45 )
46 self.assertTupleEqual(
47 tuple(itertools.chain.from_iterable(["12", "34"])), ("1", "2", "3", "4")
48 )
49 self.assertTupleEqual(
50 tuple(itertools.chain.from_iterable(["12", "34", ""])), ("1", "2", "3", "4")
51 )
52 self.assertTupleEqual(
53 tuple(itertools.chain.from_iterable(["12", "34", "5"])),
54 ("1", "2", "3", "4", "5"),
55 )
56 self.assertTupleEqual(
57 tuple(itertools.chain.from_iterable(["12", "34", "56"])),
58 ("1", "2", "3", "4", "5", "6"),
59 )
60
61
62class CycleTests(unittest.TestCase):
63 def test_dunder_init_with_seq_does_calls_dunder_iter(self):
64 class C:
65 __iter__ = Mock(name="__iter__", return_value=[].__iter__())
66
67 instance = C()
68 itertools.cycle(instance)
69 C.__iter__.assert_called_once()
70
71 def test_dunder_init_with_seq_dunder_iter_returning_non_iter_raises_type_error(
72 self,
73 ):
74 class C:
75 def __iter__(self):
76 return 5
77
78 instance = C()
79 with self.assertRaisesRegex(
80 TypeError, "iter\\(\\) returned non-iterator of type 'int'"
81 ):
82 itertools.cycle(instance)
83
84 def test_dunder_iter_yields_self(self):
85 result = itertools.cycle([])
86 self.assertIs(result.__iter__(), result)
87
88 def test_dunder_next_iterates_through_sequence(self):
89 result = itertools.cycle([1, 2, 3])
90 self.assertEqual(next(result), 1)
91 self.assertEqual(next(result), 2)
92 self.assertEqual(next(result), 3)
93
94 def test_dunder_next_cycles_through_sequence(self):
95 result = itertools.cycle([1, 2, 3])
96 self.assertEqual(next(result), 1)
97 self.assertEqual(next(result), 2)
98 self.assertEqual(next(result), 3)
99 self.assertEqual(next(result), 1)
100 self.assertEqual(next(result), 2)
101 self.assertEqual(next(result), 3)
102 self.assertEqual(next(result), 1)
103
104
105class CountTests(unittest.TestCase):
106 def test_count_with_int_returns_iterator(self):
107 iterator = itertools.count(start=7, step=-2)
108 list = [next(iterator) for _ in range(5)]
109 self.assertEqual(list, [7, 5, 3, 1, -1])
110
111 def test_count_with_non_number_raises_type_error(self):
112 with self.assertRaisesRegex(TypeError, "a number is required"):
113 itertools.count(start="a", step=".")
114
115
116class IsliceTests(unittest.TestCase):
117 def test_too_few_arguments_raises_type_error(self):
118 self.assertRaises(TypeError, itertools.islice, [])
119
120 def test_too_many_arguments_raises_type_error(self):
121 self.assertRaises(TypeError, itertools.islice, [], 1, 2, 3, 4)
122
123 def test_single_slice_non_int_arg_raises_value_error(self):
124 with self.assertRaises(ValueError) as ctx:
125 itertools.islice([], [])
126 self.assertEqual(
127 str(ctx.exception),
128 "Stop argument for islice() must be None or an integer: "
129 "0 <= x <= sys.maxsize.",
130 )
131
132 def test_single_slice_neg_arg_raises_value_error(self):
133 with self.assertRaises(ValueError) as ctx:
134 itertools.islice([], -1)
135 self.assertEqual(
136 str(ctx.exception),
137 "Stop argument for islice() must be None or an integer: "
138 "0 <= x <= sys.maxsize.",
139 )
140
141 def test_slice_non_int_stop_raises_value_error(self):
142 with self.assertRaises(ValueError) as ctx:
143 itertools.islice([], 0, [])
144 self.assertEqual(
145 str(ctx.exception),
146 "Stop argument for islice() must be None or an integer: "
147 "0 <= x <= sys.maxsize.",
148 )
149
150 def test_slice_neg_one_stop_raises_value_error(self):
151 with self.assertRaises(ValueError) as ctx:
152 itertools.islice([], 0, -1)
153 self.assertEqual(
154 str(ctx.exception),
155 "Stop argument for islice() must be None or an integer: "
156 "0 <= x <= sys.maxsize.",
157 )
158
159 def test_slice_neg_two_stop_raises_value_error(self):
160 with self.assertRaises(ValueError) as ctx:
161 itertools.islice([], 0, -2)
162 self.assertEqual(
163 str(ctx.exception),
164 "Indices for islice() must be None or an integer: 0 <= x <= sys.maxsize.",
165 )
166
167 def test_slice_non_int_start_raises_value_error(self):
168 with self.assertRaises(ValueError) as ctx:
169 itertools.islice([], [], 1)
170 self.assertEqual(
171 str(ctx.exception),
172 "Indices for islice() must be None or an integer: 0 <= x <= sys.maxsize.",
173 )
174
175 def test_slice_neg_int_start_raises_value_error(self):
176 with self.assertRaises(ValueError) as ctx:
177 itertools.islice([], -1, 1)
178 self.assertEqual(
179 str(ctx.exception),
180 "Indices for islice() must be None or an integer: 0 <= x <= sys.maxsize.",
181 )
182
183 def test_slice_non_int_step_raises_value_error(self):
184 with self.assertRaises(ValueError) as ctx:
185 itertools.islice([], 0, 1, [])
186 self.assertEqual(
187 str(ctx.exception), "Step for islice() must be a positive integer or None."
188 )
189
190 def test_slice_neg_int_step_raises_value_error(self):
191 with self.assertRaises(ValueError) as ctx:
192 itertools.islice([], 0, 1, -1)
193 self.assertEqual(
194 str(ctx.exception), "Step for islice() must be a positive integer or None."
195 )
196
197 def test_slice_with_none_stop_treats_stop_as_end(self):
198 islice = itertools.islice([0, 1, 2, 3], None)
199 self.assertEqual(next(islice), 0)
200 self.assertEqual(next(islice), 1)
201 self.assertEqual(next(islice), 2)
202 self.assertEqual(next(islice), 3)
203 with self.assertRaises(StopIteration):
204 next(islice)
205
206 def test_slice_with_zero_stop_immediately_stops(self):
207 islice = itertools.islice([0, 1, 2, 3], 0)
208 with self.assertRaises(StopIteration):
209 next(islice)
210
211 def test_slice_with_none_start_starts_at_zero(self):
212 islice = itertools.islice([0, 1, 2, 3], None, 2)
213 self.assertEqual(next(islice), 0)
214 self.assertEqual(next(islice), 1)
215 with self.assertRaises(StopIteration):
216 next(islice)
217
218 def test_slice_with_none_step_steps_by_one(self):
219 islice = itertools.islice([0, 1, 2, 3], 1, 3, None)
220 self.assertEqual(next(islice), 1)
221 self.assertEqual(next(islice), 2)
222 with self.assertRaises(StopIteration):
223 next(islice)
224
225 def test_slice_with_all_none_args_acts_as_the_given_iterable(self):
226 islice = itertools.islice([0, 1, 2, 3], None, None, None)
227 self.assertEqual(next(islice), 0)
228 self.assertEqual(next(islice), 1)
229 self.assertEqual(next(islice), 2)
230 self.assertEqual(next(islice), 3)
231 with self.assertRaises(StopIteration):
232 next(islice)
233
234 def test_slice_with_one_arg_stops_iterable_at_stop(self):
235 islice = itertools.islice([0, 1, 2, 3], 2)
236 self.assertEqual(next(islice), 0)
237 self.assertEqual(next(islice), 1)
238 with self.assertRaises(StopIteration):
239 next(islice)
240
241 def test_slice_with_one_arg_stops_at_end_of_iterable(self):
242 islice = itertools.islice([0, 1, 2], 4)
243 self.assertEqual(next(islice), 0)
244 self.assertEqual(next(islice), 1)
245 self.assertEqual(next(islice), 2)
246 with self.assertRaises(StopIteration):
247 next(islice)
248
249 def test_slice_with_start_and_stop_respects_slice_bounds(self):
250 islice = itertools.islice([0, 1, 2, 3], 1, 3)
251 self.assertEqual(next(islice), 1)
252 self.assertEqual(next(islice), 2)
253 with self.assertRaises(StopIteration):
254 next(islice)
255
256 def test_slice_with_step_arg_respects_slice(self):
257 islice = itertools.islice([0, 1, 2, 3, 4, 5, 6, 7, 8], 1, 6, 2)
258 self.assertEqual(next(islice), 1)
259 self.assertEqual(next(islice), 3)
260 self.assertEqual(next(islice), 5)
261 with self.assertRaises(StopIteration):
262 next(islice)
263
264 def test_slice_with_step_arg_ends_at_end_of_iterator(self):
265 islice = itertools.islice([0, 1, 2, 3, 4], 1, 6, 2)
266 self.assertEqual(next(islice), 1)
267 self.assertEqual(next(islice), 3)
268 with self.assertRaises(StopIteration):
269 next(islice)
270
271 def test_slice_with_step_arg_and_unspecified_stop_ends_at_end_of_iterator(self):
272 islice = itertools.islice([0, 1, 2, 3, 4], 1, None, 2)
273 self.assertEqual(next(islice), 1)
274 self.assertEqual(next(islice), 3)
275 with self.assertRaises(StopIteration):
276 next(islice)
277
278 def test_slice_calls_next_until_stop_is_reached(self):
279 class RaisesAtFive:
280 i = 0
281
282 def __iter__(self):
283 return self
284
285 def __next__(self):
286 if self.i < 5:
287 result = (0, 1, 2, 3, 4)[self.i]
288 self.i += 1
289 return result
290 raise UserWarning(f"Called with {self.i}")
291
292 islice = itertools.islice(RaisesAtFive(), 2, 6, 2)
293 self.assertEqual(next(islice), 2)
294 self.assertEqual(next(islice), 4)
295 with self.assertRaises(UserWarning) as ctx:
296 next(islice)
297 self.assertEqual(str(ctx.exception), "Called with 5")
298
299
300class PermutationsTests(unittest.TestCase):
301 def test_too_few_arguments_raises_type_error(self):
302 self.assertRaises(TypeError, itertools.permutations)
303
304 def test_too_many_arguments_raises_type_error(self):
305 self.assertRaises(TypeError, itertools.permutations, "1", "2", "3", "4")
306
307 def test_non_int_r_type_error(self):
308 self.assertRaises(TypeError, itertools.permutations, "1", 1.0)
309
310 def test_empty_returns_single_empty_permutation(self):
311 self.assertTupleEqual(tuple(itertools.permutations(())), ((),))
312
313 def test_r_defaults_to_length(self):
314 len1 = itertools.permutations("1")
315 self.assertTrue(all(map(lambda x: len(x) == 1, len1)))
316
317 len2 = itertools.permutations("12")
318 self.assertTrue(all(map(lambda x: len(x) == 2, len2)))
319
320 len3 = itertools.permutations("123")
321 self.assertTrue(all(map(lambda x: len(x) == 3, len3)))
322
323 def test_r_zero_returns_single_empty_permutation(self):
324 self.assertTupleEqual(tuple(itertools.permutations("A", 0)), ((),))
325
326 def test_r_gt_length_returns_stopped_iterator(self):
327 self.assertTupleEqual(tuple(itertools.permutations("A", 2)), ())
328
329 def test_r_lt_length_returns_items_with_length_r(self):
330 result = tuple(itertools.permutations("ABC", 2))
331 self.assertTupleEqual(
332 result,
333 (("A", "B"), ("A", "C"), ("B", "A"), ("B", "C"), ("C", "A"), ("C", "B")),
334 )
335
336 def test_ordinary_iterable(self):
337 result = tuple(itertools.permutations(range(3)))
338 self.assertTupleEqual(
339 result, ((0, 1, 2), (0, 2, 1), (1, 0, 2), (1, 2, 0), (2, 0, 1), (2, 1, 0))
340 )
341
342
343class ProductTests(unittest.TestCase):
344 def test_empty_returns_stopped_iterable(self):
345 self.assertTupleEqual(tuple(itertools.product("")), ())
346 self.assertTupleEqual(tuple(itertools.product("ab", "")), ())
347 self.assertTupleEqual(tuple(itertools.product("", "ab", "")), ())
348 self.assertTupleEqual(tuple(itertools.product("ab", "", "12")), ())
349
350 def test_no_args_returns_empty(self):
351 self.assertTupleEqual(tuple(itertools.product()), ((),))
352
353 def test_not_an_integer_raises_type_error(self):
354 with self.assertRaises(TypeError):
355 itertools.product("ab", repeat=None)
356
357 def test_repeat_zero_returns_empty(self):
358 self.assertTupleEqual(tuple(itertools.product("", repeat=0)), ((),))
359 self.assertTupleEqual(tuple(itertools.product("ab", repeat=0)), ((),))
360 self.assertTupleEqual(tuple(itertools.product("ab", "12", repeat=0)), ((),))
361
362 def test_repeat_returns_repeated_iterables(self):
363 self.assertTupleEqual(
364 tuple(itertools.product("a", "1", repeat=2)), (("a", "1", "a", "1"),)
365 )
366
367 def test_different_lengths(self):
368 self.assertTupleEqual(
369 tuple(itertools.product("a", "12")), (("a", "1"), ("a", "2"))
370 )
371 self.assertTupleEqual(
372 tuple(itertools.product("ab", "1")), (("a", "1"), ("b", "1"))
373 )
374 self.assertTupleEqual(
375 tuple(itertools.product("ab", "1", "!")), (("a", "1", "!"), ("b", "1", "!"))
376 )
377
378
379class RepeatTests(unittest.TestCase):
380 def test_dunder_init_with_non_int_times_raises_type_error(self):
381 self.assertRaises(TypeError, itertools.repeat, 5, "not_an_int")
382
383 def test_dunder_init_with_negative_times_repeats_zero_times(self):
384 iterator = itertools.repeat(5, -1)
385 self.assertEqual(tuple(iterator), ())
386
387 def test_dunder_init_with_zero_times_repeats_zero_times(self):
388 iterator = itertools.repeat(5, 0)
389 self.assertEqual(tuple(iterator), ())
390
391 def test_dunder_init_with_three_times_repeats_three_times(self):
392 iterator = itertools.repeat(5, 3)
393 self.assertEqual(next(iterator), 5)
394 self.assertEqual(next(iterator), 5)
395 self.assertEqual(next(iterator), 5)
396 self.assertRaises(StopIteration, next, iterator)
397
398 def test_dunder_init_with_no_times_repeats(self):
399 iterator = itertools.repeat(5)
400 self.assertEqual(next(iterator), 5)
401 self.assertEqual(next(iterator), 5)
402 self.assertEqual(next(iterator), 5)
403
404
405class ZipLongestTests(unittest.TestCase):
406 def test_dunder_init_with_no_seqs_returns_empty_iterator(self):
407 iterator = itertools.zip_longest()
408 self.assertEqual(tuple(iterator), ())
409
410 def test_dunder_init_calls_dunder_iter_on_seqs(self):
411 class C:
412 __iter__ = Mock(name="__iter__", return_value=().__iter__())
413
414 itertools.zip_longest(C(), C())
415 self.assertEqual(C.__iter__.call_count, 2)
416
417 def test_dunder_next_fills_with_none(self):
418 right = itertools.zip_longest([], [1, 2, 3])
419 self.assertEqual(list(right), [(None, 1), (None, 2), (None, 3)])
420
421 left = itertools.zip_longest([1, 2, 3], [])
422 self.assertEqual(list(left), [(1, None), (2, None), (3, None)])
423
424 def test_dunder_next_fills_with_fill_value(self):
425 right = itertools.zip_longest("ab", [1, 2, 3], fillvalue="X")
426 self.assertEqual(list(right), [("a", 1), ("b", 2), ("X", 3)])
427
428
429class AccumulateTests(unittest.TestCase):
430 def test_accumulate_with_iterable_accumulates(self):
431 self.assertTupleEqual(tuple(itertools.accumulate([1, 2, 3, 4])), (1, 3, 6, 10))
432
433 def test_accumulate_with_func_arg_none_uses_addition(self):
434 self.assertTupleEqual(
435 tuple(itertools.accumulate([1, 2, 3, 4], None)), (1, 3, 6, 10)
436 )
437
438 def test_accumulate_with_func_arg_uses_func(self):
439 self.assertTupleEqual(
440 tuple(itertools.accumulate([1, 2, 3, 4], operator.mul)), (1, 2, 6, 24)
441 )
442
443 def test_accumulate_with_initial_arg_none_uses_addition(self):
444 self.assertTupleEqual(
445 tuple(itertools.accumulate([1, 2, 3, 4], initial=None)), (1, 3, 6, 10)
446 )
447
448 def test_accumulate_with_initial_arg_uses_initial(self):
449 self.assertTupleEqual(
450 tuple(itertools.accumulate([1, 2, 3, 4], initial=100)),
451 (100, 101, 103, 106, 110),
452 )
453
454 def test_accumulate_with_initial_arg_uses_initial_with_str(self):
455 self.assertTupleEqual(
456 tuple(itertools.accumulate("ab", initial="c")),
457 ("c", "ca", "cab"),
458 )
459
460 def test_accumulate_with_initial_arg_uses_initial_with_range(self):
461 self.assertTupleEqual(
462 tuple(itertools.accumulate(range(3), initial=9)),
463 (9, 9, 10, 12),
464 )
465
466 def test_accumulate_with_func_and_initial_args_uses_both(self):
467 self.assertTupleEqual(
468 tuple(itertools.accumulate([1, 2, 3, 4], operator.mul, initial=10)),
469 (10, 10, 20, 60, 240),
470 )
471
472 def test_accumulate_with_initial_arg_and_empty_list_uses_initial(self):
473 self.assertTupleEqual(tuple(itertools.accumulate([], initial=100)), (100,))
474
475 def test_accumulate_with_empty_iterable_returns_stopped_iterator(self):
476 iterator = itertools.accumulate([])
477 self.assertRaises(StopIteration, next, iterator)
478
479 def test_accumulate_with_no_iterable_raises_type_error(self):
480 self.assertRaises(TypeError, itertools.accumulate)
481
482 def test_accumulate_dunder_next_with_non_callable_func_raises_type_error(self):
483 iterator = itertools.accumulate([1, 2, 3, 4], "this is not a function")
484 next(iterator)
485 self.assertRaisesRegex(
486 TypeError, "'str' object is not callable", next, iterator
487 )
488
489
490class GroupbyTests(unittest.TestCase):
491 def test_groupby_returns_groups(self):
492 it = itertools.groupby("AAAABBBCCD")
493 self.assertEqual(next(it)[0], "A")
494 self.assertEqual(next(it)[0], "B")
495 self.assertEqual(next(it)[0], "C")
496 self.assertEqual(next(it)[0], "D")
497
498 def test_groupby_returns_all_elements_in_all_groups(self):
499 it = itertools.groupby("AAAABBBCCD")
500 group = next(it)[1]
501 self.assertTupleEqual(tuple(group), ("A", "A", "A", "A"))
502 group = next(it)[1]
503 self.assertTupleEqual(tuple(group), ("B", "B", "B"))
504 group = next(it)[1]
505 self.assertTupleEqual(tuple(group), ("C", "C"))
506 group = next(it)[1]
507 self.assertEqual(next(group), "D")
508
509 def test_groupby_raises_stopiteration_after_all_groups(self):
510 it = itertools.groupby("AAAABBBCCD")
511 for _i in range(4):
512 next(it)
513 self.assertRaises(StopIteration, next, it)
514
515 def test_groupby_first_group_raises_stopiteration_after_all_elements(self):
516 it = itertools.groupby("AAAABBBCCD")
517 group = next(it)[1]
518 for _i in range(4):
519 next(group)
520 self.assertRaises(StopIteration, next, group)
521
522 def test_groupby_last_group_raises_stopiteration_after_all_elements(self):
523 it = itertools.groupby("AAAABBBCCD")
524 for _i in range(3):
525 next(it)
526 group = next(it)[1]
527 next(group)
528 self.assertRaises(StopIteration, next, group)
529
530 def test_groupby_group_raises_stopiteration_after_advancing_groupby(self):
531 it = itertools.groupby("AAAABBBCCD")
532 group = next(it)[1]
533 next(it)
534 self.assertRaises(StopIteration, next, group)
535
536 def test_groupby_handles_none_values(self):
537 it = itertools.groupby([1, None, None, None])
538 group = next(it)[1]
539 self.assertEqual(next(group), 1)
540 group = next(it)[1]
541 self.assertTupleEqual(tuple(group), (None, None, None))
542
543 def test_groupby_with_key_func_returns_keys_and_groups(self):
544 def keyfunc(value):
545 return 4 if value == 4 else 1
546
547 it = itertools.groupby([1, 2, 3, 4], keyfunc)
548 grouper = next(it)
549 self.assertEqual(grouper[0], 1)
550 self.assertTupleEqual(tuple(grouper[1]), (1, 2, 3))
551 grouper = next(it)
552 self.assertEqual(grouper[0], 4)
553 self.assertEqual(next(grouper[1]), 4)
554
555 def test_groupby_terminates(self):
556 vals = [1, 2, "hello", "bla", 5]
557 groups = itertools.groupby(vals, lambda x: type(x))
558 self.assertEqual([1, 2], list(next(groups)[1]))
559 self.assertEqual(["hello", "bla"], list(next(groups)[1]))
560 self.assertEqual([5], list(next(groups)[1]))
561 self.assertRaises(StopIteration, next, groups)
562
563
564class FilterFalseTests(unittest.TestCase):
565 def test_filterfalse_with_no_predicate_returns_false_values(self):
566 it = itertools.filterfalse(None, range(10))
567 self.assertTupleEqual(tuple(it), (0,))
568
569 def test_filterfalse_with_predicate_returns_filtered_values(self):
570 it = itertools.filterfalse(lambda x: x % 2, range(10))
571 self.assertTupleEqual(tuple(it), (0, 2, 4, 6, 8))
572
573 def test_filterfalse_with_no_sequence_raises_typeerror(self):
574 with self.assertRaises(TypeError):
575 itertools.filterfalse(None, None)
576
577
578class TeeTests(unittest.TestCase):
579 def test_default_n_retuns_two_iterators(self):
580 its = itertools.tee([1, 2, 3, 4, 5])
581 self.assertEqual(len(its), 2)
582
583 def test_tee_returns_multiple_working_iterators(self):
584 its = itertools.tee([2, 4, 6, 8, 10])
585 self.assertTupleEqual(tuple(its[0]), (2, 4, 6, 8, 10))
586 self.assertTupleEqual(tuple(its[1]), (2, 4, 6, 8, 10))
587
588 def test_tee_with_long_iterable(self):
589 its = itertools.tee(list(range(200)))
590 self.assertTupleEqual(tuple(its[0]), tuple(range(200)))
591 self.assertTupleEqual(tuple(its[1]), tuple(range(200)))
592
593 def test_tee_with_n_equal_three_returns_three_working_iterators(self):
594 its = itertools.tee(["A", "B"], 3)
595 self.assertTupleEqual(tuple(its[0]), ("A", "B"))
596 self.assertTupleEqual(tuple(its[1]), ("A", "B"))
597 self.assertTupleEqual(tuple(its[2]), ("A", "B"))
598
599 def test_tee_with_copyable_iterator_as_class_attribute(self):
600 class CopyableRangeIterator:
601 def __init__(self, n, i=0):
602 self._n = n
603 self._i = i
604
605 def __iter__(self):
606 return self
607
608 def __next__(self):
609 if self._i == self._n:
610 raise StopIteration
611 value = self._i
612 self._i += 1
613 return value
614
615 def __copy__(self):
616 return self.__class__(self._n)
617
618 its = itertools.tee(CopyableRangeIterator(2))
619 self.assertTrue(isinstance(its[0], CopyableRangeIterator))
620 self.assertTupleEqual(tuple(its[0]), (0, 1))
621 self.assertTupleEqual(tuple(its[1]), (0, 1))
622
623 def test_tee_with_copyable_iterator_as_instance_attribute(self):
624 class CopyableRangeIterator:
625 def __init__(self, n, i=0):
626 self._n = n
627 self._i = i
628
629 def __iter__(self):
630 return self
631
632 def __next__(self):
633 if self._i == self._n:
634 raise StopIteration
635 value = self._i
636 self._i += 1
637 return value
638
639 # This is a bit contrived
640 it = CopyableRangeIterator(3)
641 it.__copy__ = lambda: CopyableRangeIterator(it._n)
642 its = itertools.tee(it)
643 self.assertTrue(isinstance(its[0], CopyableRangeIterator))
644 self.assertTupleEqual(tuple(its[0]), (0, 1, 2))
645 self.assertTupleEqual(tuple(its[1]), (0, 1, 2))
646
647 def test_tee_with_non_integer_n_raises_typeerror(self):
648 with self.assertRaises(TypeError):
649 itertools.tee([1, 2, 3, 4, 5], "2")
650
651 def test_tee_with_n_lessthan_zero_raises_valueerror(self):
652 with self.assertRaises(ValueError):
653 itertools.tee([1, 2, 3, 4, 5], -2)
654
655 def test_tee_with_n_equal_zero_returns_empty_tuple(self):
656 self.assertEqual(itertools.tee([1, 2, 3, 4, 5], 0), ())
657
658
659class DropWhileTests(unittest.TestCase):
660 def test_dropwhile_passing_none_predicate_raises_typeerror(self):
661 it = itertools.dropwhile(None, [1, 2, 3])
662 self.assertRaises(TypeError, next, it)
663
664 def test_dropwhile_passing_none_iterator_raises_typeerror(self):
665 with self.assertRaises(TypeError):
666 itertools.dropwhile(lambda x: x % 2, None)
667
668 def test_dropwhile_returns_correct_elements_dropping_start(self):
669 it = itertools.dropwhile(lambda x: x < 5, [1, 4, 6, 4, 1])
670 self.assertTupleEqual(tuple(it), (6, 4, 1))
671
672 def test_dropwhile_with_true_predicate_drops_all_elements(self):
673 it = itertools.dropwhile(lambda x: True, [1, 4, 6, 4, 1])
674 self.assertTupleEqual(tuple(it), ())
675
676 def test_dropwhile_with_false_predicate_returns_all_elements(self):
677 it = itertools.dropwhile(lambda x: False, [1, 4, 6, 4, 1])
678 self.assertTupleEqual(tuple(it), (1, 4, 6, 4, 1))
679
680
681class TakeWhileTests(unittest.TestCase):
682 def test_takewhile_passing_none_predicate_raises_typeerror(self):
683 it = itertools.takewhile(None, [1, 2, 3])
684 self.assertRaises(TypeError, next, it)
685
686 def test_takewhile_passing_none_iterator_raises_typeerror(self):
687 with self.assertRaises(TypeError):
688 itertools.takewhile(lambda x: x % 2, None)
689
690 def test_takewhile_returns_correct_elements_dropping_end(self):
691 it = itertools.takewhile(lambda x: x < 5, [1, 4, 6, 4, 1])
692 self.assertTupleEqual(tuple(it), (1, 4))
693
694 def test_takewhile_with_true_predicate_returns_all_elements(self):
695 it = itertools.takewhile(lambda x: True, [1, 4, 6, 4, 1])
696 self.assertTupleEqual(tuple(it), (1, 4, 6, 4, 1))
697
698 def test_takewhile_with_false_predicate_drops_all_elements(self):
699 it = itertools.takewhile(lambda x: False, [1, 4, 6, 4, 1])
700 self.assertTupleEqual(tuple(it), ())
701
702
703class StarMapTests(unittest.TestCase):
704 def test_starmap_returns_arguments_mapped_onto_function(self):
705 it = itertools.starmap(pow, [(2, 5), (3, 2), (10, 3)])
706 self.assertTupleEqual(tuple(it), (32, 9, 1000))
707
708 def test_starmap_handles_tuple_subclass(self):
709 class C(tuple):
710 pass
711
712 points = [C((3, 4)), C((9, 10)), C((-2, -1))]
713 self.assertTrue(isinstance(points[0], tuple))
714 self.assertFalse(points[0] is tuple)
715
716 it = itertools.starmap(lambda x, y: x ** 2 + y ** 2, points)
717 self.assertTupleEqual(tuple(it), (25, 181, 5))
718
719 def test_starmap_passing_empty_sequence_returns_empty_iterator(self):
720 self.assertEqual(tuple(itertools.starmap(pow, [])), ())
721
722 def test_starmap_passing_none_function_raises_typeerror(self):
723 it = itertools.starmap(None, [(2, 5), (3, 2), (10, 3)])
724 self.assertRaises(TypeError, next, it)
725
726 def test_starmap_passing_none_iterable_raises_typeerror(self):
727 with self.assertRaises(TypeError):
728 itertools.starmap(pow, None)
729
730 def test_starmap_passing_non_sequences_raises_typeerror(self):
731 it = itertools.starmap(None, [1, 2])
732 self.assertRaisesRegex(TypeError, "'int' object is not iterable", next, it)
733
734
735class CombinationsTests(unittest.TestCase):
736 def test_too_few_arguments_raises_type_error(self):
737 self.assertRaises(TypeError, itertools.combinations)
738
739 def test_too_many_arguments_raises_type_error(self):
740 self.assertRaises(TypeError, itertools.combinations, "1", "2", "3", "4")
741
742 def test_non_int_r_type_error(self):
743 self.assertRaises(TypeError, itertools.combinations, "1", 1.0)
744
745 def test_empty_returns_single_empty_combination(self):
746 self.assertTupleEqual(tuple(itertools.combinations((), 0)), ((),))
747
748 def test_r_zero_returns_single_empty_combination(self):
749 self.assertTupleEqual(tuple(itertools.combinations("A", 0)), ((),))
750
751 def test_r_gt_length_returns_stopped_iterator(self):
752 self.assertTupleEqual(tuple(itertools.combinations("A", 2)), ())
753
754 def test_r_lt_length_returns_items_with_length_r(self):
755 result = tuple(itertools.combinations("Bam!", 2))
756 self.assertTupleEqual(
757 result,
758 (("B", "a"), ("B", "m"), ("B", "!"), ("a", "m"), ("a", "!"), ("m", "!")),
759 )
760
761 def test_ordinary_iterable(self):
762 result = tuple(itertools.combinations(range(4), 3))
763 self.assertTupleEqual(result, ((0, 1, 2), (0, 1, 3), (0, 2, 3), (1, 2, 3)))
764
765
766class CombinationsWithReplacementTests(unittest.TestCase):
767 def test_too_few_arguments_raises_type_error(self):
768 self.assertRaises(TypeError, itertools.combinations_with_replacement)
769
770 def test_too_many_arguments_raises_type_error(self):
771 self.assertRaises(
772 TypeError, itertools.combinations_with_replacement, "1", "2", "3", "4"
773 )
774
775 def test_non_int_r_type_error(self):
776 self.assertRaises(TypeError, itertools.combinations_with_replacement, "1", 1.0)
777
778 def test_empty_returns_single_empty_combination(self):
779 self.assertTupleEqual(
780 tuple(itertools.combinations_with_replacement((), 0)), ((),)
781 )
782
783 def test_r_zero_returns_single_empty_combination(self):
784 self.assertTupleEqual(
785 tuple(itertools.combinations_with_replacement("A", 0)), ((),)
786 )
787
788 def test_r_gt_length_returns_items_with_length_r(self):
789 self.assertTupleEqual(
790 tuple(itertools.combinations_with_replacement("A", 2)), (("A", "A"),)
791 )
792
793 def test_r_positive_and_length_zero_returns_stopped_iterator(self):
794 self.assertTupleEqual(tuple(itertools.combinations("A", 2)), ())
795
796 def test_r_lt_length_returns_items_with_length_r(self):
797 result = tuple(itertools.combinations_with_replacement("CBA", 2))
798 self.assertTupleEqual(
799 result,
800 (("C", "C"), ("C", "B"), ("C", "A"), ("B", "B"), ("B", "A"), ("A", "A")),
801 )
802
803 def test_ordinary_iterable(self):
804 result = tuple(itertools.combinations_with_replacement(range(4), 1))
805 self.assertTupleEqual(result, ((0,), (1,), (2,), (3,)))
806
807
808class CompressTests(unittest.TestCase):
809 def test_empty_data_returns_stopped_iterator(self):
810 self.assertEqual(tuple(itertools.compress([], [1, 0])), ())
811
812 def test_empty_selectors_returns_stopped_iterator(self):
813 self.assertEqual(tuple(itertools.compress("TYLER", [])), ())
814
815 def test_numeric_selector_values(self):
816 result = "".join(itertools.compress("T_Y_L_E_R", [1, 0, 1, 0, 1, 0, 1, 0, 1]))
817 self.assertEqual(result, "TYLER")
818
819 def test_boolean_selector_values(self):
820 result = tuple(itertools.compress([2, 3, 4], [True, False, False]))
821 self.assertTupleEqual(result, (2,))
822
823 def test_other_selector_values(self):
824 result = "".join(
825 itertools.compress("FOUBAR", [None, [1], (), "hello", b"", {1: 2}])
826 )
827 self.assertEqual(result, "OBR")
828
829 def test_data_length_gt_selector_length(self):
830 result = tuple(itertools.compress([2, 3, 4, 5, 6], [1, 0, 1]))
831 self.assertTupleEqual(result, (2, 4))
832
833 def test_data_length_lt_selector_length(self):
834 result = tuple(itertools.compress([2, 3, 4], [1, 0, 1, 1, 1, 1]))
835 self.assertTupleEqual(result, (2, 4))
836
837 def test_passing_non_iterable_data_raises_typeerror(self):
838 self.assertRaises(TypeError, itertools.compress, None, [1, 2, 3])
839
840 def test_passing_non_iterable_selectors_raises_typeerror(self):
841 self.assertRaises(TypeError, itertools.compress, [1, 2, 3], None)
842
843
844if __name__ == "__main__":
845 unittest.main()