fix vm.to_index now use number protocol

qingshi163 · youknowone · commit 1ed18c012af1 · 2022-05-30T07:36:08.000+09: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
@@ -134,9 +134,11 @@ impl PyNumber<'_> {
                     ),
                     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)) =
@@ -173,9 +175,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)) {
@@ -189,15 +195,23 @@ impl PyNumber<'_> {
                     ),
                     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()
-            )))
-        }
+            ))
+        })
     }
 
     pub fn float_opt(&self, vm: &VirtualMachine) -> PyResult<Option<PyRef<PyFloat>>> {
diff --git a/vm/src/vm/vm_ops.rs b/vm/src/vm/vm_ops.rs
@@ -3,36 +3,18 @@ use crate::{
     builtins::{PyInt, PyIntRef, PyStrInterned},
     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, identifier!(self, __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]
