Merge pull request RustPython#3939 from daeun503/update-array-index

update array.index for optional range

Author: youknowone

Lib/test/test_array.py

@@ -928,8 +928,6 @@ def test_extended_set_del_slice(self):
                     del a[start:stop:step]
                     self.assertEqual(a, array.array(self.typecode, L))
 
-    # TODO: RUSTPYTHON
-    @unittest.expectedFailure
     def test_index(self):
         example = 2*self.example
         a = array.array(self.typecode, example)
@@ -1184,11 +1182,6 @@ class UnicodeTest(StringTest, unittest.TestCase):
     outside = str('\x33')
     minitemsize = 2
 
-    # TODO: RUSTPYTHON
-    @unittest.expectedFailure
-    def test_index(self):  # XXX: RUSTPYTHON; the method also need to be removed when done
-        super().test_index()
-
     def test_unicode(self):
         self.assertRaises(TypeError, array.array, 'b', 'foo')
 

stdlib/src/array.rs

@@ -230,12 +230,18 @@ mod array {
                     }
                 }
 
-                fn index(&self, obj: PyObjectRef, vm: &VirtualMachine) -> PyResult<usize> {
+                fn index(
+                    &self,
+                    obj: PyObjectRef,
+                    start: usize,
+                    stop: usize,
+                    vm: &VirtualMachine
+                ) -> PyResult<usize> {
                     match self {
                         $(ArrayContentType::$n(v) => {
                             if let Ok(val) = <$t>::try_into_from_object(vm, obj) {
-                                if let Some(pos) = v.iter().position(|&a| a == val) {
-                                    return Ok(pos);
+                                if let Some(pos) = v.iter().take(stop as _).skip(start as _).position(|&elem| elem == val) {
+                                    return Ok(pos + start);
                                 }
                             }
                             Err(vm.new_value_error("array.index(x): x not in array".to_owned()))
@@ -858,8 +864,21 @@ mod array {
         }
 
         #[pymethod]
-        fn index(&self, x: PyObjectRef, vm: &VirtualMachine) -> PyResult<usize> {
-            self.read().index(x, vm)
+        fn index(
+            &self,
+            x: PyObjectRef,
+            start: OptionalArg<PyObjectRef>,
+            stop: OptionalArg<PyObjectRef>,
+            vm: &VirtualMachine,
+        ) -> PyResult<usize> {
+            let len = self.len();
+            let saturate = |obj: PyObjectRef, len| -> PyResult<_> {
+                obj.try_into_value(vm)
+                    .map(|int: PyIntRef| int.as_bigint().saturated_at(len))
+            };
+            let start = start.map_or(Ok(0), |obj| saturate(obj, len))?;
+            let stop = stop.map_or(Ok(len), |obj| saturate(obj, len))?;
+            self.read().index(x, start, stop, vm)
         }
 
         #[pymethod]