From eab3609bde7867e9f31d870f9e97b62e36b21717 Mon Sep 17 00:00:00 2001 From: Kangzhi Shi Date: Sat, 25 Sep 2021 11:04:35 +0200 Subject: [PATCH] fix __length_hint__ --- vm/src/builtins/bytearray.rs | 8 +++-- vm/src/builtins/bytes.rs | 6 ++-- vm/src/builtins/dict.rs | 13 +++----- vm/src/builtins/enumerate.rs | 21 ++++++------ vm/src/builtins/iter.rs | 65 +++++++++++++++++++----------------- vm/src/builtins/list.rs | 17 +++------- vm/src/builtins/pystr.rs | 10 +++--- vm/src/builtins/set.rs | 6 ++-- vm/src/builtins/tuple.rs | 12 +++---- vm/src/stdlib/collections.rs | 10 +++--- 10 files changed, 77 insertions(+), 91 deletions(-) diff --git a/vm/src/builtins/bytearray.rs b/vm/src/builtins/bytearray.rs index 8b1a8011c..ecf518352 100644 --- a/vm/src/builtins/bytearray.rs +++ b/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 { diff --git a/vm/src/builtins/bytes.rs b/vm/src/builtins/bytes.rs index 200937a94..32ac819f2 100644 --- a/vm/src/builtins/bytes.rs +++ b/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)] diff --git a/vm/src/builtins/dict.rs b/vm/src/builtins/dict.rs index d4ff3bd86..dd688496c 100644 --- a/vm/src/builtins/dict.rs +++ b/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) } } diff --git a/vm/src/builtins/enumerate.rs b/vm/src/builtins/enumerate.rs index 7ad93f55f..cf6a522d4 100644 --- a/vm/src/builtins/enumerate.rs +++ b/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 { + 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)] diff --git a/vm/src/builtins/iter.rs b/vm/src/builtins/iter.rs index b9e53f62c..32ae0b06c 100644 --- a/vm/src/builtins/iter.rs +++ b/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 PositionIterInternal { 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 PositionIterInternal { 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 PositionIterInternal { } } - pub fn length_hint(&self, f: F, vm: &VirtualMachine) -> PyObjectRef + pub fn length_hint(&self, f: F) -> usize where - F: FnOnce(&T) -> Option, - { - 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(&self, f: F, vm: &VirtualMachine) -> PyObjectRef - where - F: FnOnce(&T) -> Option, + 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(&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, 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 { diff --git a/vm/src/builtins/list.rs b/vm/src/builtins/list.rs index 665658e05..8f9e0a8b6 100644 --- a/vm/src/builtins/list.rs +++ b/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)] diff --git a/vm/src/builtins/pystr.rs b/vm/src/builtins/pystr.rs index 71eb569b2..a8672f25b 100644 --- a/vm/src/builtins/pystr.rs +++ b/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)] diff --git a/vm/src/builtins/set.rs b/vm/src/builtins/set.rs index e57269165..c6047c149 100644 --- a/vm/src/builtins/set.rs +++ b/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)] diff --git a/vm/src/builtins/tuple.rs b/vm/src/builtins/tuple.rs index c8aa6b95e..e23d2cc62 100644 --- a/vm/src/builtins/tuple.rs +++ b/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)] diff --git a/vm/src/stdlib/collections.rs b/vm/src/stdlib/collections.rs index 8809b2a4a..9ae8bac31 100644 --- a/vm/src/stdlib/collections.rs +++ b/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)]