diff --git a/Lib/test/test_itertools.py b/Lib/test/test_itertools.py index b0ddbb0c8..ead224c6a 100644 --- a/Lib/test/test_itertools.py +++ b/Lib/test/test_itertools.py @@ -12,6 +12,8 @@ from functools import reduce import sys import struct import threading +import gc + maxsize = support.MAX_Py_ssize_t minsize = -maxsize-1 @@ -193,7 +195,6 @@ class TestBasicOps(unittest.TestCase): self.assertRaises(TypeError, list, oper(chain(2, 3))) for proto in range(pickle.HIGHEST_PROTOCOL + 1): self.pickletest(proto, chain('abc', 'def'), compare=list('abcdef')) - # TODO: RUSTPYTHON @unittest.expectedFailure def test_chain_setstate(self): @@ -208,7 +209,6 @@ class TestBasicOps(unittest.TestCase): it = chain() it.__setstate__((iter(['abc', 'def']), iter(['ghi']))) self.assertEqual(list(it), ['ghi', 'a', 'b', 'c', 'd', 'e', 'f']) - # TODO: RUSTPYTHON @unittest.expectedFailure def test_combinations(self): @@ -1053,6 +1053,25 @@ class TestBasicOps(unittest.TestCase): self.assertEqual(next(it), (1, 2)) self.assertRaises(RuntimeError, next, it) + def test_pairwise(self): + self.assertEqual(list(pairwise('')), []) + self.assertEqual(list(pairwise('a')), []) + self.assertEqual(list(pairwise('ab')), + [('a', 'b')]), + self.assertEqual(list(pairwise('abcde')), + [('a', 'b'), ('b', 'c'), ('c', 'd'), ('d', 'e')]) + self.assertEqual(list(pairwise(range(10_000))), + list(zip(range(10_000), range(1, 10_000)))) + + with self.assertRaises(TypeError): + pairwise() # too few arguments + with self.assertRaises(TypeError): + pairwise('abc', 10) # too many arguments + with self.assertRaises(TypeError): + pairwise(iterable='abc') # keyword arguments + with self.assertRaises(TypeError): + pairwise(None) # non-iterable argument + def test_product(self): for args, result in [ ([], [()]), # zero iterables @@ -1609,6 +1628,51 @@ class TestBasicOps(unittest.TestCase): self.assertRaises(StopIteration, next, f(lambda x:x, [])) self.assertRaises(StopIteration, next, f(lambda x:x, StopNow())) + @support.cpython_only + def test_combinations_result_gc(self): + # bpo-42536: combinations's tuple-reuse speed trick breaks the GC's + # assumptions about what can be untracked. Make sure we re-track result + # tuples whenever we reuse them. + it = combinations([None, []], 1) + next(it) + gc.collect() + # That GC collection probably untracked the recycled internal result + # tuple, which has the value (None,). Make sure it's re-tracked when + # it's mutated and returned from __next__: + self.assertTrue(gc.is_tracked(next(it))) + + @support.cpython_only + def test_combinations_with_replacement_result_gc(self): + # Ditto for combinations_with_replacement. + it = combinations_with_replacement([None, []], 1) + next(it) + gc.collect() + self.assertTrue(gc.is_tracked(next(it))) + + @support.cpython_only + def test_permutations_result_gc(self): + # Ditto for permutations. + it = permutations([None, []], 1) + next(it) + gc.collect() + self.assertTrue(gc.is_tracked(next(it))) + + @support.cpython_only + def test_product_result_gc(self): + # Ditto for product. + it = product([None, []]) + next(it) + gc.collect() + self.assertTrue(gc.is_tracked(next(it))) + + @support.cpython_only + def test_zip_longest_result_gc(self): + # Ditto for zip_longest. + it = zip_longest([[]]) + gc.collect() + self.assertTrue(gc.is_tracked(next(it))) + + class TestExamples(unittest.TestCase): def test_accumulate(self): @@ -1848,6 +1912,10 @@ class TestGC(unittest.TestCase): a = [] self.makecycle(islice([a]*2, None), a) + def test_pairwise(self): + a = [] + self.makecycle(pairwise([a]*5), a) + def test_permutations(self): a = [] self.makecycle(permutations([1,2,a,3], 3), a) @@ -1946,6 +2014,7 @@ def L(seqn): class TestVariousIteratorArgs(unittest.TestCase): + def test_accumulate(self): s = [1,2,3,4,5] r = [1,3,6,10,15] @@ -2055,6 +2124,17 @@ class TestVariousIteratorArgs(unittest.TestCase): self.assertRaises(TypeError, islice, N(s), 10) self.assertRaises(ZeroDivisionError, list, islice(E(s), 10)) + def test_pairwise(self): + for s in ("123", "", range(1000), ('do', 1.2), range(2000,2200,5)): + for g in (G, I, Ig, S, L, R): + seq = list(g(s)) + expected = list(zip(seq, seq[1:])) + actual = list(pairwise(g(s))) + self.assertEqual(actual, expected) + self.assertRaises(TypeError, pairwise, X(s)) + self.assertRaises(TypeError, pairwise, N(s)) + self.assertRaises(ZeroDivisionError, list, pairwise(E(s))) + def test_starmap(self): for s in (range(10), range(0), range(100), (7,11), range(20,50,5)): for g in (G, I, Ig, S, L, R): @@ -2356,7 +2436,7 @@ Samuele ... "Count how many times the predicate is true" ... return sum(map(pred, iterable)) ->>> def padnone(iterable): +>>> def pad_none(iterable): ... "Returns the sequence elements and then returns None indefinitely" ... return chain(iterable, repeat(None)) @@ -2378,15 +2458,6 @@ Samuele ... else: ... return starmap(func, repeat(args, times)) ->>> def pairwise(iterable): -... "s -> (s0,s1), (s1,s2), (s2, s3), ..." -... a, b = tee(iterable) -... try: -... next(b) -... except StopIteration: -... pass -... return zip(a, b) - >>> def grouper(n, iterable, fillvalue=None): ... "grouper(3, 'ABCDEFG', 'x') --> ABC DEF Gxx" ... args = [iter(iterable)] * n @@ -2517,16 +2588,7 @@ True >>> take(5, map(int, repeatfunc(random.random))) [0, 0, 0, 0, 0] ->>> list(pairwise('abcd')) -[('a', 'b'), ('b', 'c'), ('c', 'd')] - ->>> list(pairwise([])) -[] - ->>> list(pairwise('a')) -[] - ->>> list(islice(padnone('abc'), 0, 6)) +>>> list(islice(pad_none('abc'), 0, 6)) ['a', 'b', 'c', None, None, None] >>> list(ncycles('abc', 3)) diff --git a/vm/src/stdlib/itertools.rs b/vm/src/stdlib/itertools.rs index 48b73adb1..f10565112 100644 --- a/vm/src/stdlib/itertools.rs +++ b/vm/src/stdlib/itertools.rs @@ -804,10 +804,20 @@ mod decl { #[derive(Debug)] struct PyItertoolsAccumulate { iterable: PyObjectRef, - binop: PyObjectRef, + binop: Option, + initial: Option, acc_value: PyRwLock>, } + #[derive(FromArgs)] + struct AccumulateArgs { + iterable: PyObjectRef, + #[pyarg(any, optional)] + func: OptionalOption, + #[pyarg(named, optional)] + initial: OptionalOption, + } + impl PyValue for PyItertoolsAccumulate { fn class(_vm: &VirtualMachine) -> &PyTypeRef { Self::static_type() @@ -819,15 +829,15 @@ mod decl { #[pyslot] fn tp_new( cls: PyTypeRef, - iterable: PyObjectRef, - binop: OptionalArg, + args: AccumulateArgs, vm: &VirtualMachine, ) -> PyResult> { - let iter = get_iter(vm, iterable)?; + let iter = get_iter(vm, args.iterable)?; PyItertoolsAccumulate { iterable: iter, - binop: binop.unwrap_or_none(vm), + binop: args.func.flatten(), + initial: args.initial.flatten(), acc_value: PyRwLock::new(None), } .into_ref_with_type(vm, cls) @@ -836,17 +846,19 @@ mod decl { impl PyIter for PyItertoolsAccumulate { fn next(zelf: &PyRef, vm: &VirtualMachine) -> PyResult { let iterable = &zelf.iterable; - let obj = call_next(vm, iterable)?; let acc_value = zelf.acc_value.read().clone(); let next_acc_value = match acc_value { - None => obj, + None => match &zelf.initial { + None => call_next(vm, iterable)?, + Some(obj) => obj.clone(), + }, Some(value) => { - if vm.is_none(&zelf.binop) { - vm._add(&value, &obj)? - } else { - vm.invoke(&zelf.binop, vec![value, obj])? + let obj = call_next(vm, iterable)?; + match &zelf.binop { + None => vm._add(&value, &obj)?, + Some(op) => vm.invoke(op, vec![value, obj])?, } } }; @@ -1387,7 +1399,7 @@ mod decl { } #[derive(FromArgs)] - struct ZiplongestArgs { + struct ZipLongestArgs { #[pyarg(named, optional)] fillvalue: OptionalArg, } @@ -1398,7 +1410,7 @@ mod decl { fn tp_new( cls: PyTypeRef, iterables: Args, - args: ZiplongestArgs, + args: ZipLongestArgs, vm: &VirtualMachine, ) -> PyResult> { let fillvalue = args.fillvalue.unwrap_or_none(vm); @@ -1442,4 +1454,47 @@ mod decl { } } } + + #[pyattr] + #[pyclass(name = "pairwise")] + #[derive(Debug)] + struct PyItertoolsPairwise { + iterator: PyObjectRef, + old: PyRwLock>, + } + + impl PyValue for PyItertoolsPairwise { + fn class(_vm: &VirtualMachine) -> &PyTypeRef { + Self::static_type() + } + } + + #[pyimpl(with(PyIter))] + impl PyItertoolsPairwise { + #[pyslot] + fn tp_new( + cls: PyTypeRef, + iterable: PyObjectRef, + vm: &VirtualMachine, + ) -> PyResult> { + let iterator = get_iter(vm, iterable)?; + + PyItertoolsPairwise { + iterator, + old: PyRwLock::new(None), + } + .into_ref_with_type(vm, cls) + } + } + impl PyIter for PyItertoolsPairwise { + fn next(zelf: &PyRef, vm: &VirtualMachine) -> PyResult { + let old = match zelf.old.read().clone() { + None => call_next(vm, &zelf.iterator)?, + Some(obj) => obj, + }; + let new = call_next(vm, &zelf.iterator)?; + *zelf.old.write() = Some(new.clone()); + Ok(vm.ctx.new_tuple(vec![old, new])) + } + } }