Merge pull request RustPython#2473 from rickygao/itertools-support

Follow the latest `itertools` module

Author: coolreader18

Lib/test/test_itertools.py

@@ -12,6 +12,8 @@
 import sys
 import struct
 import threading
+import gc
+
 maxsize = support.MAX_Py_ssize_t
 minsize = -maxsize-1
 
@@ -193,7 +195,6 @@ def test_chain_reducible(self):
             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 @@ def test_chain_setstate(self):
         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 @@ def run(r1, r2):
         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 @@ def test_StopIteration(self):
             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 @@ def test_islice(self):
         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 @@ def test_islice(self):
             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 @@ def test_permutations_sizeof(self):
 ...     "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 @@ def test_permutations_sizeof(self):
 ...     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 @@ def test_permutations_sizeof(self):
 >>> 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))

vm/src/stdlib/itertools.rs

@@ -804,10 +804,20 @@ mod decl {
     #[derive(Debug)]
     struct PyItertoolsAccumulate {
         iterable: PyObjectRef,
-        binop: PyObjectRef,
+        binop: Option<PyObjectRef>,
+        initial: Option<PyObjectRef>,
         acc_value: PyRwLock<Option<PyObjectRef>>,
     }
 
+    #[derive(FromArgs)]
+    struct AccumulateArgs {
+        iterable: PyObjectRef,
+        #[pyarg(any, optional)]
+        func: OptionalOption<PyObjectRef>,
+        #[pyarg(named, optional)]
+        initial: OptionalOption<PyObjectRef>,
+    }
+
     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<PyObjectRef>,
+            args: AccumulateArgs,
             vm: &VirtualMachine,
         ) -> PyResult<PyRef<Self>> {
-            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<Self>, 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<PyObjectRef>,
     }
@@ -1398,7 +1410,7 @@ mod decl {
         fn tp_new(
             cls: PyTypeRef,
             iterables: Args,
-            args: ZiplongestArgs,
+            args: ZipLongestArgs,
             vm: &VirtualMachine,
         ) -> PyResult<PyRef<Self>> {
             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<Option<PyObjectRef>>,
+    }
+
+    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<PyRef<Self>> {
+            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<Self>, 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]))
+        }
+    }
 }