fix __length_hint__

vm/src/builtins/bytearray.rs

@@ -1,5 +1,7 @@
 //! Implementation of the python bytearray object.
-use super::{PyBytes, PyBytesRef, PyDictRef, PyIntRef, PyStrRef, PyTupleRef, PyTypeRef};
+use super::{
+    PositionIterInternal, PyBytes, PyBytesRef, PyDictRef, PyIntRef, PyStrRef, PyTupleRef, PyTypeRef,
+};
 use crate::common::{
     borrow::{BorrowedValue, BorrowedValueMut},
     lock::{
@@ -742,8 +744,8 @@ impl PyValue for PyByteArrayIterator {
 #[pyimpl(with(SlotIterator))]
 impl PyByteArrayIterator {
     #[pymethod(magic)]
-    fn length_hint(&self, vm: &VirtualMachine) -> PyObjectRef {
-        self.internal.read().length_hint(|obj| Some(obj.len()), vm)
+    fn length_hint(&self) -> usize {
+        self.internal.read().length_hint(|obj| obj.len())
     }
     #[pymethod(magic)]
     fn reduce(&self, vm: &VirtualMachine) -> PyObjectRef {

vm/src/builtins/bytes.rs

@@ -1,4 +1,4 @@
-use super::{PyDictRef, PyInt, PyIntRef, PyStrRef, PyTupleRef, PyTypeRef, int};
+use super::{PositionIterInternal, PyDictRef, PyIntRef, PyStrRef, PyTupleRef, PyTypeRef};
 use crate::{
     anystr::{self, AnyStr},
     bytesinner::{
@@ -595,8 +595,8 @@ impl PyValue for PyBytesIterator {
 #[pyimpl(with(SlotIterator))]
 impl PyBytesIterator {
     #[pymethod(magic)]
-    fn length_hint(&self, vm: &VirtualMachine) -> PyObjectRef {
-        self.internal.read().length_hint(|obj| Some(obj.len()), vm)
+    fn length_hint(&self) -> usize {
+        self.internal.read().length_hint(|obj| obj.len())
     }
 
     #[pymethod(magic)]

vm/src/builtins/dict.rs

@@ -1,4 +1,4 @@
-use super::{IterStatus, PySet, PyStrRef, PyTypeRef};
+use super::{IterStatus, PositionIterInternal, PySet, PyStrRef, PyTypeRef};
 use crate::{
     builtins::PyBaseExceptionRef,
     common::ascii,
@@ -14,7 +14,6 @@ use crate::{
     PyAttributes, PyClassDef, PyClassImpl, PyComparisonValue, PyContext, PyObjectRef, PyRef,
     PyResult, PyValue, TryFromObject, TypeProtocol,
 };
-use crossbeam_utils::atomic::AtomicCell;
 use rustpython_common::lock::PyRwLock;
 use std::fmt;
 use std::mem::size_of;
@@ -726,10 +725,8 @@ macro_rules! dict_iterator {
             }
 
             #[pymethod(magic)]
-            fn length_hint(&self, vm: &VirtualMachine) -> PyObjectRef {
-                self.internal
-                    .read()
-                    .length_hint(|obj| Some(obj.entries.len()), vm)
+            fn length_hint(&self) -> usize {
+                self.internal.read().length_hint(|_| self.size.entries_size)
             }
         }
 
@@ -786,10 +783,10 @@ macro_rules! dict_iterator {
             }
 
             #[pymethod(magic)]
-            fn length_hint(&self, vm: &VirtualMachine) -> PyObjectRef {
+            fn length_hint(&self) -> usize {
                 self.internal
                     .read()
-                    .rev_length_hint(|_| Some(self.size.entries_size), vm)
+                    .rev_length_hint(|_| self.size.entries_size)
             }
         }
 

vm/src/builtins/enumerate.rs

@@ -1,17 +1,12 @@
-use super::{
-    int,
-    IterStatus::{self, Active, Exhausted},
-    PyInt, PyIntRef, PyTypeRef,
-};
+use super::{IterStatus, PositionIterInternal, PyIntRef, PyTypeRef};
 use crate::common::lock::PyRwLock;
 use crate::{
     function::OptionalArg,
     protocol::{PyIter, PyIterReturn},
     slots::{IteratorIterable, SlotConstructor, SlotIterator},
     IntoPyObject, ItemProtocol, PyClassImpl, PyContext, PyObjectRef, PyRef, PyResult, PyValue,
-    TypeProtocol, VirtualMachine,
+    VirtualMachine,
 };
-use crossbeam_utils::atomic::AtomicCell;
 use num_bigint::BigInt;
 use num_traits::Zero;
 
@@ -91,10 +86,14 @@ impl PyReverseSequenceIterator {
     }
 
     #[pymethod(magic)]
-    fn length_hint(&self, vm: &VirtualMachine) -> PyObjectRef {
-        self.internal
-            .read()
-            .rev_length_hint(|obj| vm.obj_len(obj).ok(), vm)
+    fn length_hint(&self, vm: &VirtualMachine) -> PyResult<usize> {
+        let internal = self.internal.read();
+        if let IterStatus::Active(obj) = &internal.status {
+            if internal.position <= vm.obj_len(obj)? {
+                return Ok(internal.position + 1);
+            }
+        }
+        Ok(0)
     }
 
     #[pymethod(magic)]

vm/src/builtins/iter.rs

@@ -10,7 +10,7 @@ use crate::{
     ItemProtocol, PyClassImpl, PyContext, PyObjectRef, PyRef, PyResult, PyValue, TypeProtocol,
     VirtualMachine,
 };
-use crossbeam_utils::atomic::AtomicCell;
+use rustpython_common::lock::{PyRwLock, PyRwLockUpgradableReadGuard};
 
 /// Marks status of iterator.
 #[derive(Debug, Clone)]
@@ -95,6 +95,10 @@ impl<T> PositionIterInternal<T> {
                     self.status = IterStatus::Exhausted;
                     Err(vm.new_stop_iteration())
                 }
+                Err(e) if e.isinstance(&vm.ctx.exceptions.runtime_error) => {
+                    self.status = IterStatus::Exhausted;
+                    Err(e)
+                }
                 Err(e) => Err(e),
                 Ok(ret) => {
                     self.position += 1;
@@ -120,6 +124,10 @@ impl<T> PositionIterInternal<T> {
                     self.status = IterStatus::Exhausted;
                     Err(vm.new_stop_iteration())
                 }
+                Err(e) if e.isinstance(&vm.ctx.exceptions.runtime_error) => {
+                    self.status = IterStatus::Exhausted;
+                    Err(e)
+                }
                 Err(e) => Err(e),
                 Ok(ret) => {
                     if self.position == 0 {
@@ -135,35 +143,27 @@ impl<T> PositionIterInternal<T> {
         }
     }
 
-    pub fn length_hint<F>(&self, f: F, vm: &VirtualMachine) -> PyObjectRef
+    pub fn length_hint<F>(&self, f: F) -> usize
     where
-        F: FnOnce(&T) -> Option<usize>,
-    {
-        let len = if let IterStatus::Active(obj) = &self.status {
-            if let Some(obj_len) = f(obj) {
-                obj_len.saturating_sub(self.position)
-            } else {
-                return vm.ctx.not_implemented();
-            }
-        } else {
-            0
-        };
-        PyInt::from(len).into_object(vm)
-    }
-
-    pub fn rev_length_hint<F>(&self, f: F, vm: &VirtualMachine) -> PyObjectRef
-    where
-        F: FnOnce(&T) -> Option<usize>,
+        F: FnOnce(&T) -> usize,
     {
         if let IterStatus::Active(obj) = &self.status {
-            if let Some(obj_len) = f(obj) {
-                if self.position <= obj_len {
-                    return PyInt::from(self.position + 1).into_object(vm);
-                }
-            }
-            // FIXME: return NotImplemented?
+            f(obj).saturating_sub(self.position)
+        } else {
+            0
         }
-        PyInt::from(0).into_object(vm)
+    }
+
+    pub fn rev_length_hint<F>(&self, f: F) -> usize
+    where
+        F: FnOnce(&T) -> usize,
+    {
+        if let IterStatus::Active(obj) = &self.status {
+            if self.position <= f(obj) {
+                return self.position + 1;
+            }
+        }
+        0
     }
 }
 
@@ -189,9 +189,14 @@ impl PySequenceIterator {
 
     #[pymethod(magic)]
     fn length_hint(&self, vm: &VirtualMachine) -> PyObjectRef {
-        self.internal
-            .read()
-            .length_hint(|obj| vm.obj_len(obj).ok(), vm)
+        let internal = self.internal.read();
+        if let IterStatus::Active(obj) = &internal.status {
+            vm.obj_len(obj)
+                .map(|x| PyInt::from(x).into_object(vm))
+                .unwrap_or_else(|_| vm.ctx.not_implemented())
+        } else {
+            PyInt::from(0).into_object(vm)
+        }
     }
 
     #[pymethod(magic)]
@@ -240,12 +245,10 @@ impl PyCallableIterator {
 impl IteratorIterable for PyCallableIterator {}
 impl SlotIterator for PyCallableIterator {
     fn next(zelf: &PyRef<Self>, vm: &VirtualMachine) -> PyResult {
-        // let mut status = zelf.status.write();
         let status = zelf.status.upgradable_read();
         if let IterStatus::Active(callable) = &*status {
             let ret = callable.invoke((), vm)?;
             if vm.bool_eq(&ret, &zelf.sentinel)? {
-                // *status = IterStatus::Exhausted;
                 *PyRwLockUpgradableReadGuard::upgrade(status) = IterStatus::Exhausted;
                 Err(vm.new_stop_iteration())
             } else {

vm/src/builtins/list.rs

@@ -1,8 +1,4 @@
-use super::{
-    int,
-    iter::IterStatus::{self, Active, Exhausted},
-    PyGenericAlias, PyInt, PySliceRef, PyTypeRef,
-};
+use super::{PositionIterInternal, PyGenericAlias, PySliceRef, PyTypeRef};
 use crate::common::lock::{
     PyMappedRwLockReadGuard, PyRwLock, PyRwLockReadGuard, PyRwLockWriteGuard,
 };
@@ -19,7 +15,6 @@ use crate::{
     PyClassDef, PyClassImpl, PyComparisonValue, PyContext, PyObjectRef, PyRef, PyResult, PyValue,
     TryFromObject, TypeProtocol,
 };
-use crossbeam_utils::atomic::AtomicCell;
 use std::fmt;
 use std::iter::FromIterator;
 use std::mem::size_of;
@@ -479,8 +474,8 @@ impl PyValue for PyListIterator {
 #[pyimpl(with(SlotIterator))]
 impl PyListIterator {
     #[pymethod(magic)]
-    fn length_hint(&self, vm: &VirtualMachine) -> PyObjectRef {
-        self.internal.read().length_hint(|obj| Some(obj.len()), vm)
+    fn length_hint(&self) -> usize {
+        self.internal.read().length_hint(|obj| obj.len())
     }
 
     #[pymethod(magic)]
@@ -526,10 +521,8 @@ impl PyValue for PyListReverseIterator {
 #[pyimpl(with(SlotIterator))]
 impl PyListReverseIterator {
     #[pymethod(magic)]
-    fn length_hint(&self, vm: &VirtualMachine) -> PyObjectRef {
-        self.internal
-            .read()
-            .rev_length_hint(|obj| Some(obj.len()), vm)
+    fn length_hint(&self) -> usize {
+        self.internal.read().rev_length_hint(|obj| obj.len())
     }
 
     #[pymethod(magic)]

vm/src/builtins/pystr.rs

@@ -1,7 +1,7 @@
 use super::{
-    int::{try_to_primitive, PyInt, PyIntRef},
-    iter::IterStatus::{self, Active, Exhausted},
-    PyBytesRef, PyDict, PyTypeRef,
+    int::{PyInt, PyIntRef},
+    iter::IterStatus::{self, Exhausted},
+    PositionIterInternal, PyBytesRef, PyDict, PyTypeRef,
 };
 use crate::{
     anystr::{self, adjust_indices, AnyStr, AnyStrContainer, AnyStrWrapper},
@@ -181,8 +181,8 @@ impl PyValue for PyStrIterator {
 #[pyimpl(with(SlotIterator))]
 impl PyStrIterator {
     #[pymethod(magic)]
-    fn length_hint(&self, vm: &VirtualMachine) -> PyObjectRef {
-        self.internal.read().length_hint(|obj| Some(obj.len()), vm)
+    fn length_hint(&self) -> usize {
+        self.internal.read().length_hint(|obj| obj.len())
     }
 
     #[pymethod(magic)]

vm/src/builtins/set.rs

@@ -838,10 +838,8 @@ impl PyValue for PySetIterator {
 #[pyimpl(with(SlotIterator))]
 impl PySetIterator {
     #[pymethod(magic)]
-    fn length_hint(&self, vm: &VirtualMachine) -> PyObjectRef {
-        self.internal
-            .read()
-            .length_hint(|_| Some(self.size.entries_size), vm)
+    fn length_hint(&self) -> usize {
+        self.internal.read().length_hint(|_| self.size.entries_size)
     }
 
     #[pymethod(magic)]

vm/src/builtins/tuple.rs

@@ -1,8 +1,4 @@
-use super::{
-    int,
-    iter::IterStatus::{self, Active, Exhausted},
-    PyInt, PyTypeRef,
-};
+use super::{PositionIterInternal, PyTypeRef};
 use crate::common::hash::PyHash;
 use crate::{
     function::OptionalArg,
@@ -19,7 +15,7 @@ use crate::{
     PyContext, PyObjectRef, PyRef, PyResult, PyValue, TransmuteFromObject, TryFromObject,
     TypeProtocol,
 };
-use crossbeam_utils::atomic::AtomicCell;
+use rustpython_common::lock::PyRwLock;
 use std::fmt;
 use std::marker::PhantomData;
 
@@ -331,8 +327,8 @@ impl PyValue for PyTupleIterator {
 #[pyimpl(with(SlotIterator))]
 impl PyTupleIterator {
     #[pymethod(magic)]
-    fn length_hint(&self, vm: &VirtualMachine) -> PyObjectRef {
-        self.internal.read().length_hint(|obj| Some(obj.len()), vm)
+    fn length_hint(&self) -> usize {
+        self.internal.read().length_hint(|obj| obj.len())
     }
 
     #[pymethod(magic)]

vm/src/stdlib/collections.rs

@@ -612,8 +612,8 @@ mod _collections {
         }
 
         #[pymethod(magic)]
-        fn length_hint(&self, vm: &VirtualMachine) -> PyObjectRef {
-            self.internal.read().length_hint(|obj| Some(obj.len()), vm)
+        fn length_hint(&self) -> usize {
+            self.internal.read().length_hint(|obj| obj.len())
         }
 
         #[pymethod(magic)]
@@ -683,10 +683,8 @@ mod _collections {
     #[pyimpl(with(SlotIterator, SlotConstructor))]
     impl PyReverseDequeIterator {
         #[pymethod(magic)]
-        fn length_hint(&self, vm: &VirtualMachine) -> PyObjectRef {
-            self.internal
-                .read()
-                .rev_length_hint(|obj| Some(obj.len()), vm)
+        fn length_hint(&self) -> usize {
+            self.internal.read().rev_length_hint(|obj| obj.len())
         }
 
         #[pymethod(magic)]