diff --git a/vm/src/obj/objlist.rs b/vm/src/obj/objlist.rs index 28d09cd0b..ec0e0dda6 100644 --- a/vm/src/obj/objlist.rs +++ b/vm/src/obj/objlist.rs @@ -17,7 +17,7 @@ use crate::bytesinner; use crate::common::cell::{PyRwLock, PyRwLockReadGuard, PyRwLockWriteGuard}; use crate::function::OptionalArg; use crate::pyobject::{ - BorrowValue, PyClassImpl, PyComparisonValue, PyContext, PyIterable, PyObjectRef, PyRef, + BorrowValue, Either, PyClassImpl, PyComparisonValue, PyContext, PyIterable, PyObjectRef, PyRef, PyResult, PyValue, TryFromObject, TypeProtocol, }; use crate::sequence::{self, SimpleSeq}; @@ -192,7 +192,12 @@ impl PyList { #[pymethod(name = "__getitem__")] fn getitem(zelf: PyRef, needle: PyObjectRef, vm: &VirtualMachine) -> PyResult { - get_item(vm, zelf.as_object(), &zelf.borrow_value(), needle) + Ok( + match get_item(vm, zelf.as_object(), &zelf.borrow_value(), needle)? { + Either::A(obj) => obj, + Either::B(vec) => vm.ctx.new_list(vec), + }, + ) } #[pymethod(name = "__iter__")] diff --git a/vm/src/obj/objsequence.rs b/vm/src/obj/objsequence.rs index 273b7dd3f..9ef4d2c0f 100644 --- a/vm/src/obj/objsequence.rs +++ b/vm/src/obj/objsequence.rs @@ -4,12 +4,10 @@ use num_bigint::BigInt; use num_traits::{One, Signed, ToPrimitive, Zero}; use super::objint::{PyInt, PyIntRef}; -use super::objlist::PyList; use super::objsingletons::PyNone; use super::objslice::{PySlice, PySliceRef}; -use super::objtuple::PyTuple; use crate::function::OptionalArg; -use crate::pyobject::{BorrowValue, PyObject, PyObjectRef, PyResult, TryFromObject, TypeProtocol}; +use crate::pyobject::{BorrowValue, Either, PyObjectRef, PyResult, TryFromObject, TypeProtocol}; use crate::vm::VirtualMachine; pub trait PySliceableSequenceMut { @@ -465,34 +463,23 @@ pub fn get_item( sequence: &PyObjectRef, elements: &[PyObjectRef], subscript: PyObjectRef, -) -> PyResult { +) -> PyResult>> { if let Some(i) = subscript.payload::() { let value = i.borrow_value().to_isize().ok_or_else(|| { vm.new_index_error("cannot fit 'int' into an index-sized integer".to_owned()) })?; let pos_index = get_pos(value, elements.len()) .ok_or_else(|| vm.new_index_error("Index out of bounds!".to_owned()))?; - return Ok(elements[pos_index].clone()); - } - - let slice = subscript.payload::().ok_or_else(|| { - vm.new_type_error(format!( - "{} indices must be integers or slices", - sequence.lease_class().name - )) - })?; - let items = if sequence.payload_is::() { - PyObject::new( - PyList::from(elements.get_slice_items(vm, slice)?), - sequence.class(), - None, - ) - } else if sequence.payload_is::() { - vm.ctx.new_tuple(elements.get_slice_items(vm, slice)?) + Ok(Either::A(elements[pos_index].clone())) } else { - panic!("sequence get_item called for non-sequence") - }; - Ok(items) + let slice = subscript.payload::().ok_or_else(|| { + vm.new_type_error(format!( + "{} indices must be integers or slices", + sequence.lease_class().name + )) + })?; + Ok(Either::B(elements.get_slice_items(vm, slice)?)) + } } //Check if given arg could be used with PySliceableSequence.get_slice_range() diff --git a/vm/src/obj/objtuple.rs b/vm/src/obj/objtuple.rs index d2f3f8d03..67befce63 100644 --- a/vm/src/obj/objtuple.rs +++ b/vm/src/obj/objtuple.rs @@ -6,7 +6,7 @@ use super::objsequence::get_item; use super::objtype::PyClassRef; use crate::function::OptionalArg; use crate::pyobject::{ - self, BorrowValue, IdProtocol, IntoPyObject, + self, BorrowValue, Either, IdProtocol, IntoPyObject, PyArithmaticValue::{self, *}, PyClassImpl, PyComparisonValue, PyContext, PyObjectRef, PyRef, PyResult, PyValue, }; @@ -168,7 +168,12 @@ impl PyTuple { #[pymethod(name = "__getitem__")] fn getitem(zelf: PyRef, needle: PyObjectRef, vm: &VirtualMachine) -> PyResult { - get_item(vm, zelf.as_object(), &zelf.elements, needle) + Ok( + match get_item(vm, zelf.as_object(), &zelf.elements, needle)? { + Either::A(obj) => obj, + Either::B(vec) => vm.ctx.new_tuple(vec), + }, + ) } #[pymethod(name = "index")]