refactor length now use sequence and mapping protocol

qingshi163 · qingshi163 · commit 2a8373cdbf46 · 2021-11-16T22:06:05.000+02:00
diff --git a/vm/src/protocol/mapping.rs b/vm/src/protocol/mapping.rs
@@ -43,6 +43,17 @@ impl PyMapping<PyObjectRef> {
         }
         PyMappingMethods::default()
     }
+
+    pub fn length(&self, vm: &VirtualMachine) -> PyResult<usize> {
+        if let Some(f) = self.methods(vm).length {
+            f(self.0.clone(), vm)
+        } else {
+            Err(vm.new_type_error(format!(
+                "object of type '{}' has no len() or not a mapping",
+                self.0.class().name()
+            )))
+        }
+    }
 }
 
 impl<T> PyMapping<T>
diff --git a/vm/src/protocol/object.rs b/vm/src/protocol/object.rs
@@ -6,7 +6,7 @@ use crate::{
     bytesinner::ByteInnerNewOptions,
     common::{hash::PyHash, str::to_ascii},
     function::{IntoPyObject, OptionalArg},
-    protocol::PyIter,
+    protocol::{PyIter, PyMapping, PySequence},
     pyref_type_error,
     types::{Constructor, PyComparisonOp},
     utils::Either,
@@ -404,11 +404,11 @@ impl PyObject {
     // int PyObject_TypeCheck(PyObject *o, PyTypeObject *type)
 
     pub fn length(&self, vm: &VirtualMachine) -> PyResult<usize> {
-        vm.obj_len_opt(self).unwrap_or_else(|| {
-            Err(vm.new_type_error(format!(
-                "object of type '{}' has no len()",
-                self.class().name()
-            )))
-        })
+        let seq = PySequence::from(self);
+        if let Ok(len) = seq.length(vm) {
+            Ok(len)
+        } else {
+            PyMapping::try_from_object(vm, self.to_owned())?.length(vm)
+        }
     }
 }
diff --git a/vm/src/protocol/sequence.rs b/vm/src/protocol/sequence.rs
@@ -297,7 +297,7 @@ impl PySequence<'_> {
         if let Some(tuple) = self.obj.downcast_ref_if_exact::<PyTuple>(vm) {
             Ok(tuple.to_owned())
         } else if let Some(list) = self.obj.downcast_ref_if_exact::<PyList>(vm) {
-            Ok(vm.ctx.new_tuple(list.borrow_vec().to_vec()).into())
+            Ok(vm.ctx.new_tuple(list.borrow_vec().to_vec()))
         } else {
             let iter = self.obj.to_owned().get_iter(vm)?;
             let iter = iter.iter(vm)?;
diff --git a/vm/src/types/slot.rs b/vm/src/types/slot.rs
@@ -163,7 +163,7 @@ macro_rules! then_some_closure {
 }
 
 fn slot_length(obj: PyObjectRef, vm: &VirtualMachine) -> PyResult<usize> {
-    let ret = vm.call_special_method(obj.to_owned(), "__len__", ())?;
+    let ret = vm.call_special_method(obj, "__len__", ())?;
     let len = ret.payload::<PyInt>().ok_or_else(|| {
         vm.new_type_error(format!(
             "'{}' object cannot be interpreted as an integer",
diff --git a/vm/src/vm.rs b/vm/src/vm.rs
@@ -31,7 +31,7 @@ use crate::{
     TypeProtocol,
 };
 use crossbeam_utils::atomic::AtomicCell;
-use num_traits::{Signed, ToPrimitive};
+use num_traits::ToPrimitive;
 use std::borrow::Cow;
 use std::cell::{Cell, Ref, RefCell};
 use std::collections::{HashMap, HashSet};
@@ -1673,41 +1673,12 @@ impl VirtualMachine {
             .invoke((), self)
     }
 
-    pub fn obj_len_opt(&self, obj: &PyObject) -> Option<PyResult<usize>> {
-        self.get_special_method(obj.to_owned(), "__len__")
-            .map(Result::ok)
-            .transpose()
-            .map(|meth| {
-                let len = meth?.invoke((), self)?;
-                let len = len
-                    .payload_if_subclass::<PyInt>(self)
-                    .ok_or_else(|| {
-                        self.new_type_error(format!(
-                            "'{}' object cannot be interpreted as an integer",
-                            len.class().name()
-                        ))
-                    })?
-                    .as_bigint();
-                if len.is_negative() {
-                    return Err(self.new_value_error("__len__() should return >= 0".to_owned()));
-                }
-                let len = len.to_isize().ok_or_else(|| {
-                    self.new_overflow_error(
-                        "cannot fit 'int' into an index-sized integer".to_owned(),
-                    )
-                })?;
-                Ok(len as usize)
-            })
-    }
-
     pub fn length_hint_opt(&self, iter: PyObjectRef) -> PyResult<Option<usize>> {
-        if let Some(len) = self.obj_len_opt(&iter) {
-            match len {
-                Ok(len) => return Ok(Some(len)),
-                Err(e) => {
-                    if !e.isinstance(&self.ctx.exceptions.type_error) {
-                        return Err(e);
-                    }
+        match iter.length(self) {
+            Ok(len) => return Ok(Some(len)),
+            Err(e) => {
+                if !e.isinstance(&self.ctx.exceptions.type_error) {
+                    return Err(e);
                 }
             }
         }

Original file line number	Diff line number	Diff line change
`@@ -43,6 +43,17 @@ impl PyMapping<PyObjectRef> {`
`43`	`43`	`}`
`44`	`44`	`PyMappingMethods::default()`
`45`	`45`	`}`
	`46`	`+`
	`47`	`+ pub fn length(&self, vm: &VirtualMachine) -> PyResult<usize> {`
	`48`	`+ if let Some(f) = self.methods(vm).length {`
	`49`	`+ f(self.0.clone(), vm)`
	`50`	`+ } else {`
	`51`	`+ Err(vm.new_type_error(format!(`
	`52`	`+ "object of type '{}' has no len() or not a mapping",`
	`53`	`+ self.0.class().name()`
	`54`	`+ )))`
	`55`	`+ }`
	`56`	`+ }`
`46`	`57`	`}`
`47`	`58`
`48`	`59`	`impl<T> PyMapping<T>`
Original file line number	Diff line number	Diff line change
`@@ -163,7 +163,7 @@ macro_rules! then_some_closure {`
`163`	`163`	`}`
`164`	`164`
`165`	`165`	`fn slot_length(obj: PyObjectRef, vm: &VirtualMachine) -> PyResult<usize> {`
`166`		`- let ret = vm.call_special_method(obj.to_owned(), "__len__", ())?;`
	`166`	`+ let ret = vm.call_special_method(obj, "__len__", ())?;`
`167`	`167`	`let len = ret.payload::<PyInt>().ok_or_else(\|\| {`
`168`	`168`	`vm.new_type_error(format!(`
`169`	`169`	`"'{}' object cannot be interpreted as an integer",`