Skip to content

Commit 742ea0c

Browse files
qingshi163youknowone
qingshi163
and
youknowone
authored
Reimpl Buffer Protocol and memoryview support ndarray with shape, stride and suboffset (#3340)
* Fix buffer protocol and memoryview * Fix io and array the right way to use buffer protocol * fix rebase * fix memoryview and introduce VecBuffer * fix deadlock * fix memoryview avoid double release buffer * impl ndim buffer support * first implement nd array for buffer and memoryview * fix slice adjust_indices * fix adjust_indices introduct SaturatedSliceIterator * fix memoryview cmp * fix adjust_indices positive_order * fix bug mark passed test * fix clippy * fix set_item_slice deadlock, optimize buffer * Delete @test_113974_tmp * Remove vec_buffer_type * rusty names * impl memoryview multi index * add comments Co-authored-by: Jeong YunWon <jeong@youknowone.org>
1 parent 5ceb2be commit 742ea0c

17 files changed

Lines changed: 1376 additions & 1030 deletions

File tree

‎Lib/test/test_gzip.py‎

Lines changed: 0 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -88,8 +88,6 @@ def test_write_read_with_pathlike_file(self):
8888
self.assertEqual(d, data1 * 51)
8989
self.assertIsInstance(f.name, str)
9090

91-
# TODO: RUSTPYTHON
92-
@unittest.expectedFailure
9391
# The following test_write_xy methods test that write accepts
9492
# the corresponding bytes-like object type as input
9593
# and that the data written equals bytes(xy) in all cases.

‎Lib/test/test_memoryview.py‎

Lines changed: 0 additions & 2 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")

‎extra_tests/snippets/memoryview.py‎

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -45,6 +45,8 @@ def test_slice():
4545
assert m[::2][1:-1].tobytes() == b'357'
4646
assert m[::2][::2].tobytes() == b'159'
4747
assert m[::2][1::2].tobytes() == b'37'
48+
assert m[::-1].tobytes() == b'987654321'
49+
assert m[::-2].tobytes() == b'97531'
4850

4951
test_slice()
5052

@@ -56,9 +58,9 @@ def test_resizable():
5658
m.release()
5759
b.append(6)
5860
m2 = memoryview(b)
59-
m4 = memoryview(b)
61+
m4 = memoryview(m2)
6062
assert_raises(BufferError, lambda: b.append(5))
61-
m3 = memoryview(b)
63+
m3 = memoryview(m2)
6264
assert_raises(BufferError, lambda: b.append(5))
6365
m2.release()
6466
assert_raises(BufferError, lambda: b.append(5))

‎stdlib/src/array.rs‎

Lines changed: 34 additions & 40 deletions
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,7 @@ pub(crate) use array::make_module;
33
#[pymodule(name = "array")]
44
mod array {
55
use crate::common::{
6-
borrow::{BorrowedValue, BorrowedValueMut},
6+
atomic::{self, AtomicUsize},
77
lock::{
88
PyMappedRwLockReadGuard, PyMappedRwLockWriteGuard, PyRwLock, PyRwLockReadGuard,
99
PyRwLockWriteGuard,
@@ -20,18 +20,17 @@ mod array {
2020
ArgBytesLike, ArgIntoFloat, ArgIterable, IntoPyObject, IntoPyResult, OptionalArg,
2121
},
2222
protocol::{
23-
BufferInternal, BufferOptions, BufferResizeGuard, PyBuffer, PyIterReturn,
23+
BufferDescriptor, BufferMethods, BufferResizeGuard, PyBuffer, PyIterReturn,
2424
PyMappingMethods,
2525
},
2626
sliceable::{PySliceableSequence, PySliceableSequenceMut, SaturatedSlice, SequenceIndex},
2727
types::{
2828
AsBuffer, AsMapping, Comparable, Constructor, IterNext, IterNextIterable, Iterable,
2929
PyComparisonOp,
3030
},
31-
IdProtocol, PyComparisonValue, PyObject, PyObjectRef, PyObjectView, PyRef, PyResult,
32-
PyValue, TryFromObject, TypeProtocol, VirtualMachine,
31+
IdProtocol, PyComparisonValue, PyObject, PyObjectRef, PyObjectView, PyObjectWrap, PyRef,
32+
PyResult, PyValue, TryFromObject, TypeProtocol, VirtualMachine,
3333
};
34-
use crossbeam_utils::atomic::AtomicCell;
3534
use itertools::Itertools;
3635
use num_traits::ToPrimitive;
3736
use std::cmp::Ordering;
@@ -615,7 +614,7 @@ mod array {
615614
#[derive(Debug, PyValue)]
616615
pub struct PyArray {
617616
array: PyRwLock<ArrayContentType>,
618-
exports: AtomicCell<usize>,
617+
exports: AtomicUsize,
619618
}
620619

621620
pub type PyArrayRef = PyRef<PyArray>;
@@ -624,7 +623,7 @@ mod array {
624623
fn from(array: ArrayContentType) -> Self {
625624
PyArray {
626625
array: PyRwLock::new(array),
627-
exports: AtomicCell::new(0),
626+
exports: AtomicUsize::new(0),
628627
}
629628
}
630629
}
@@ -1220,40 +1219,35 @@ mod array {
12201219
fn as_buffer(zelf: &PyObjectView<Self>, _vm: &VirtualMachine) -> PyResult<PyBuffer> {
12211220
let array = zelf.read();
12221221
let buf = PyBuffer::new(
1223-
zelf.as_object().to_owned(),
1224-
PyArrayBufferInternal(zelf.to_owned()),
1225-
BufferOptions {
1226-
readonly: false,
1227-
len: array.len(),
1228-
itemsize: array.itemsize(),
1229-
format: array.typecode_str().into(),
1230-
..Default::default()
1231-
},
1222+
zelf.to_owned().into(),
1223+
BufferDescriptor::format(
1224+
array.len() * array.itemsize(),
1225+
false,
1226+
array.itemsize(),
1227+
array.typecode_str().into(),
1228+
),
1229+
&BUFFER_METHODS,
12321230
);
12331231
Ok(buf)
12341232
}
12351233
}
12361234

1237-
#[derive(Debug)]
1238-
struct PyArrayBufferInternal(PyRef<PyArray>);
1239-
1240-
impl BufferInternal for PyArrayBufferInternal {
1241-
fn obj_bytes(&self) -> BorrowedValue<[u8]> {
1242-
self.0.get_bytes().into()
1243-
}
1244-
1245-
fn obj_bytes_mut(&self) -> BorrowedValueMut<[u8]> {
1246-
self.0.get_bytes_mut().into()
1247-
}
1248-
1249-
fn release(&self) {
1250-
self.0.exports.fetch_sub(1);
1251-
}
1252-
1253-
fn retain(&self) {
1254-
self.0.exports.fetch_add(1);
1255-
}
1256-
}
1235+
static BUFFER_METHODS: BufferMethods = BufferMethods {
1236+
obj_bytes: |buffer| buffer.obj_as::<PyArray>().get_bytes().into(),
1237+
obj_bytes_mut: |buffer| buffer.obj_as::<PyArray>().get_bytes_mut().into(),
1238+
release: |buffer| {
1239+
buffer
1240+
.obj_as::<PyArray>()
1241+
.exports
1242+
.fetch_sub(1, atomic::Ordering::Release);
1243+
},
1244+
retain: |buffer| {
1245+
buffer
1246+
.obj_as::<PyArray>()
1247+
.exports
1248+
.fetch_add(1, atomic::Ordering::Release);
1249+
},
1250+
};
12571251

12581252
impl AsMapping for PyArray {
12591253
fn as_mapping(_zelf: &PyObjectView<Self>, _vm: &VirtualMachine) -> PyMappingMethods {
@@ -1290,7 +1284,7 @@ mod array {
12901284
impl Iterable for PyArray {
12911285
fn iter(zelf: PyRef<Self>, vm: &VirtualMachine) -> PyResult {
12921286
Ok(PyArrayIter {
1293-
position: AtomicCell::new(0),
1287+
position: AtomicUsize::new(0),
12941288
array: zelf,
12951289
}
12961290
.into_object(vm))
@@ -1302,7 +1296,7 @@ mod array {
13021296

13031297
fn try_resizable(&'a self, vm: &VirtualMachine) -> PyResult<Self::Resizable> {
13041298
let w = self.write();
1305-
if self.exports.load() == 0 {
1299+
if self.exports.load(atomic::Ordering::SeqCst) == 0 {
13061300
Ok(w)
13071301
} else {
13081302
Err(vm.new_buffer_error(
@@ -1316,7 +1310,7 @@ mod array {
13161310
#[pyclass(name = "array_iterator")]
13171311
#[derive(Debug, PyValue)]
13181312
pub struct PyArrayIter {
1319-
position: AtomicCell<usize>,
1313+
position: AtomicUsize,
13201314
array: PyArrayRef,
13211315
}
13221316

@@ -1326,7 +1320,7 @@ mod array {
13261320
impl IterNextIterable for PyArrayIter {}
13271321
impl IterNext for PyArrayIter {
13281322
fn next(zelf: &PyObjectView<Self>, vm: &VirtualMachine) -> PyResult<PyIterReturn> {
1329-
let pos = zelf.position.fetch_add(1);
1323+
let pos = zelf.position.fetch_add(1, atomic::Ordering::SeqCst);
13301324
let r = if let Some(item) = zelf.array.read().getitem_by_idx(pos, vm)? {
13311325
PyIterReturn::Return(item)
13321326
} else {

‎vm/src/builtins/bytearray.rs‎

Lines changed: 47 additions & 42 deletions
Original file line numberDiff line numberDiff line change
@@ -3,13 +3,6 @@ use super::{
33
PositionIterInternal, PyBytes, PyBytesRef, PyDictRef, PyIntRef, PyStrRef, PyTuple, PyTupleRef,
44
PyTypeRef,
55
};
6-
use crate::common::{
7-
borrow::{BorrowedValue, BorrowedValueMut},
8-
lock::{
9-
PyMappedRwLockReadGuard, PyMappedRwLockWriteGuard, PyMutex, PyRwLock, PyRwLockReadGuard,
10-
PyRwLockWriteGuard,
11-
},
12-
};
136
use crate::{
147
anystr::{self, AnyStr},
158
builtins::PyType,
@@ -18,9 +11,17 @@ use crate::{
1811
ByteInnerNewOptions, ByteInnerPaddingOptions, ByteInnerSplitOptions,
1912
ByteInnerTranslateOptions, DecodeArgs, PyBytesInner,
2013
},
14+
common::{
15+
atomic::{AtomicUsize, Ordering},
16+
lock::{
17+
PyMappedRwLockReadGuard, PyMappedRwLockWriteGuard, PyMutex, PyRwLock,
18+
PyRwLockReadGuard, PyRwLockWriteGuard,
19+
},
20+
},
2121
function::{ArgBytesLike, ArgIterable, FuncArgs, IntoPyObject, OptionalArg, OptionalOption},
2222
protocol::{
23-
BufferInternal, BufferOptions, BufferResizeGuard, PyBuffer, PyIterReturn, PyMappingMethods,
23+
BufferDescriptor, BufferMethods, BufferResizeGuard, PyBuffer, PyIterReturn,
24+
PyMappingMethods,
2425
},
2526
sliceable::{PySliceableSequence, PySliceableSequenceMut, SequenceIndex},
2627
types::{
@@ -29,10 +30,9 @@ use crate::{
2930
},
3031
utils::Either,
3132
IdProtocol, PyClassDef, PyClassImpl, PyComparisonValue, PyContext, PyObject, PyObjectRef,
32-
PyRef, PyResult, PyValue, TypeProtocol, VirtualMachine,
33+
PyObjectView, PyObjectWrap, PyRef, PyResult, PyValue, TypeProtocol, VirtualMachine,
3334
};
3435
use bstr::ByteSlice;
35-
use crossbeam_utils::atomic::AtomicCell;
3636
use std::mem::size_of;
3737

3838
/// "bytearray(iterable_of_ints) -> bytearray\n\
@@ -50,7 +50,7 @@ use std::mem::size_of;
5050
#[derive(Debug, Default)]
5151
pub struct PyByteArray {
5252
inner: PyRwLock<PyBytesInner>,
53-
exports: AtomicCell<usize>,
53+
exports: AtomicUsize,
5454
}
5555

5656
pub type PyByteArrayRef = PyRef<PyByteArray>;
@@ -63,7 +63,7 @@ impl PyByteArray {
6363
fn from_inner(inner: PyBytesInner) -> Self {
6464
PyByteArray {
6565
inner: PyRwLock::new(inner),
66-
exports: AtomicCell::new(0),
66+
exports: AtomicUsize::new(0),
6767
}
6868
}
6969

@@ -118,6 +118,12 @@ impl PyByteArray {
118118
Ok(())
119119
}
120120

121+
#[cfg(debug_assertions)]
122+
#[pyproperty]
123+
fn exports(&self) -> usize {
124+
self.exports.load(Ordering::Relaxed)
125+
}
126+
121127
#[inline]
122128
fn inner(&self) -> PyRwLockReadGuard<'_, PyBytesInner> {
123129
self.inner.read()
@@ -708,36 +714,35 @@ impl Comparable for PyByteArray {
708714
}
709715
}
710716

711-
impl AsBuffer for PyByteArray {
712-
fn as_buffer(zelf: &crate::PyObjectView<Self>, _vm: &VirtualMachine) -> PyResult<PyBuffer> {
713-
let buffer = PyBuffer::new(
714-
zelf.as_object().to_owned(),
715-
zelf.to_owned(),
716-
BufferOptions {
717-
readonly: false,
718-
len: zelf.len(),
719-
..Default::default()
720-
},
721-
);
722-
Ok(buffer)
723-
}
724-
}
725-
726-
impl BufferInternal for PyRef<PyByteArray> {
727-
fn obj_bytes(&self) -> BorrowedValue<[u8]> {
728-
self.borrow_buf().into()
729-
}
730-
731-
fn obj_bytes_mut(&self) -> BorrowedValueMut<[u8]> {
732-
PyRwLockWriteGuard::map(self.inner_mut(), |inner| &mut *inner.elements).into()
733-
}
734-
735-
fn release(&self) {
736-
self.exports.fetch_sub(1);
737-
}
717+
static BUFFER_METHODS: BufferMethods = BufferMethods {
718+
obj_bytes: |buffer| buffer.obj_as::<PyByteArray>().borrow_buf().into(),
719+
obj_bytes_mut: |buffer| {
720+
PyMappedRwLockWriteGuard::map(buffer.obj_as::<PyByteArray>().borrow_buf_mut(), |x| {
721+
x.as_mut_slice()
722+
})
723+
.into()
724+
},
725+
release: |buffer| {
726+
buffer
727+
.obj_as::<PyByteArray>()
728+
.exports
729+
.fetch_sub(1, Ordering::Release);
730+
},
731+
retain: |buffer| {
732+
buffer
733+
.obj_as::<PyByteArray>()
734+
.exports
735+
.fetch_add(1, Ordering::Release);
736+
},
737+
};
738738

739-
fn retain(&self) {
740-
self.exports.fetch_add(1);
739+
impl AsBuffer for PyByteArray {
740+
fn as_buffer(zelf: &PyObjectView<Self>, _vm: &VirtualMachine) -> PyResult<PyBuffer> {
741+
Ok(PyBuffer::new(
742+
zelf.to_owned().into_object(),
743+
BufferDescriptor::simple(zelf.len(), false),
744+
&BUFFER_METHODS,
745+
))
741746
}
742747
}
743748

@@ -746,7 +751,7 @@ impl<'a> BufferResizeGuard<'a> for PyByteArray {
746751

747752
fn try_resizable(&'a self, vm: &VirtualMachine) -> PyResult<Self::Resizable> {
748753
let w = self.inner.upgradable_read();
749-
if self.exports.load() == 0 {
754+
if self.exports.load(Ordering::SeqCst) == 0 {
750755
Ok(parking_lot::lock_api::RwLockUpgradableReadGuard::upgrade(w))
751756
} else {
752757
Err(vm

0 commit comments

Comments
 (0)