Merge pull request RustPython#3868 from maong0927/arrayiter_reduce

add arrayiter.__reduce__

Author: youknowone

Lib/test/test_array.py

@@ -304,8 +304,6 @@ def test_pickle_for_empty_array(self):
             self.assertEqual(a.x, b.x)
             self.assertEqual(type(a), type(b))
 
-    # TODO: RUSTPYTHON
-    @unittest.expectedFailure
     def test_iterator_pickle(self):
         orig = array.array(self.typecode, self.example)
         data = list(orig)

stdlib/src/array.rs

@@ -6,15 +6,15 @@ mod array {
         common::{
             atomic::{self, AtomicUsize},
             lock::{
-                PyMappedRwLockReadGuard, PyMappedRwLockWriteGuard, PyRwLock, PyRwLockReadGuard,
-                PyRwLockWriteGuard,
+                PyMappedRwLockReadGuard, PyMappedRwLockWriteGuard, PyMutex, PyRwLock,
+                PyRwLockReadGuard, PyRwLockWriteGuard,
             },
             str::wchar_t,
         },
         vm::{
             builtins::{
-                PyByteArray, PyBytes, PyBytesRef, PyDictRef, PyFloat, PyInt, PyIntRef, PyList,
-                PyListRef, PyStr, PyStrRef, PyTupleRef, PyTypeRef,
+                PositionIterInternal, PyByteArray, PyBytes, PyBytesRef, PyDictRef, PyFloat, PyInt,
+                PyIntRef, PyList, PyListRef, PyStr, PyStrRef, PyTupleRef, PyTypeRef,
             },
             class_or_notimplemented,
             convert::{ToPyObject, ToPyResult, TryFromObject},
@@ -1258,9 +1258,8 @@ mod array {
 
     impl Iterable for PyArray {
         fn iter(zelf: PyRef<Self>, vm: &VirtualMachine) -> PyResult {
-            Ok(PyArrayIterator {
-                position: AtomicUsize::new(0),
-                array: zelf,
+            Ok(PyArrayIter {
+                internal: PyMutex::new(PositionIterInternal::new(zelf, 0)),
             }
             .into_pyobject(vm))
         }
@@ -1278,24 +1277,36 @@ mod array {
     #[pyattr]
     #[pyclass(name = "arrayiterator")]
     #[derive(Debug, PyPayload)]
-    pub struct PyArrayIterator {
-        position: AtomicUsize,
-        array: PyArrayRef,
+    pub struct PyArrayIter {
+        internal: PyMutex<PositionIterInternal<PyArrayRef>>,
     }
 
-    #[pyimpl(with(IterNext))]
-    impl PyArrayIterator {}
+    #[pyimpl(with(IterNext), flags(HAS_DICT))]
+    impl PyArrayIter {
+        #[pymethod(magic)]
+        fn setstate(&self, state: PyObjectRef, vm: &VirtualMachine) -> PyResult<()> {
+            self.internal
+                .lock()
+                .set_state(state, |obj, pos| pos.min(obj.len()), vm)
+        }
 
-    impl IterNextIterable for PyArrayIterator {}
-    impl IterNext for PyArrayIterator {
+        #[pymethod(magic)]
+        fn reduce(&self, vm: &VirtualMachine) -> PyTupleRef {
+            self.internal
+                .lock()
+                .builtins_iter_reduce(|x| x.clone().into(), vm)
+        }
+    }
+
+    impl IterNextIterable for PyArrayIter {}
+    impl IterNext for PyArrayIter {
         fn next(zelf: &Py<Self>, vm: &VirtualMachine) -> PyResult<PyIterReturn> {
-            let pos = zelf.position.fetch_add(1, atomic::Ordering::SeqCst);
-            let r = if let Some(item) = zelf.array.read().get(pos, vm) {
-                PyIterReturn::Return(item?)
-            } else {
-                PyIterReturn::StopIteration(None)
-            };
-            Ok(r)
+            zelf.internal.lock().next(|array, pos| {
+                Ok(match array.read().get(pos, vm) {
+                    Some(item) => PyIterReturn::Return(item?),
+                    None => PyIterReturn::StopIteration(None),
+                })
+            })
         }
     }