fix vm.to_index now use number protocol

qingshi163 · qingshi163 · commit 40a889049a33 · 2022-05-22T21:00:18.000+02:00
diff --git a/Lib/test/test_index.py b/Lib/test/test_index.py
@@ -71,8 +71,6 @@ def __index__(self):
         self.assertIs(type(direct_index), int)
         #self.assertIs(type(operator_index), int)
 
-    # TODO: RUSTPYTHON
-    @unittest.expectedFailure
     def test_index_returns_int_subclass(self):
         class BadInt:
             def __index__(self):
diff --git a/vm/src/protocol/number.rs b/vm/src/protocol/number.rs
@@ -129,9 +129,11 @@ impl PyNumber<'_> {
                     format!("__int__ returned non-int (type {})", ret.class()),
                     1,
                     vm,
-                )?
+                )?;
+                Ok(vm.ctx.new_int(ret.as_bigint().clone()))
+            } else {
+                Ok(ret)
             }
-            Ok(ret)
         } else if self.methods(vm).index.is_some() {
             self.index(vm)
         } else if let Ok(Ok(f)) = vm.get_special_method(self.obj.to_owned(), "__trunc__") {
@@ -165,9 +167,13 @@ impl PyNumber<'_> {
         }
     }
 
-    pub fn index(&self, vm: &VirtualMachine) -> PyResult<PyIntRef> {
+    pub fn index_opt(&self, vm: &VirtualMachine) -> PyResult<Option<PyIntRef>> {
         if self.obj.class().is(PyInt::class(vm)) {
-            Ok(unsafe { self.obj.to_owned().downcast_unchecked::<PyInt>() })
+            Ok(Some(unsafe {
+                self.obj.to_owned().downcast_unchecked::<PyInt>()
+            }))
+        } else if let Some(i) = self.obj.downcast_ref::<PyInt>() {
+            Ok(Some(i.to_owned()))
         } else if let Some(f) = self.methods(vm).index {
             let ret = f(self, vm)?;
             if !ret.class().is(PyInt::class(vm)) {
@@ -176,15 +182,23 @@ impl PyNumber<'_> {
                     format!("__index__ returned non-int (type {})", ret.class()),
                     1,
                     vm,
-                )?
+                )?;
+                Ok(Some(vm.ctx.new_int(ret.as_bigint().clone())))
+            } else {
+                Ok(Some(ret))
             }
-            Ok(ret)
         } else {
-            Err(vm.new_type_error(format!(
+            Ok(None)
+        }
+    }
+
+    pub fn index(&self, vm: &VirtualMachine) -> PyResult<PyIntRef> {
+        self.index_opt(vm)?.ok_or_else(|| {
+            vm.new_type_error(format!(
                 "'{}' object cannot be interpreted as an integer",
                 self.obj.class()
-            )))
-        }
+            ))
+        })
     }
 }
 
diff --git a/vm/src/types/slot.rs b/vm/src/types/slot.rs
@@ -1,6 +1,6 @@
 use crate::common::{hash::PyHash, lock::PyRwLock};
 use crate::{
-    builtins::{PyInt, PyStrRef, PyType, PyTypeRef},
+    builtins::{PyFloat, PyInt, PyStrRef, PyType, PyTypeRef},
     bytecode::ComparisonOperator,
     convert::{ToPyObject, ToPyResult},
     function::Either,
diff --git a/vm/src/vm/vm_ops.rs b/vm/src/vm/vm_ops.rs
@@ -3,34 +3,18 @@ use crate::{
     builtins::{PyInt, PyIntRef},
     function::PyArithmeticValue,
     object::{AsObject, PyObject, PyObjectRef, PyResult},
-    protocol::PyIterReturn,
+    protocol::{PyIterReturn, PyNumber},
     types::PyComparisonOp,
 };
 
 /// Collection of operators
 impl VirtualMachine {
     pub fn to_index_opt(&self, obj: PyObjectRef) -> Option<PyResult<PyIntRef>> {
-        match obj.downcast() {
-            Ok(val) => Some(Ok(val)),
-            Err(obj) => self.get_method(obj, "__index__").map(|index| {
-                // TODO: returning strict subclasses of int in __index__ is deprecated
-                self.invoke(&index?, ())?.downcast().map_err(|bad| {
-                    self.new_type_error(format!(
-                        "__index__ returned non-int (type {})",
-                        bad.class().name()
-                    ))
-                })
-            }),
-        }
+        PyNumber::from(obj.as_ref()).index_opt(self).transpose()
     }
 
     pub fn to_index(&self, obj: &PyObject) -> PyResult<PyIntRef> {
-        self.to_index_opt(obj.to_owned()).unwrap_or_else(|| {
-            Err(self.new_type_error(format!(
-                "'{}' object cannot be interpreted as an integer",
-                obj.class().name()
-            )))
-        })
+        PyNumber::from(obj).index(self)
     }
 
     #[inline]
