Skip to content

Commit 1922725

Browse files
qingshi163qingshi163
committed
Implement array setitem by slice
1 parent 92d8e2c commit 1922725

File tree

3 files changed

+146
-19
lines changed

3 files changed

+146
-19
lines changed

‎Lib/test/test_array.py

Lines changed: 4 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -25,7 +25,10 @@ class ArraySubclassWithKwargs(array.array):
2525
def __init__(self, typecode, newarg=None):
2626
array.array.__init__(self)
2727

28-
typecodes = 'ubBhHiIlLfdqQ'
28+
# TODO: RUSTPYTHON
29+
# We did not support typecode u for unicode yet
30+
# typecodes = 'ubBhHiIlLfdqQ'
31+
typecodes = 'bBhHiIlLfdqQ'
2932

3033
class MiscTest(unittest.TestCase):
3134

@@ -801,8 +804,6 @@ def test_extended_getslice(self):
801804
self.assertEqual(list(a[start:stop:step]),
802805
list(a)[start:stop:step])
803806

804-
# TODO: RUSTPYTHON
805-
@unittest.expectedFailure
806807
def test_setslice(self):
807808
a = array.array(self.typecode, self.example)
808809
a[:1] = a
@@ -1235,8 +1236,6 @@ def test_delslice(self):
12351236
a = array.array(self.typecode, range(10))
12361237
del a[9::1<<333]
12371238

1238-
# TODO: RUSTPYTHON
1239-
@unittest.expectedFailure
12401239
def test_assignment(self):
12411240
a = array.array(self.typecode, range(10))
12421241
a[::2] = array.array(self.typecode, [42]*5)

‎tests/snippets/stdlib_array.py

Lines changed: 15 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -19,4 +19,18 @@
1919
b = a
2020
assert a.__ne__(b) is False
2121
b = array("B", [3, 2, 1, 0])
22-
assert a.__ne__(b) is True
22+
assert a.__ne__(b) is True
23+
24+
# slice assignment step overflow behaviour test
25+
T = 'I'
26+
a = array(T, range(10))
27+
b = array(T, [100])
28+
a[::9999999999] = b
29+
assert a == array(T, [100, 1, 2, 3, 4, 5, 6, 7, 8, 9])
30+
a[::-9999999999] = b
31+
assert a == array(T, [100, 1, 2, 3, 4, 5, 6, 7, 8, 100])
32+
c = array(T)
33+
a[0:0:9999999999] = c
34+
assert a == array(T, [100, 1, 2, 3, 4, 5, 6, 7, 8, 100])
35+
a[0:0:-9999999999] = c
36+
assert a == array(T, [100, 1, 2, 3, 4, 5, 6, 7, 8, 100])

‎vm/src/stdlib/array.rs

Lines changed: 127 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -4,19 +4,21 @@ use crate::common::cell::{
44
};
55
use crate::function::OptionalArg;
66
use crate::obj::objbytes::PyBytesRef;
7-
use crate::obj::objsequence::PySliceableSequence;
7+
use crate::obj::objsequence::{PySliceableSequence, get_slice_range};
88
use crate::obj::objslice::PySliceRef;
99
use crate::obj::objstr::PyStringRef;
1010
use crate::obj::objtype::PyClassRef;
1111
use crate::obj::{objbool, objiter};
1212
use crate::pyobject::{
1313
BorrowValue, Either, IntoPyObject, PyArithmaticValue, PyClassImpl, PyComparisonValue,
14-
PyIterable, PyObject, PyObjectRef, PyRef, PyResult, PyValue, TryFromObject,
14+
PyIterable, PyObject, PyObjectRef, PyRef, PyResult, PyValue, TryFromObject, TypeProtocol
1515
};
1616
use crate::VirtualMachine;
1717
use crossbeam_utils::atomic::AtomicCell;
1818
use itertools::Itertools;
1919
use std::fmt;
20+
use num_bigint::BigInt;
21+
use num_traits::{Zero, One, Signed, ToPrimitive};
2022

2123
struct ArrayTypeSpecifierError {
2224
_priv: (),
@@ -33,7 +35,7 @@ impl fmt::Display for ArrayTypeSpecifierError {
3335

3436
macro_rules! def_array_enum {
3537
($(($n:ident, $t:ident, $c:literal)),*$(,)?) => {
36-
#[derive(Debug)]
38+
#[derive(Debug, Clone)]
3739
enum ArrayContentType {
3840
$($n(Vec<$t>),)*
3941
}
@@ -223,19 +225,115 @@ macro_rules! def_array_enum {
223225
}
224226
}
225227

226-
fn setitem(&mut self, needle: Either<isize, PySliceRef>, value: PyObjectRef, vm: &VirtualMachine) -> PyResult<()> {
227-
match needle {
228-
Either::A(i) => {
229-
let i = self.idx(i, "array assignment", vm)?;
230-
match self {
231-
$(ArrayContentType::$n(v) => { v[i] = TryFromObject::try_from_object(vm, value)? },)*
228+
fn setitem_by_slice(&mut self, slice: PySliceRef, items: ArrayContentType, vm: &VirtualMachine) -> PyResult<()> {
229+
let start = slice.start_index(vm)?;
230+
let stop = slice.stop_index(vm)?;
231+
let step = slice.step_index(vm)?.unwrap_or_else(BigInt::one);
232+
233+
if step.is_zero() {
234+
return Err(vm.new_value_error("slice step cannot be zero".to_owned()));
235+
}
236+
237+
match self {
238+
$(ArrayContentType::$n(elements) => if let ArrayContentType::$n(items) = items {
239+
if step == BigInt::one() {
240+
let range = get_slice_range(&start, &stop, elements.len());
241+
let range = if range.end < range.start {
242+
range.start..range.start
243+
} else {
244+
range
245+
};
246+
elements.splice(range, items);
247+
return Ok(());
232248
}
233-
Ok(())
234-
}
235-
Either::B(_slice) => Err(vm.new_not_implemented_error("array slice is not implemented".to_owned())),
249+
250+
let (start, stop, step, is_negative_step) = if step.is_negative() {
251+
(
252+
stop.map(|x| if x == -BigInt::one() {
253+
elements.len() + BigInt::one()
254+
} else {
255+
x + 1
256+
}),
257+
start.map(|x| if x == -BigInt::one() {
258+
BigInt::from(elements.len())
259+
} else {
260+
x + 1
261+
}),
262+
-step,
263+
true
264+
)
265+
} else {
266+
(start, stop, step, false)
267+
};
268+
269+
let range = get_slice_range(&start, &stop, elements.len());
270+
let range = if range.end < range.start {
271+
range.start..range.start
272+
} else {
273+
range
274+
};
275+
276+
// step is not negative here
277+
if let Some(step) = step.to_usize() {
278+
let slicelen = if range.end > range.start {
279+
(range.end - range.start - 1) / step + 1
280+
} else {
281+
0
282+
};
283+
284+
if slicelen == items.len() {
285+
if is_negative_step {
286+
for (i, item) in range.rev().step_by(step).zip(items) {
287+
elements[i] = item;
288+
}
289+
} else {
290+
for (i, item) in range.step_by(step).zip(items) {
291+
elements[i] = item;
292+
}
293+
}
294+
Ok(())
295+
} else {
296+
Err(vm.new_value_error(format!(
297+
"attempt to assign sequence of size {} to extended slice of size {}",
298+
items.len(), slicelen
299+
)))
300+
}
301+
} else {
302+
// edge case, step is too big for usize
303+
// same behaviour as CPython
304+
let slicelen = if range.start < range.end { 1 } else { 0 };
305+
if match items.len() {
306+
0 => slicelen == 0,
307+
1 => {
308+
elements[
309+
if is_negative_step { range.end - 1 } else { range.start }
310+
] = items[0];
311+
true
312+
},
313+
_ => false,
314+
} {
315+
Ok(())
316+
} else {
317+
Err(vm.new_value_error(format!(
318+
"attempt to assign sequence of size {} to extended slice of size {}",
319+
items.len(), slicelen
320+
)))
321+
}
322+
}
323+
} else {
324+
Err(vm.new_type_error("bad argument type for built-in operation".to_owned()))
325+
},)*
236326
}
237327
}
238328

329+
fn setitem_by_idx(&mut self, i: isize, value: PyObjectRef, vm: &VirtualMachine) -> PyResult<()> {
330+
let i = self.idx(i, "array assignment", vm)?;
331+
match self {
332+
$(ArrayContentType::$n(v) => { v[i] = TryFromObject::try_from_object(vm, value)? },)*
333+
}
334+
Ok(())
335+
}
336+
239337
fn iter<'a>(&'a self, vm: &'a VirtualMachine) -> impl Iterator<Item = PyObjectRef> + 'a {
240338
let mut i = 0;
241339
std::iter::from_fn(move || {
@@ -440,7 +538,23 @@ impl PyArray {
440538
obj: PyObjectRef,
441539
vm: &VirtualMachine,
442540
) -> PyResult<()> {
443-
self.borrow_value_mut().setitem(needle, obj, vm)
541+
match needle {
542+
Either::A(i) => self.borrow_value_mut().setitem_by_idx(i, obj, vm),
543+
Either::B(slice) => {
544+
let items = match obj.payload::<PyArray>() {
545+
// TODO: is there a way to check ref equality between [self] and [obj]
546+
// so we can do clone() only when they are referring the same PyObject
547+
Some(array) => array.borrow_value().clone(),
548+
None => {
549+
return Err(vm.new_type_error(format!(
550+
"can only assign array (not \"{}\") to array slice",
551+
obj.class().name
552+
)));
553+
}
554+
};
555+
self.borrow_value_mut().setitem_by_slice(slice, items, vm)
556+
}
557+
}
444558
}
445559

446560
#[pymethod(name = "__eq__")]

0 commit comments

Comments
 (0)