Implement array setitem by slice

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)

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
+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])

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<isize, PySliceRef>, 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<Item = PyObjectRef> + '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::<PyArray>() {
+                    // 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__")]