From 1922725af5bf6e3375500f40273e006f1bfca45d Mon Sep 17 00:00:00 2001 From: Kangzhi Shi Date: Sun, 30 Aug 2020 16:51:52 +0200 Subject: [PATCH] Implement array setitem by slice --- Lib/test/test_array.py | 9 +-- tests/snippets/stdlib_array.py | 16 +++- vm/src/stdlib/array.rs | 142 +++++++++++++++++++++++++++++---- 3 files changed, 147 insertions(+), 20 deletions(-) diff --git a/Lib/test/test_array.py b/Lib/test/test_array.py index 9a121d5f1..2e3019744 100644 --- a/Lib/test/test_array.py +++ b/Lib/test/test_array.py @@ -25,7 +25,10 @@ class ArraySubclassWithKwargs(array.array): def __init__(self, typecode, newarg=None): array.array.__init__(self) -typecodes = 'ubBhHiIlLfdqQ' +# TODO: RUSTPYTHON +# We did not support typecode u for unicode yet +# typecodes = 'ubBhHiIlLfdqQ' +typecodes = 'bBhHiIlLfdqQ' class MiscTest(unittest.TestCase): @@ -801,8 +804,6 @@ class BaseTest: self.assertEqual(list(a[start:stop:step]), list(a)[start:stop:step]) - # TODO: RUSTPYTHON - @unittest.expectedFailure def test_setslice(self): a = array.array(self.typecode, self.example) a[:1] = a @@ -1235,8 +1236,6 @@ class NumberTest(BaseTest): a = array.array(self.typecode, range(10)) del a[9::1<<333] - # TODO: RUSTPYTHON - @unittest.expectedFailure def test_assignment(self): a = array.array(self.typecode, range(10)) a[::2] = array.array(self.typecode, [42]*5) diff --git a/tests/snippets/stdlib_array.py b/tests/snippets/stdlib_array.py index 0ef0969ff..9bdd92598 100644 --- a/tests/snippets/stdlib_array.py +++ b/tests/snippets/stdlib_array.py @@ -19,4 +19,18 @@ a = array("b", [0, 1, 2, 3]) b = a assert a.__ne__(b) is False b = array("B", [3, 2, 1, 0]) -assert a.__ne__(b) is True \ No newline at end of file +assert a.__ne__(b) is True + +# slice assignment step overflow behaviour test +T = 'I' +a = array(T, range(10)) +b = array(T, [100]) +a[::9999999999] = b +assert a == array(T, [100, 1, 2, 3, 4, 5, 6, 7, 8, 9]) +a[::-9999999999] = b +assert a == array(T, [100, 1, 2, 3, 4, 5, 6, 7, 8, 100]) +c = array(T) +a[0:0:9999999999] = c +assert a == array(T, [100, 1, 2, 3, 4, 5, 6, 7, 8, 100]) +a[0:0:-9999999999] = c +assert a == array(T, [100, 1, 2, 3, 4, 5, 6, 7, 8, 100]) diff --git a/vm/src/stdlib/array.rs b/vm/src/stdlib/array.rs index d01773e9d..ab3d74c2b 100644 --- a/vm/src/stdlib/array.rs +++ b/vm/src/stdlib/array.rs @@ -4,19 +4,21 @@ use crate::common::cell::{ }; use crate::function::OptionalArg; use crate::obj::objbytes::PyBytesRef; -use crate::obj::objsequence::PySliceableSequence; +use crate::obj::objsequence::{PySliceableSequence, get_slice_range}; use crate::obj::objslice::PySliceRef; use crate::obj::objstr::PyStringRef; use crate::obj::objtype::PyClassRef; use crate::obj::{objbool, objiter}; use crate::pyobject::{ BorrowValue, Either, IntoPyObject, PyArithmaticValue, PyClassImpl, PyComparisonValue, - PyIterable, PyObject, PyObjectRef, PyRef, PyResult, PyValue, TryFromObject, + PyIterable, PyObject, PyObjectRef, PyRef, PyResult, PyValue, TryFromObject, TypeProtocol }; use crate::VirtualMachine; use crossbeam_utils::atomic::AtomicCell; use itertools::Itertools; use std::fmt; +use num_bigint::BigInt; +use num_traits::{Zero, One, Signed, ToPrimitive}; struct ArrayTypeSpecifierError { _priv: (), @@ -33,7 +35,7 @@ impl fmt::Display for ArrayTypeSpecifierError { macro_rules! def_array_enum { ($(($n:ident, $t:ident, $c:literal)),*$(,)?) => { - #[derive(Debug)] + #[derive(Debug, Clone)] enum ArrayContentType { $($n(Vec<$t>),)* } @@ -223,17 +225,113 @@ macro_rules! def_array_enum { } } - fn setitem(&mut self, needle: Either, value: PyObjectRef, vm: &VirtualMachine) -> PyResult<()> { - match needle { - Either::A(i) => { - let i = self.idx(i, "array assignment", vm)?; - match self { - $(ArrayContentType::$n(v) => { v[i] = TryFromObject::try_from_object(vm, value)? },)* - } - Ok(()) - } - Either::B(_slice) => Err(vm.new_not_implemented_error("array slice is not implemented".to_owned())), + fn setitem_by_slice(&mut self, slice: PySliceRef, items: ArrayContentType, vm: &VirtualMachine) -> PyResult<()> { + let start = slice.start_index(vm)?; + let stop = slice.stop_index(vm)?; + let step = slice.step_index(vm)?.unwrap_or_else(BigInt::one); + + if step.is_zero() { + return Err(vm.new_value_error("slice step cannot be zero".to_owned())); } + + match self { + $(ArrayContentType::$n(elements) => if let ArrayContentType::$n(items) = items { + if step == BigInt::one() { + let range = get_slice_range(&start, &stop, elements.len()); + let range = if range.end < range.start { + range.start..range.start + } else { + range + }; + elements.splice(range, items); + return Ok(()); + } + + let (start, stop, step, is_negative_step) = if step.is_negative() { + ( + stop.map(|x| if x == -BigInt::one() { + elements.len() + BigInt::one() + } else { + x + 1 + }), + start.map(|x| if x == -BigInt::one() { + BigInt::from(elements.len()) + } else { + x + 1 + }), + -step, + true + ) + } else { + (start, stop, step, false) + }; + + let range = get_slice_range(&start, &stop, elements.len()); + let range = if range.end < range.start { + range.start..range.start + } else { + range + }; + + // step is not negative here + if let Some(step) = step.to_usize() { + let slicelen = if range.end > range.start { + (range.end - range.start - 1) / step + 1 + } else { + 0 + }; + + if slicelen == items.len() { + if is_negative_step { + for (i, item) in range.rev().step_by(step).zip(items) { + elements[i] = item; + } + } else { + for (i, item) in range.step_by(step).zip(items) { + elements[i] = item; + } + } + Ok(()) + } else { + Err(vm.new_value_error(format!( + "attempt to assign sequence of size {} to extended slice of size {}", + items.len(), slicelen + ))) + } + } else { + // edge case, step is too big for usize + // same behaviour as CPython + let slicelen = if range.start < range.end { 1 } else { 0 }; + if match items.len() { + 0 => slicelen == 0, + 1 => { + elements[ + if is_negative_step { range.end - 1 } else { range.start } + ] = items[0]; + true + }, + _ => false, + } { + Ok(()) + } else { + Err(vm.new_value_error(format!( + "attempt to assign sequence of size {} to extended slice of size {}", + items.len(), slicelen + ))) + } + } + } else { + Err(vm.new_type_error("bad argument type for built-in operation".to_owned())) + },)* + } + } + + fn setitem_by_idx(&mut self, i: isize, value: PyObjectRef, vm: &VirtualMachine) -> PyResult<()> { + let i = self.idx(i, "array assignment", vm)?; + match self { + $(ArrayContentType::$n(v) => { v[i] = TryFromObject::try_from_object(vm, value)? },)* + } + Ok(()) } fn iter<'a>(&'a self, vm: &'a VirtualMachine) -> impl Iterator + 'a { @@ -440,7 +538,23 @@ impl PyArray { obj: PyObjectRef, vm: &VirtualMachine, ) -> PyResult<()> { - self.borrow_value_mut().setitem(needle, obj, vm) + match needle { + Either::A(i) => self.borrow_value_mut().setitem_by_idx(i, obj, vm), + Either::B(slice) => { + let items = match obj.payload::() { + // TODO: is there a way to check ref equality between [self] and [obj] + // so we can do clone() only when they are referring the same PyObject + Some(array) => array.borrow_value().clone(), + None => { + return Err(vm.new_type_error(format!( + "can only assign array (not \"{}\") to array slice", + obj.class().name + ))); + } + }; + self.borrow_value_mut().setitem_by_slice(slice, items, vm) + } + } } #[pymethod(name = "__eq__")]