getitem for bytes/bytearray

Author: youknowone

Lib/test/test_bytes.py

@@ -934,8 +934,6 @@ def test_translate(self):
 class BytesTest(BaseBytesTest, unittest.TestCase):
     type2test = bytes
 
-    # TODO: RUSTPYTHON
-    @unittest.expectedFailure
     def test_getitem_error(self):
         b = b'python'
         msg = "byte indices must be integers or slices"
@@ -1133,8 +1131,6 @@ class BufferBlocked(bytearray):
 class ByteArrayTest(BaseBytesTest, unittest.TestCase):
     type2test = bytearray
 
-    # TODO: RUSTPYTHON
-    @unittest.expectedFailure
     def test_getitem_error(self):
         b = bytearray(b'python')
         msg = "bytearray indices must be integers or slices"

vm/src/bytesinner.rs

@@ -295,19 +295,17 @@ impl PyBytesInner {
         })
     }
 
-    pub fn getitem(&self, needle: SequenceIndex, vm: &VirtualMachine) -> PyResult {
-        match needle {
-            SequenceIndex::Int(int) => {
-                if let Some(idx) = self.elements.wrap_index(int) {
-                    Ok(vm.ctx.new_int(self.elements[idx]))
-                } else {
-                    Err(vm.new_index_error("index out of range".to_owned()))
-                }
-            }
-            SequenceIndex::Slice(slice) => {
-                Ok(vm.ctx.new_bytes(self.elements.get_slice_items(vm, &slice)?))
-            }
-        }
+    pub fn getitem(
+        &self,
+        name: &'static str,
+        needle: PyObjectRef,
+        vm: &VirtualMachine,
+    ) -> PyResult {
+        let obj = match self.elements.get_item(vm, needle, name)? {
+            Either::A(byte) => vm.new_pyobj(byte),
+            Either::B(bytes) => vm.ctx.new_bytes(bytes),
+        };
+        Ok(obj)
     }
 
     pub fn setindex(

vm/src/obj/objbytearray.rs

@@ -152,8 +152,8 @@ impl PyByteArray {
     }
 
     #[pymethod(name = "__getitem__")]
-    fn getitem(&self, needle: SequenceIndex, vm: &VirtualMachine) -> PyResult {
-        self.borrow_value().getitem(needle, vm)
+    fn getitem(&self, needle: PyObjectRef, vm: &VirtualMachine) -> PyResult {
+        self.borrow_value().getitem("bytearray", needle, vm)
     }
 
     #[pymethod(name = "__setitem__")]

vm/src/obj/objbytes.rs

@@ -19,7 +19,6 @@ use crate::pyobject::{
     BorrowValue, Either, IntoPyObject, PyClassImpl, PyComparisonValue, PyContext, PyIterable,
     PyObjectRef, PyRef, PyResult, PyValue, TryFromObject,
 };
-use crate::sliceable::SequenceIndex;
 use crate::slots::{Comparable, Hashable, PyComparisonOp};
 use crate::vm::VirtualMachine;
 use rustpython_common::hash::PyHash;
@@ -139,8 +138,8 @@ impl PyBytes {
     }
 
     #[pymethod(name = "__getitem__")]
-    fn getitem(&self, needle: SequenceIndex, vm: &VirtualMachine) -> PyResult {
-        self.inner.getitem(needle, vm)
+    fn getitem(&self, needle: PyObjectRef, vm: &VirtualMachine) -> PyResult {
+        self.inner.getitem("byte", needle, vm)
     }
 
     #[pymethod(name = "isalnum")]