Implement missing Iterator & ReverseIterator and its methods (RustPython#2906)

Implement missing methods of string iterators.

Author: Snowapril

Lib/test/test_iter.py

@@ -313,8 +313,6 @@ def test_iter_range(self):
         self.check_for_loop(iter(range(10)), list(range(10)))
 
     # Test a string
-    # TODO: RUSTPYTHON
-    @unittest.expectedFailure
     def test_iter_string(self):
         self.check_for_loop(iter("abcde"), ["a", "b", "c", "d", "e"])
 

vm/src/builtins/pystr.rs

@@ -3,6 +3,7 @@ use std::ops::Range;
 use std::string::ToString;
 use std::{char, ffi, fmt};
 
+use crossbeam_utils::atomic::AtomicCell;
 use itertools::Itertools;
 use num_traits::ToPrimitive;
 use unic_ucd_bidi::BidiClass;
@@ -12,7 +13,11 @@ use unicode_casing::CharExt;
 
 use super::bytes::PyBytesRef;
 use super::dict::PyDict;
-use super::int::{PyInt, PyIntRef};
+use super::int::{try_to_primitive, PyInt, PyIntRef};
+use super::iter::{
+    IterStatus,
+    IterStatus::{Active, Exhausted},
+};
 use super::pytype::PyTypeRef;
 use crate::anystr::{self, adjust_indices, AnyStr, AnyStrContainer, AnyStrWrapper};
 use crate::exceptions::IntoPyException;
@@ -113,6 +118,7 @@ impl TryIntoRef<PyStr> for &str {
 pub struct PyStrIterator {
     string: PyStrRef,
     position: PyAtomic<usize>,
+    status: AtomicCell<IterStatus>,
 }
 
 impl PyValue for PyStrIterator {
@@ -122,20 +128,66 @@ impl PyValue for PyStrIterator {
 }
 
 #[pyimpl(with(PyIter))]
-impl PyStrIterator {}
+impl PyStrIterator {
+    #[pymethod(magic)]
+    fn length_hint(&self) -> usize {
+        match self.status.load() {
+            Active => {
+                let pos = self.position.load(atomic::Ordering::SeqCst);
+                self.string.len().saturating_sub(pos)
+            }
+            Exhausted => 0,
+        }
+    }
+
+    #[pymethod(magic)]
+    fn setstate(&self, state: PyObjectRef, vm: &VirtualMachine) -> PyResult<()> {
+        // When we're exhausted, just return.
+        if let Exhausted = self.status.load() {
+            return Ok(());
+        }
+        let pos = state
+            .payload::<PyInt>()
+            .ok_or_else(|| vm.new_type_error("an integer is required.".to_owned()))?;
+        let pos = std::cmp::min(
+            try_to_primitive(pos.as_bigint(), vm).unwrap_or(0),
+            self.string.len(),
+        );
+        self.position.store(pos, atomic::Ordering::SeqCst);
+        Ok(())
+    }
+
+    #[pymethod(magic)]
+    fn reduce(&self, vm: &VirtualMachine) -> PyResult {
+        let iter = vm.get_attribute(vm.builtins.clone(), "iter")?;
+        Ok(vm.ctx.new_tuple(match self.status.load() {
+            Exhausted => vec![iter, vm.ctx.new_tuple(vec![vm.ctx.new_str("")])],
+            Active => vec![
+                iter,
+                vm.ctx.new_tuple(vec![self.string.clone().into_object()]),
+                vm.ctx
+                    .new_int(self.position.load(atomic::Ordering::Relaxed)),
+            ],
+        }))
+    }
+}
 
 impl PyIter for PyStrIterator {
     fn next(zelf: &PyRef<Self>, vm: &VirtualMachine) -> PyResult {
+        if let Exhausted = zelf.status.load() {
+            return Err(vm.new_stop_iteration());
+        }
         let value = &*zelf.string.value;
         let mut start = zelf.position.load(atomic::Ordering::SeqCst);
         loop {
             if start == value.len() {
+                zelf.status.store(Exhausted);
                 return Err(vm.new_stop_iteration());
             }
-            let ch = value[start..]
-                .chars()
-                .next()
-                .ok_or_else(|| vm.new_stop_iteration())?;
+            let ch = value[start..].chars().next().ok_or_else(|| {
+                zelf.status.store(Exhausted);
+                vm.new_stop_iteration()
+            })?;
 
             match zelf.position.compare_exchange_weak(
                 start,
@@ -1123,6 +1175,7 @@ impl Iterable for PyStr {
         Ok(PyStrIterator {
             position: Radium::new(0),
             string: zelf,
+            status: AtomicCell::new(Active),
         }
         .into_object(vm))
     }