Skip to content

Commit 275b727

Browse files
qingshi163qingshi163
committed
memoryview __setitem__
1 parent cf35723 commit 275b727

File tree

4 files changed

+174
-35
lines changed

4 files changed

+174
-35
lines changed

‎Lib/test/test_memoryview.py

Lines changed: 11 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -71,8 +71,6 @@ def setitem(value):
7171
m = None
7272
self.assertEqual(sys.getrefcount(b), oldrefcount)
7373

74-
# TODO: RUSTPYTHON
75-
@unittest.expectedFailure
7674
def test_setitem_writable(self):
7775
if not self.rw_type:
7876
self.skipTest("no writable type to test")
@@ -114,11 +112,13 @@ def setitem(key, value):
114112
self.assertRaises(TypeError, setitem, "a", b"a")
115113
# Not implemented: multidimensional slices
116114
slices = (slice(0,1,1), slice(0,1,2))
117-
self.assertRaises(NotImplementedError, setitem, slices, b"a")
115+
# TODO: RUSTPYTHON
116+
# self.assertRaises(NotImplementedError, setitem, slices, b"a")
118117
# Trying to resize the memory object
119118
exc = ValueError if m.format == 'c' else TypeError
120-
self.assertRaises(exc, setitem, 0, b"")
121-
self.assertRaises(exc, setitem, 0, b"ab")
119+
# TODO: RUSTPYTHON
120+
# self.assertRaises(exc, setitem, 0, b"")
121+
# self.assertRaises(exc, setitem, 0, b"ab")
122122
self.assertRaises(ValueError, setitem, slice(1,1), b"a")
123123
self.assertRaises(ValueError, setitem, slice(0,2), b"a")
124124

@@ -272,11 +272,12 @@ def _check_released(self, m, tp):
272272
with check: m.itemsize
273273
with check: m.ndim
274274
with check: m.readonly
275-
with check: m.shape
276-
with check: m.strides
277-
with check:
278-
with m:
279-
pass
275+
# TODO: RUSTPYTHON
276+
# with check: m.shape
277+
# with check: m.strides
278+
# with check:
279+
# with m:
280+
# pass
280281
# str() and repr() still function
281282
self.assertIn("released memory", str(m))
282283
self.assertIn("released memory", repr(m))
@@ -298,8 +299,6 @@ def test_contextmanager(self):
298299
with m:
299300
m.release()
300301

301-
# TODO: RUSTPYTHON
302-
@unittest.expectedFailure
303302
def test_release(self):
304303
for tp in self._types:
305304
b = tp(self._source)
@@ -435,11 +434,9 @@ class BaseArrayMemoryTests(AbstractMemoryTests):
435434
itemsize = array.array('i').itemsize
436435
format = 'i'
437436

438-
@unittest.skip('XXX test should be adapted for non-byte buffers')
439437
def test_getbuffer(self):
440438
pass
441439

442-
@unittest.skip('XXX NotImplementedError: tolist() only supports byte views')
443440
def test_tolist(self):
444441
pass
445442

@@ -503,7 +500,6 @@ def test_constructor(self):
503500
class ArrayMemoryviewTest(unittest.TestCase,
504501
BaseMemoryviewTests, BaseArrayMemoryTests):
505502

506-
@unittest.skip("TODO: RUSTPYTHON")
507503
def test_array_assign(self):
508504
# Issue #4569: segfault when mutating a memoryview with itemsize != 1
509505
a = array.array('i', range(10))

‎vm/src/obj/objmemory.rs

Lines changed: 156 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -1,23 +1,24 @@
1-
use super::objtype::PyTypeRef;
21
use std::{fmt::Debug, ops::Deref};
32

3+
use crate::common::borrow::{BorrowedValue, BorrowedValueMut};
4+
use crate::common::hash::PyHash;
45
use crate::obj::objbytes::{PyBytes, PyBytesRef};
56
use crate::obj::objlist::{PyList, PyListRef};
6-
use crate::obj::{objslice::PySliceRef, objstr::PyStr};
7+
use crate::obj::objslice::PySliceRef;
8+
use crate::obj::objstr::PyStr;
9+
use crate::obj::objtype::PyTypeRef;
710
use crate::pyobject::{
8-
IdProtocol, PyClassImpl, PyContext, PyObjectRef, PyRef, PyResult, PyThreadingConstraint,
9-
PyValue, TypeProtocol,
11+
IdProtocol, PyClassImpl, PyComparisonValue, PyContext, PyObjectRef, PyRef, PyResult,
12+
PyThreadingConstraint, PyValue, TypeProtocol,
1013
};
11-
use crate::sliceable::{saturate_range, wrap_index, SequenceIndex};
14+
use crate::sliceable::{convert_slice, saturate_range, wrap_index, SequenceIndex};
1215
use crate::slots::{BufferProtocol, Comparable, Hashable, PyComparisonOp};
1316
use crate::stdlib::pystruct::_struct::FormatSpec;
1417
use crate::VirtualMachine;
15-
use crate::{common::hash::PyHash, pyobject::PyComparisonValue};
1618
use crossbeam_utils::atomic::AtomicCell;
1719
use itertools::Itertools;
1820
use num_bigint::BigInt;
1921
use num_traits::{One, Signed, ToPrimitive, Zero};
20-
use rustpython_common::borrow::{BorrowedValue, BorrowedValueMut};
2122

2223
#[derive(Debug)]
2324
pub struct BufferRef(Box<dyn Buffer>);
@@ -62,6 +63,10 @@ pub trait Buffer: Debug + PyThreadingConstraint {
6263
Some(self.obj_bytes_mut())
6364
}
6465

66+
fn to_contiguous(&self) -> Vec<u8> {
67+
self.obj_bytes().to_vec()
68+
}
69+
6570
fn try_resizable(&self, vm: &VirtualMachine) -> PyResult<()> {
6671
if self.is_resizable() {
6772
Ok(())
@@ -371,26 +376,154 @@ impl PyMemoryView {
371376
}
372377
}
373378

379+
fn setitem_by_idx(
380+
zelf: PyRef<Self>,
381+
i: isize,
382+
value: PyObjectRef,
383+
vm: &VirtualMachine,
384+
) -> PyResult<()> {
385+
let i = zelf
386+
.get_pos(i)
387+
.ok_or_else(|| vm.new_index_error("index out of range".to_owned()))?;
388+
let itemsize = zelf.options.itemsize;
389+
let data = zelf.format_spec.pack(&[value], vm)?;
390+
zelf.obj_bytes_mut()[i..i + itemsize].copy_from_slice(&data);
391+
Ok(())
392+
}
393+
394+
fn setitem_by_slice(
395+
zelf: PyRef<Self>,
396+
slice: PySliceRef,
397+
items: PyObjectRef,
398+
vm: &VirtualMachine,
399+
) -> PyResult<()> {
400+
let items = try_buffer_from_object(vm, &items)?;
401+
let options = items.get_options();
402+
let len = options.len;
403+
let itemsize = options.itemsize;
404+
405+
if itemsize != zelf.options.itemsize {
406+
return Err(vm.new_type_error(format!(
407+
"memoryview: invalid type for format '{}'",
408+
zelf.options.format
409+
)));
410+
}
411+
412+
let diff_err = || {
413+
Err(vm.new_value_error(
414+
"memoryview assignment: lvalue and rvalue have different structures".to_owned(),
415+
))
416+
};
417+
418+
if options.format != zelf.options.format {
419+
return diff_err();
420+
}
421+
422+
let (range, step, is_negative_step) = convert_slice(&slice, zelf.options.len, vm)?;
423+
424+
let bytes = items.to_contiguous();
425+
assert_eq!(bytes.len(), len * itemsize);
426+
427+
if !is_negative_step && step == Some(1) {
428+
if range.end - range.start != len {
429+
return diff_err();
430+
}
431+
432+
if let Some(mut buffer) = zelf.as_contiguous_mut() {
433+
buffer[range.start * itemsize..range.end * itemsize].copy_from_slice(&bytes);
434+
return Ok(());
435+
}
436+
}
437+
438+
if let Some(step) = step {
439+
let slicelen = if range.end > range.start {
440+
(range.end - range.start - 1) / step + 1
441+
} else {
442+
0
443+
};
444+
445+
if slicelen != len {
446+
return diff_err();
447+
}
448+
449+
let indexes = if is_negative_step {
450+
itertools::Either::Left(range.rev().step_by(step))
451+
} else {
452+
itertools::Either::Right(range.step_by(step))
453+
};
454+
455+
let item_index = (0..len).step_by(itemsize);
456+
457+
let mut buffer = zelf.obj_bytes_mut();
458+
459+
indexes
460+
.map(|i| zelf.get_pos(i as isize).unwrap())
461+
.zip(item_index)
462+
.for_each(|(i, item_i)| {
463+
buffer[i..i + itemsize].copy_from_slice(&bytes[item_i..item_i + itemsize]);
464+
});
465+
Ok(())
466+
} else {
467+
let slicelen = if range.start < range.end { 1 } else { 0 };
468+
if match len {
469+
0 => slicelen == 0,
470+
1 => {
471+
let mut buffer = zelf.obj_bytes_mut();
472+
let i = zelf.get_pos(range.start as isize).unwrap();
473+
buffer[i..i + itemsize].copy_from_slice(&bytes);
474+
true
475+
}
476+
_ => false,
477+
} {
478+
Ok(())
479+
} else {
480+
diff_err()
481+
}
482+
}
483+
}
484+
485+
#[pymethod(magic)]
486+
fn setitem(
487+
zelf: PyRef<Self>,
488+
needle: SequenceIndex,
489+
value: PyObjectRef,
490+
vm: &VirtualMachine,
491+
) -> PyResult<()> {
492+
zelf.try_not_released(vm)?;
493+
if zelf.options.readonly {
494+
return Err(vm.new_type_error("cannot modify read-only memory".to_owned()));
495+
}
496+
match needle {
497+
SequenceIndex::Int(i) => Self::setitem_by_idx(zelf, i, value, vm),
498+
SequenceIndex::Slice(slice) => Self::setitem_by_slice(zelf, slice, value, vm),
499+
}
500+
}
501+
374502
#[pymethod(magic)]
375503
fn len(&self, vm: &VirtualMachine) -> PyResult<usize> {
376504
self.try_not_released(vm).map(|_| self.options.len)
377505
}
378506

379-
#[pymethod]
380-
fn tobytes(zelf: PyRef<Self>, vm: &VirtualMachine) -> PyResult<PyBytesRef> {
381-
zelf.try_not_released(vm)?;
507+
fn to_bytes_vec(zelf: &PyRef<Self>) -> Vec<u8> {
382508
if let Some(bytes) = zelf.as_contiguous() {
383-
Ok(PyBytes::from(bytes.to_vec()).into_ref(vm))
509+
bytes.to_vec()
384510
} else {
385511
let bytes = &*zelf.obj_bytes();
386-
let bytes = (0..zelf.options.len)
512+
let len = zelf.options.len;
513+
let itemsize = zelf.options.itemsize;
514+
(0..len)
387515
.map(|i| zelf.get_pos(i as isize).unwrap())
388-
.flat_map(|i| (i..i + zelf.options.itemsize).map(|i| bytes[i]))
389-
.collect::<Vec<u8>>();
390-
Ok(PyBytes::from(bytes).into_ref(vm))
516+
.flat_map(|i| bytes[i..i + itemsize].to_vec())
517+
.collect()
391518
}
392519
}
393520

521+
#[pymethod]
522+
fn tobytes(zelf: PyRef<Self>, vm: &VirtualMachine) -> PyResult<PyBytesRef> {
523+
zelf.try_not_released(vm)?;
524+
Ok(PyBytes::from(Self::to_bytes_vec(&zelf)).into_ref(vm))
525+
}
526+
394527
#[pymethod]
395528
fn tolist(zelf: PyRef<Self>, vm: &VirtualMachine) -> PyResult<PyListRef> {
396529
zelf.try_not_released(vm)?;
@@ -433,7 +566,7 @@ impl PyMemoryView {
433566
.into_ref(vm))
434567
}
435568

436-
#[pymethod]
569+
#[pymethod(magic)]
437570
fn repr(zelf: PyRef<Self>) -> String {
438571
if zelf.released.load() {
439572
format!("<released memory at 0x{:x}>", zelf.get_id())
@@ -443,10 +576,12 @@ impl PyMemoryView {
443576
}
444577

445578
fn eq(zelf: &PyRef<Self>, other: &PyObjectRef, vm: &VirtualMachine) -> PyResult<bool> {
446-
zelf.try_not_released(vm)?;
447579
if zelf.is(other) {
448580
return Ok(true);
449581
}
582+
if zelf.released.load() {
583+
return Ok(false);
584+
}
450585
let options_cmp = |a: &BufferOptions, b: &BufferOptions| -> bool {
451586
a.len == b.len && a.itemsize == b.itemsize
452587
};
@@ -535,6 +670,10 @@ impl Buffer for PyMemoryViewRef {
535670
&mut x[self.start..self.stop]
536671
}))
537672
}
673+
674+
fn to_contiguous(&self) -> Vec<u8> {
675+
PyMemoryView::to_bytes_vec(self)
676+
}
538677
}
539678

540679
impl Comparable for PyMemoryView {

‎vm/src/sliceable.rs

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -434,14 +434,18 @@ fn saturate_big_index(slice_pos: &BigInt, len: usize) -> usize {
434434
}
435435
}
436436

437-
pub fn saturate_range(start: &Option<BigInt>, stop: &Option<BigInt>, len: usize) -> Range<usize> {
437+
pub(crate) fn saturate_range(
438+
start: &Option<BigInt>,
439+
stop: &Option<BigInt>,
440+
len: usize,
441+
) -> Range<usize> {
438442
let start = start.as_ref().map_or(0, |x| saturate_big_index(x, len));
439443
let stop = stop.as_ref().map_or(len, |x| saturate_big_index(x, len));
440444

441445
start..stop
442446
}
443447

444-
fn convert_slice(
448+
pub(crate) fn convert_slice(
445449
slice: &PySlice,
446450
len: usize,
447451
vm: &VirtualMachine,

‎vm/src/stdlib/pystruct.rs

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -119,7 +119,7 @@ pub(crate) mod _struct {
119119
Ok(FormatSpec { endianness, codes })
120120
}
121121

122-
fn pack(&self, args: &[PyObjectRef], vm: &VirtualMachine) -> PyResult<Vec<u8>> {
122+
pub fn pack(&self, args: &[PyObjectRef], vm: &VirtualMachine) -> PyResult<Vec<u8>> {
123123
// Create data vector:
124124
let mut data = Vec::<u8>::new();
125125

0 commit comments

Comments
 (0)