diff --git a/Lib/test/test_bytes.py b/Lib/test/test_bytes.py index 06c01a0b4..f6a1b3cfd 100644 --- a/Lib/test/test_bytes.py +++ b/Lib/test/test_bytes.py @@ -1522,8 +1522,6 @@ class ByteArrayTest(BaseBytesTest, unittest.TestCase): b.append(Indexable(ord('A'))) self.assertEqual(b, b'A') - # TODO: RUSTPYTHON - @unittest.expectedFailure def test_insert(self): b = bytearray(b'msssspp') b.insert(1, ord('i')) diff --git a/vm/src/bytesinner.rs b/vm/src/bytesinner.rs index 03ec5e1ae..c1b8670e1 100644 --- a/vm/src/bytesinner.rs +++ b/vm/src/bytesinner.rs @@ -11,7 +11,9 @@ use crate::obj::objint::{self, PyInt, PyIntRef}; use crate::obj::objlist::PyList; use crate::obj::objmemory::PyMemoryView; use crate::obj::objnone::PyNoneRef; -use crate::obj::objsequence::{PySliceableSequence, PySliceableSequenceMut, SequenceIndex}; +use crate::obj::objsequence::{ + get_bounded_pos, PySliceableSequence, PySliceableSequenceMut, SequenceIndex, +}; use crate::obj::objslice::PySliceRef; use crate::obj::objstr::{self, PyString, PyStringRef}; use crate::pyobject::{ @@ -422,6 +424,18 @@ impl PyBytesInner { } } + pub fn insert( + &mut self, + index: isize, + object: PyObjectRef, + vm: &VirtualMachine, + ) -> PyResult<()> { + let value = Self::value_try_from_object(vm, object)?; + let index = get_bounded_pos(index, self.len()); + self.elements.insert(index, value); + Ok(()) + } + pub fn isalnum(&self) -> bool { !self.elements.is_empty() && self diff --git a/vm/src/obj/objbytearray.rs b/vm/src/obj/objbytearray.rs index b226704ff..369f36609 100644 --- a/vm/src/obj/objbytearray.rs +++ b/vm/src/obj/objbytearray.rs @@ -531,32 +531,8 @@ impl PyByteArray { } #[pymethod(name = "insert")] - fn insert(&self, mut index: isize, x: PyIntRef, vm: &VirtualMachine) -> PyResult<()> { - let bytes = &mut self.borrow_value_mut().elements; - let len = bytes - .len() - .to_isize() - .ok_or_else(|| vm.new_overflow_error("bytearray too big".to_owned()))?; - - let x = x.borrow_value().byte_or(vm)?; - - if index >= len { - bytes.push(x); - return Ok(()); - } - - if index < 0 { - index += len; - index = index.max(0); - } - - let index = index - .to_usize() - .ok_or_else(|| vm.new_overflow_error("overflow in index calculation".to_owned()))?; - - bytes.insert(index, x); - - Ok(()) + fn insert(&self, index: isize, value: PyObjectRef, vm: &VirtualMachine) -> PyResult<()> { + self.borrow_value_mut().insert(index, value, vm) } #[pymethod(name = "pop")] diff --git a/vm/src/obj/objsequence.rs b/vm/src/obj/objsequence.rs index 65671b41e..3b24c3079 100644 --- a/vm/src/obj/objsequence.rs +++ b/vm/src/obj/objsequence.rs @@ -420,6 +420,25 @@ pub fn get_pos(p: isize, len: usize) -> Option { } } +// return pos is in range [0, len] inclusive +pub fn get_bounded_pos(p: isize, len: usize) -> usize { + if p.is_negative() { + let p = -p as usize; + if p >= len { + 0 + } else { + len - p + } + } else { + let p = p as usize; + if p >= len { + len + } else { + p + } + } +} + pub fn get_slice_pos(slice_pos: &BigInt, len: usize) -> usize { if let Some(pos) = slice_pos.to_isize() { if let Some(index) = get_pos(pos, len) {