Skip to content

Commit a38965f

Browse files
authored
Merge pull request RustPython#3939 from daeun503/update-array-index
update array.index for optional range
2 parents 86bbfef + e6452d8 commit a38965f

File tree

2 files changed

+24
-12
lines changed

2 files changed

+24
-12
lines changed

Lib/test/test_array.py

Lines changed: 0 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -928,8 +928,6 @@ def test_extended_set_del_slice(self):
928928
del a[start:stop:step]
929929
self.assertEqual(a, array.array(self.typecode, L))
930930

931-
# TODO: RUSTPYTHON
932-
@unittest.expectedFailure
933931
def test_index(self):
934932
example = 2*self.example
935933
a = array.array(self.typecode, example)
@@ -1184,11 +1182,6 @@ class UnicodeTest(StringTest, unittest.TestCase):
11841182
outside = str('\x33')
11851183
minitemsize = 2
11861184

1187-
# TODO: RUSTPYTHON
1188-
@unittest.expectedFailure
1189-
def test_index(self): # XXX: RUSTPYTHON; the method also need to be removed when done
1190-
super().test_index()
1191-
11921185
def test_unicode(self):
11931186
self.assertRaises(TypeError, array.array, 'b', 'foo')
11941187

stdlib/src/array.rs

Lines changed: 24 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -230,12 +230,18 @@ mod array {
230230
}
231231
}
232232

233-
fn index(&self, obj: PyObjectRef, vm: &VirtualMachine) -> PyResult<usize> {
233+
fn index(
234+
&self,
235+
obj: PyObjectRef,
236+
start: usize,
237+
stop: usize,
238+
vm: &VirtualMachine
239+
) -> PyResult<usize> {
234240
match self {
235241
$(ArrayContentType::$n(v) => {
236242
if let Ok(val) = <$t>::try_into_from_object(vm, obj) {
237-
if let Some(pos) = v.iter().position(|&a| a == val) {
238-
return Ok(pos);
243+
if let Some(pos) = v.iter().take(stop as _).skip(start as _).position(|&elem| elem == val) {
244+
return Ok(pos + start);
239245
}
240246
}
241247
Err(vm.new_value_error("array.index(x): x not in array".to_owned()))
@@ -858,8 +864,21 @@ mod array {
858864
}
859865

860866
#[pymethod]
861-
fn index(&self, x: PyObjectRef, vm: &VirtualMachine) -> PyResult<usize> {
862-
self.read().index(x, vm)
867+
fn index(
868+
&self,
869+
x: PyObjectRef,
870+
start: OptionalArg<PyObjectRef>,
871+
stop: OptionalArg<PyObjectRef>,
872+
vm: &VirtualMachine,
873+
) -> PyResult<usize> {
874+
let len = self.len();
875+
let saturate = |obj: PyObjectRef, len| -> PyResult<_> {
876+
obj.try_into_value(vm)
877+
.map(|int: PyIntRef| int.as_bigint().saturated_at(len))
878+
};
879+
let start = start.map_or(Ok(0), |obj| saturate(obj, len))?;
880+
let stop = stop.map_or(Ok(len), |obj| saturate(obj, len))?;
881+
self.read().index(x, start, stop, vm)
863882
}
864883

865884
#[pymethod]

0 commit comments

Comments
 (0)