Skip to content

Commit 52cb6fa

Browse files
qingshi163qingshi163
committed
Fix buffer protocol and memoryview
1 parent 4ede6b1 commit 52cb6fa

File tree

9 files changed

+342
-360
lines changed

9 files changed

+342
-360
lines changed

‎stdlib/src/array.rs

Lines changed: 5 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -1231,24 +1231,21 @@ mod array {
12311231
}
12321232
}
12331233

1234-
#[derive(Debug)]
1235-
struct PyArrayBufferInternal(PyRef<PyArray>);
1236-
1237-
impl BufferInternal for PyArrayBufferInternal {
1234+
impl BufferInternal for PyArray {
12381235
fn obj_bytes(&self) -> BorrowedValue<[u8]> {
1239-
self.0.get_bytes().into()
1236+
self.get_bytes().into()
12401237
}
12411238

12421239
fn obj_bytes_mut(&self) -> BorrowedValueMut<[u8]> {
1243-
self.0.get_bytes_mut().into()
1240+
self.get_bytes_mut().into()
12441241
}
12451242

12461243
fn release(&self) {
1247-
self.0.exports.fetch_sub(1);
1244+
self.exports.fetch_sub(1);
12481245
}
12491246

12501247
fn retain(&self) {
1251-
self.0.exports.fetch_add(1);
1248+
self.exports.fetch_add(1);
12521249
}
12531250
}
12541251

‎vm/src/builtins/bytearray.rs

Lines changed: 23 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@ use crate::common::{
99
PyMappedRwLockReadGuard, PyMappedRwLockWriteGuard, PyMutex, PyRwLock, PyRwLockReadGuard,
1010
PyRwLockWriteGuard,
1111
},
12+
static_cell,
1213
};
1314
use crate::{
1415
anystr::{self, AnyStr},
@@ -20,7 +21,7 @@ use crate::{
2021
},
2122
function::{ArgBytesLike, ArgIterable, FuncArgs, IntoPyObject, OptionalArg, OptionalOption},
2223
protocol::{
23-
BufferInternal, BufferOptions, BufferResizeGuard, PyBuffer, PyIterReturn, PyMappingMethods,
24+
BufferMethods, BufferOptions, BufferResizeGuard, PyBuffer, PyIterReturn, PyMappingMethods,
2425
},
2526
sliceable::{PySliceableSequence, PySliceableSequenceMut, SequenceIndex},
2627
types::{
@@ -33,6 +34,7 @@ use crate::{
3334
};
3435
use bstr::ByteSlice;
3536
use crossbeam_utils::atomic::AtomicCell;
37+
use rustpython_common::atomic::Radium;
3638
use std::mem::size_of;
3739

3840
/// "bytearray(iterable_of_ints) -> bytearray\n\
@@ -710,37 +712,38 @@ impl Comparable for PyByteArray {
710712

711713
impl AsBuffer for PyByteArray {
712714
fn as_buffer(zelf: &PyRef<Self>, _vm: &VirtualMachine) -> PyResult<PyBuffer> {
715+
static_cell! {
716+
static METHODS: BufferMethods;
717+
}
718+
let methods = METHODS.get_or_init(|| BufferMethods {
719+
obj_bytes: |zelf| zelf.payload::<Self>().unwrap().borrow_buf().into(),
720+
obj_bytes_mut: |zelf| {
721+
let zelf = zelf.payload::<Self>().unwrap();
722+
PyRwLockWriteGuard::map(zelf.inner_mut(), |x| &mut *x.elements).into()
723+
},
724+
release: Some(|zelf| {
725+
zelf.payload::<Self>().unwrap().exports.fetch_sub(1);
726+
}),
727+
retain: Some(|zelf| {
728+
zelf.payload::<Self>().unwrap().exports.fetch_add(1);
729+
}),
730+
contiguous: None,
731+
contiguous_mut: None,
732+
collect_bytes: None,
733+
});
713734
let buffer = PyBuffer::new(
714735
zelf.as_object().clone(),
715-
zelf.clone(),
716736
BufferOptions {
717737
readonly: false,
718738
len: zelf.len(),
719739
..Default::default()
720740
},
741+
methods,
721742
);
722743
Ok(buffer)
723744
}
724745
}
725746

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-
}
738-
739-
fn retain(&self) {
740-
self.exports.fetch_add(1);
741-
}
742-
}
743-
744747
impl<'a> BufferResizeGuard<'a> for PyByteArray {
745748
type Resizable = PyRwLockWriteGuard<'a, PyBytesInner>;
746749

‎vm/src/builtins/bytes.rs

Lines changed: 12 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,7 @@ use crate::{
1010
function::{
1111
ArgBytesLike, ArgIterable, IntoPyObject, IntoPyResult, OptionalArg, OptionalOption,
1212
},
13-
protocol::{BufferInternal, BufferOptions, PyBuffer, PyIterReturn, PyMappingMethods},
13+
protocol::{BufferMethods, BufferOptions, PyBuffer, PyIterReturn, PyMappingMethods},
1414
types::{
1515
AsBuffer, AsMapping, Callable, Comparable, Constructor, Hashable, IterNext,
1616
IterNextIterable, Iterable, PyComparisonOp, Unconstructible,
@@ -541,33 +541,30 @@ impl PyBytes {
541541
}
542542
}
543543

544+
static BUFFER_METHODS: BufferMethods = BufferMethods {
545+
obj_bytes: |zelf| zelf.payload::<PyBytes>().unwrap().as_bytes().into(),
546+
obj_bytes_mut: |_| panic!(),
547+
contiguous: None,
548+
contiguous_mut: None,
549+
collect_bytes: None,
550+
release: None,
551+
retain: None,
552+
};
553+
544554
impl AsBuffer for PyBytes {
545555
fn as_buffer(zelf: &PyRef<Self>, _vm: &VirtualMachine) -> PyResult<PyBuffer> {
546556
let buf = PyBuffer::new(
547557
zelf.as_object().clone(),
548-
zelf.clone(),
549558
BufferOptions {
550559
len: zelf.len(),
551560
..Default::default()
552561
},
562+
&BUFFER_METHODS,
553563
);
554564
Ok(buf)
555565
}
556566
}
557567

558-
impl BufferInternal for PyRef<PyBytes> {
559-
fn obj_bytes(&self) -> BorrowedValue<[u8]> {
560-
self.as_bytes().into()
561-
}
562-
563-
fn obj_bytes_mut(&self) -> BorrowedValueMut<[u8]> {
564-
unreachable!("bytes is not mutable")
565-
}
566-
567-
fn release(&self) {}
568-
fn retain(&self) {}
569-
}
570-
571568
impl AsMapping for PyBytes {
572569
fn as_mapping(_zelf: &PyRef<Self>, _vm: &VirtualMachine) -> PyMappingMethods {
573570
PyMappingMethods {

0 commit comments

Comments
 (0)