Skip to content

Commit 00a86ba

Browse files
youknowoneyouknowone
authored
Merge pull request #3029 from youknowone/pybuffer
Refactor PyBuffer
2 parents 76bb68a + f2ecab4 commit 00a86ba

File tree

14 files changed

+337
-419
lines changed

14 files changed

+337
-419
lines changed

‎Lib/test/test_httplib.py

-4
Original file line numberDiff line numberDiff line change
@@ -839,8 +839,6 @@ def test_send_type_error(self):
839839
with self.assertRaises(TypeError):
840840
conn.request('POST', 'test', conn)
841841

842-
# TODO: RUSTPYTHON
843-
@unittest.expectedFailure
844842
def test_chunked(self):
845843
expected = chunked_expected
846844
sock = FakeSocket(chunked_start + last_chunk + chunked_end)
@@ -873,8 +871,6 @@ def test_chunked(self):
873871
finally:
874872
resp.close()
875873

876-
# TODO: RUSTPYTHON
877-
@unittest.expectedFailure
878874
def test_readinto_chunked(self):
879875

880876
expected = chunked_expected

‎Lib/test/test_io.py

+2-10
Original file line numberDiff line numberDiff line change
@@ -1039,8 +1039,6 @@ def test_destructor(self):
10391039
super().test_destructor(self)
10401040

10411041
class PyIOTest(IOTest):
1042-
# TODO: RUSTPYTHON, can't resize b/c of existing exports
1043-
@unittest.expectedFailure
10441042
def test_optional_abilities(self):
10451043
super().test_optional_abilities()
10461044

@@ -2630,8 +2628,7 @@ def test_non_text_encoding_codecs_are_rejected(self):
26302628
with self.assertRaisesRegex(LookupError, "is not a text encoding"):
26312629
self.TextIOWrapper(b, encoding="hex")
26322630

2633-
# TODO: RUSTPYTHON
2634-
@unittest.expectedFailure
2631+
@unittest.skip('TODO: RUSTPYTHON')
26352632
def test_detach(self):
26362633
r = self.BytesIO()
26372634
b = self.BufferedWriter(r)
@@ -2693,8 +2690,7 @@ def test_line_buffering(self):
26932690
t.write("A\rB")
26942691
self.assertEqual(r.getvalue(), b"XY\nZA\rB")
26952692

2696-
# TODO: RUSTPYTHON
2697-
@unittest.expectedFailure
2693+
@unittest.skip('TODO: RUSTPYTHON')
26982694
def test_reconfigure_line_buffering(self):
26992695
r = self.BytesIO()
27002696
b = self.BufferedWriter(r, 1000)
@@ -3924,16 +3920,12 @@ class PyTextIOWrapperTest(TextIOWrapperTest):
39243920
def test_newlines(self):
39253921
super().test_newlines()
39263922

3927-
# TODO: RUSTPYTHON
3928-
@unittest.expectedFailure
39293923
def test_line_buffering(self):
39303924
super().test_line_buffering()
39313925

39323926
def test_seeking_too(self):
39333927
super().test_seeking_too()
39343928

3935-
# TODO: RUSTPYTHON, can't resize b/c of existing exports
3936-
@unittest.expectedFailure
39373929
def test_bufio_write_through(self):
39383930
super().test_bufio_write_through()
39393931

‎vm/src/buffer.rs

+62-90
Original file line numberDiff line numberDiff line change
@@ -3,12 +3,10 @@
33
use crate::common::borrow::{BorrowedValue, BorrowedValueMut};
44
use crate::common::rc::PyRc;
55
use crate::PyThreadingConstraint;
6-
use crate::{PyObjectRef, PyResult, TypeProtocol};
7-
use crate::{TryFromBorrowedObject, VirtualMachine};
8-
use std::{borrow::Cow, fmt::Debug, ops::Deref};
6+
use crate::{PyObjectRef, PyResult, TryFromBorrowedObject, TypeProtocol, VirtualMachine};
7+
use std::{borrow::Cow, fmt::Debug};
98

10-
pub trait PyBuffer: Debug + PyThreadingConstraint {
11-
fn get_options(&self) -> &BufferOptions;
9+
pub trait PyBufferInternal: Debug + PyThreadingConstraint {
1210
/// Get the full inner buffer of this memory. You probably want [`as_contiguous()`], as
1311
/// `obj_bytes` doesn't take into account the range a memoryview might operate on, among other
1412
/// footguns.
@@ -18,49 +16,84 @@ pub trait PyBuffer: Debug + PyThreadingConstraint {
1816
/// might operate on, among other footguns.
1917
fn obj_bytes_mut(&self) -> BorrowedValueMut<[u8]>;
2018
fn release(&self);
19+
// not included in PyBuffer protocol itself
20+
fn retain(&self);
21+
}
22+
23+
#[derive(Debug)]
24+
pub struct PyBuffer {
25+
pub obj: PyObjectRef,
26+
pub options: BufferOptions,
27+
pub(crate) internal: PyRc<dyn PyBufferInternal>,
28+
}
2129

22-
fn as_contiguous(&self) -> Option<BorrowedValue<[u8]>> {
23-
if !self.get_options().contiguous {
30+
impl PyBuffer {
31+
pub fn new(
32+
obj: PyObjectRef,
33+
buffer: impl PyBufferInternal + 'static,
34+
options: BufferOptions,
35+
) -> Self {
36+
buffer.retain();
37+
Self {
38+
obj,
39+
options,
40+
internal: PyRc::new(buffer),
41+
}
42+
}
43+
pub fn as_contiguous(&self) -> Option<BorrowedValue<[u8]>> {
44+
if !self.options.contiguous {
2445
return None;
2546
}
26-
Some(self.obj_bytes())
47+
Some(self.internal.obj_bytes())
2748
}
2849

29-
fn as_contiguous_mut(&self) -> Option<BorrowedValueMut<[u8]>> {
30-
if !self.get_options().contiguous {
50+
pub fn as_contiguous_mut(&self) -> Option<BorrowedValueMut<[u8]>> {
51+
if !self.options.contiguous {
3152
return None;
3253
}
33-
Some(self.obj_bytes_mut())
54+
Some(self.internal.obj_bytes_mut())
3455
}
3556

36-
fn to_contiguous(&self) -> Vec<u8> {
37-
self.obj_bytes().to_vec()
57+
pub fn to_contiguous(&self) -> Vec<u8> {
58+
self.internal.obj_bytes().to_vec()
59+
}
60+
61+
pub fn clone_with_options(&self, options: BufferOptions) -> Self {
62+
self.internal.retain();
63+
Self {
64+
obj: self.obj.clone(),
65+
options,
66+
internal: self.internal.clone(),
67+
}
3868
}
3969
}
4070

4171
#[derive(Debug, Clone)]
4272
pub struct BufferOptions {
43-
pub readonly: bool,
73+
// buf
4474
pub len: usize,
75+
pub readonly: bool,
4576
pub itemsize: usize,
46-
pub contiguous: bool,
4777
pub format: Cow<'static, str>,
48-
// TODO: support multiple dimension array
49-
pub ndim: usize,
78+
pub ndim: usize, // TODO: support multiple dimension array
5079
pub shape: Vec<usize>,
5180
pub strides: Vec<isize>,
81+
// suboffsets
82+
83+
// RustPython fields
84+
pub contiguous: bool,
5285
}
5386

5487
impl BufferOptions {
5588
pub const DEFAULT: Self = BufferOptions {
56-
readonly: true,
5789
len: 0,
90+
readonly: true,
5891
itemsize: 1,
59-
contiguous: true,
6092
format: Cow::Borrowed("B"),
6193
ndim: 1,
6294
shape: Vec::new(),
6395
strides: Vec::new(),
96+
contiguous: true,
6497
};
6598
}
6699

@@ -70,15 +103,12 @@ impl Default for BufferOptions {
70103
}
71104
}
72105

73-
#[derive(Debug)]
74-
pub struct PyBufferRef(Box<dyn PyBuffer>);
75-
76-
impl TryFromBorrowedObject for PyBufferRef {
106+
impl TryFromBorrowedObject for PyBuffer {
77107
fn try_from_borrowed_object(vm: &VirtualMachine, obj: &PyObjectRef) -> PyResult<Self> {
78108
let obj_cls = obj.class();
79109
for cls in obj_cls.iter_mro() {
80110
if let Some(f) = cls.slots.as_buffer.as_ref() {
81-
return f(obj, vm).map(|x| PyBufferRef(x));
111+
return f(obj, vm);
82112
}
83113
}
84114
Err(vm.new_type_error(format!(
@@ -88,76 +118,18 @@ impl TryFromBorrowedObject for PyBufferRef {
88118
}
89119
}
90120

91-
impl Drop for PyBufferRef {
92-
fn drop(&mut self) {
93-
self.0.release();
94-
}
95-
}
96-
97-
impl Deref for PyBufferRef {
98-
type Target = dyn PyBuffer;
99-
100-
fn deref(&self) -> &Self::Target {
101-
self.0.deref()
102-
}
103-
}
104-
105-
impl PyBufferRef {
106-
pub fn new(buffer: impl PyBuffer + 'static) -> Self {
107-
Self(Box::new(buffer))
108-
}
109-
110-
pub fn into_rcbuf(self) -> RcBuffer {
111-
// move self.0 out of self; PyBufferRef impls Drop so it's tricky
112-
let this = std::mem::ManuallyDrop::new(self);
113-
let buf_box = unsafe { std::ptr::read(&this.0) };
114-
RcBuffer(buf_box.into())
115-
}
116-
}
117-
118-
impl From<Box<dyn PyBuffer>> for PyBufferRef {
119-
fn from(buffer: Box<dyn PyBuffer>) -> Self {
120-
PyBufferRef(buffer)
121-
}
122-
}
123-
124-
#[derive(Debug, Clone)]
125-
pub struct RcBuffer(PyRc<dyn PyBuffer>);
126-
impl Deref for RcBuffer {
127-
type Target = dyn PyBuffer;
128-
fn deref(&self) -> &Self::Target {
129-
self.0.deref()
130-
}
131-
}
132-
133-
impl Drop for RcBuffer {
121+
// What we actually want to implement is:
122+
// impl<T> Drop for T where T: PyBufferInternal
123+
// but it is not supported by Rust
124+
impl Drop for PyBuffer {
134125
fn drop(&mut self) {
135-
// check if this is the last rc before the inner buffer gets dropped
136-
if let Some(buf) = PyRc::get_mut(&mut self.0) {
137-
buf.release()
138-
}
126+
self.internal.release();
139127
}
140128
}
141129

142-
impl PyBuffer for RcBuffer {
143-
fn get_options(&self) -> &BufferOptions {
144-
self.0.get_options()
145-
}
146-
fn obj_bytes(&self) -> BorrowedValue<[u8]> {
147-
self.0.obj_bytes()
148-
}
149-
fn obj_bytes_mut(&self) -> BorrowedValueMut<[u8]> {
150-
self.0.obj_bytes_mut()
151-
}
152-
fn release(&self) {}
153-
fn as_contiguous(&self) -> Option<BorrowedValue<[u8]>> {
154-
self.0.as_contiguous()
155-
}
156-
fn as_contiguous_mut(&self) -> Option<BorrowedValueMut<[u8]>> {
157-
self.0.as_contiguous_mut()
158-
}
159-
fn to_contiguous(&self) -> Vec<u8> {
160-
self.0.to_contiguous()
130+
impl Clone for PyBuffer {
131+
fn clone(&self) -> Self {
132+
self.clone_with_options(self.options.clone())
161133
}
162134
}
163135

‎vm/src/builtins/bytearray.rs

+14-20
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,7 @@ use super::pystr::PyStrRef;
66
use super::pytype::PyTypeRef;
77
use super::tuple::PyTupleRef;
88
use crate::anystr::{self, AnyStr};
9-
use crate::buffer::{BufferOptions, PyBuffer, ResizeGuard};
9+
use crate::buffer::{BufferOptions, PyBuffer, PyBufferInternal, ResizeGuard};
1010
use crate::bytesinner::{
1111
bytes_decode, bytes_from_object, value_from_object, ByteInnerFindOptions, ByteInnerNewOptions,
1212
ByteInnerPaddingOptions, ByteInnerSplitOptions, ByteInnerTranslateOptions, DecodeArgs,
@@ -670,41 +670,35 @@ impl Comparable for PyByteArray {
670670
}
671671

672672
impl AsBuffer for PyByteArray {
673-
fn get_buffer(zelf: &PyRef<Self>, _vm: &VirtualMachine) -> PyResult<Box<dyn PyBuffer>> {
674-
zelf.exports.fetch_add(1);
675-
let buf = ByteArrayBuffer {
676-
bytearray: zelf.clone(),
677-
options: BufferOptions {
673+
fn get_buffer(zelf: &PyRef<Self>, _vm: &VirtualMachine) -> PyResult<PyBuffer> {
674+
let buffer = PyBuffer::new(
675+
zelf.as_object().clone(),
676+
zelf.clone(),
677+
BufferOptions {
678678
readonly: false,
679679
len: zelf.len(),
680680
..Default::default()
681681
},
682-
};
683-
Ok(Box::new(buf))
682+
);
683+
Ok(buffer)
684684
}
685685
}
686686

687-
#[derive(Debug)]
688-
struct ByteArrayBuffer {
689-
bytearray: PyByteArrayRef,
690-
options: BufferOptions,
691-
}
692-
693-
impl PyBuffer for ByteArrayBuffer {
687+
impl PyBufferInternal for PyRef<PyByteArray> {
694688
fn obj_bytes(&self) -> BorrowedValue<[u8]> {
695-
self.bytearray.borrow_buf().into()
689+
self.borrow_buf().into()
696690
}
697691

698692
fn obj_bytes_mut(&self) -> BorrowedValueMut<[u8]> {
699-
PyRwLockWriteGuard::map(self.bytearray.inner_mut(), |inner| &mut *inner.elements).into()
693+
PyRwLockWriteGuard::map(self.inner_mut(), |inner| &mut *inner.elements).into()
700694
}
701695

702696
fn release(&self) {
703-
self.bytearray.exports.fetch_sub(1);
697+
self.exports.fetch_sub(1);
704698
}
705699

706-
fn get_options(&self) -> &BufferOptions {
707-
&self.options
700+
fn retain(&self) {
701+
self.exports.fetch_add(1);
708702
}
709703
}
710704

0 commit comments

Comments
 (0)