diff --git a/Lib/test/test_array.py b/Lib/test/test_array.py index f7405098b1c..b68d85fbd40 100644 --- a/Lib/test/test_array.py +++ b/Lib/test/test_array.py @@ -284,8 +284,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) diff --git a/stdlib/src/array.rs b/stdlib/src/array.rs index b7fd8502b39..6dc7fb1bdfb 100644 --- a/stdlib/src/array.rs +++ b/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, 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>, } - #[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, vm: &VirtualMachine) -> PyResult { - 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), + }) + }) } }