Merge pull request RustPython#2856 from youknowone/buffer-protocol

give an own module for buffer protocol + TryFromBorrowedObject

Author: youknowone

vm/src/buffer.rs

@@ -0,0 +1,167 @@
+//! Buffer protocol
+
+use crate::common::borrow::{BorrowedValue, BorrowedValueMut};
+use crate::common::rc::PyRc;
+use crate::PyThreadingConstraint;
+use crate::{PyObjectRef, PyResult, TypeProtocol};
+use crate::{TryFromBorrowedObject, VirtualMachine};
+use std::{borrow::Cow, fmt::Debug, ops::Deref};
+
+pub trait PyBuffer: Debug + PyThreadingConstraint {
+    fn get_options(&self) -> &BufferOptions;
+    /// Get the full inner buffer of this memory. You probably want [`as_contiguous()`], as
+    /// `obj_bytes` doesn't take into account the range a memoryview might operate on, among other
+    /// footguns.
+    fn obj_bytes(&self) -> BorrowedValue<[u8]>;
+    /// Get the full inner buffer of this memory, mutably. You probably want
+    /// [`as_contiguous_mut()`], as `obj_bytes` doesn't take into account the range a memoryview
+    /// might operate on, among other footguns.
+    fn obj_bytes_mut(&self) -> BorrowedValueMut<[u8]>;
+    fn release(&self);
+
+    fn as_contiguous(&self) -> Option<BorrowedValue<[u8]>> {
+        if !self.get_options().contiguous {
+            return None;
+        }
+        Some(self.obj_bytes())
+    }
+
+    fn as_contiguous_mut(&self) -> Option<BorrowedValueMut<[u8]>> {
+        if !self.get_options().contiguous {
+            return None;
+        }
+        Some(self.obj_bytes_mut())
+    }
+
+    fn to_contiguous(&self) -> Vec<u8> {
+        self.obj_bytes().to_vec()
+    }
+}
+
+#[derive(Debug, Clone)]
+pub struct BufferOptions {
+    pub readonly: bool,
+    pub len: usize,
+    pub itemsize: usize,
+    pub contiguous: bool,
+    pub format: Cow<'static, str>,
+    // TODO: support multiple dimension array
+    pub ndim: usize,
+    pub shape: Vec<usize>,
+    pub strides: Vec<isize>,
+}
+
+impl BufferOptions {
+    pub const DEFAULT: Self = BufferOptions {
+        readonly: true,
+        len: 0,
+        itemsize: 1,
+        contiguous: true,
+        format: Cow::Borrowed("B"),
+        ndim: 1,
+        shape: Vec::new(),
+        strides: Vec::new(),
+    };
+}
+
+impl Default for BufferOptions {
+    fn default() -> Self {
+        Self::DEFAULT
+    }
+}
+
+#[derive(Debug)]
+pub struct PyBufferRef(Box<dyn PyBuffer>);
+
+impl TryFromBorrowedObject for PyBufferRef {
+    fn try_from_borrowed_object(vm: &VirtualMachine, obj: &PyObjectRef) -> PyResult<Self> {
+        let obj_cls = obj.class();
+        for cls in obj_cls.iter_mro() {
+            if let Some(f) = cls.slots.as_buffer.as_ref() {
+                return f(obj, vm).map(|x| PyBufferRef(x));
+            }
+        }
+        Err(vm.new_type_error(format!(
+            "a bytes-like object is required, not '{}'",
+            obj_cls.name
+        )))
+    }
+}
+
+impl Drop for PyBufferRef {
+    fn drop(&mut self) {
+        self.0.release();
+    }
+}
+
+impl Deref for PyBufferRef {
+    type Target = dyn PyBuffer;
+
+    fn deref(&self) -> &Self::Target {
+        self.0.deref()
+    }
+}
+
+impl PyBufferRef {
+    pub fn new(buffer: impl PyBuffer + 'static) -> Self {
+        Self(Box::new(buffer))
+    }
+
+    pub fn into_rcbuf(self) -> RcBuffer {
+        // move self.0 out of self; PyBufferRef impls Drop so it's tricky
+        let this = std::mem::ManuallyDrop::new(self);
+        let buf_box = unsafe { std::ptr::read(&this.0) };
+        RcBuffer(buf_box.into())
+    }
+}
+
+impl From<Box<dyn PyBuffer>> for PyBufferRef {
+    fn from(buffer: Box<dyn PyBuffer>) -> Self {
+        PyBufferRef(buffer)
+    }
+}
+
+#[derive(Debug, Clone)]
+pub struct RcBuffer(PyRc<dyn PyBuffer>);
+impl Deref for RcBuffer {
+    type Target = dyn PyBuffer;
+    fn deref(&self) -> &Self::Target {
+        self.0.deref()
+    }
+}
+
+impl Drop for RcBuffer {
+    fn drop(&mut self) {
+        // check if this is the last rc before the inner buffer gets dropped
+        if let Some(buf) = PyRc::get_mut(&mut self.0) {
+            buf.release()
+        }
+    }
+}
+
+impl PyBuffer for RcBuffer {
+    fn get_options(&self) -> &BufferOptions {
+        self.0.get_options()
+    }
+    fn obj_bytes(&self) -> BorrowedValue<[u8]> {
+        self.0.obj_bytes()
+    }
+    fn obj_bytes_mut(&self) -> BorrowedValueMut<[u8]> {
+        self.0.obj_bytes_mut()
+    }
+    fn release(&self) {}
+    fn as_contiguous(&self) -> Option<BorrowedValue<[u8]>> {
+        self.0.as_contiguous()
+    }
+    fn as_contiguous_mut(&self) -> Option<BorrowedValueMut<[u8]>> {
+        self.0.as_contiguous_mut()
+    }
+    fn to_contiguous(&self) -> Vec<u8> {
+        self.0.to_contiguous()
+    }
+}
+
+pub(crate) trait ResizeGuard<'a> {
+    type Resizable: 'a;
+    fn try_resizable(&'a self, vm: &VirtualMachine) -> PyResult<Self::Resizable>;
+}

vm/src/builtins/bytearray.rs

@@ -2,27 +2,25 @@
 use super::bytes::{PyBytes, PyBytesRef};
 use super::dict::PyDictRef;
 use super::int::PyIntRef;
-use super::memory::{BufferOptions, PyBuffer, ResizeGuard};
 use super::pystr::PyStrRef;
 use super::pytype::PyTypeRef;
 use super::tuple::PyTupleRef;
 use crate::anystr::{self, AnyStr};
+use crate::buffer::{BufferOptions, PyBuffer, ResizeGuard};
 use crate::bytesinner::{
     bytes_decode, bytes_from_object, value_from_object, ByteInnerFindOptions, ByteInnerNewOptions,
     ByteInnerPaddingOptions, ByteInnerSplitOptions, ByteInnerTranslateOptions, DecodeArgs,
     PyBytesInner,
 };
-use crate::byteslike::PyBytesLike;
+use crate::byteslike::ArgBytesLike;
 use crate::common::borrow::{BorrowedValue, BorrowedValueMut};
 use crate::common::lock::{
     PyMappedRwLockReadGuard, PyMappedRwLockWriteGuard, PyRwLock, PyRwLockReadGuard,
     PyRwLockWriteGuard,
 };
 use crate::function::{FuncArgs, OptionalArg, OptionalOption};
 use crate::sliceable::{PySliceableSequence, PySliceableSequenceMut, SequenceIndex};
-use crate::slots::{
-    BufferProtocol, Comparable, Hashable, Iterable, PyComparisonOp, PyIter, Unhashable,
-};
+use crate::slots::{AsBuffer, Comparable, Hashable, Iterable, PyComparisonOp, PyIter, Unhashable};
 use crate::utils::Either;
 use crate::vm::VirtualMachine;
 use crate::{
@@ -99,7 +97,7 @@ pub(crate) fn init(context: &PyContext) {
     PyByteArrayIterator::extend_class(context, &context.types.bytearray_iterator_type);
 }
 
-#[pyimpl(flags(BASETYPE), with(Hashable, Comparable, BufferProtocol, Iterable))]
+#[pyimpl(flags(BASETYPE), with(Hashable, Comparable, AsBuffer, Iterable))]
 impl PyByteArray {
     #[pyslot]
     fn tp_new(cls: PyTypeRef, _args: FuncArgs, vm: &VirtualMachine) -> PyResult<PyRef<Self>> {
@@ -144,7 +142,7 @@ impl PyByteArray {
     }
 
     #[pymethod(magic)]
-    fn add(&self, other: PyBytesLike, vm: &VirtualMachine) -> PyObjectRef {
+    fn add(&self, other: ArgBytesLike, vm: &VirtualMachine) -> PyObjectRef {
         vm.ctx.new_bytearray(self.inner().add(&*other.borrow_buf()))
     }
 
@@ -192,7 +190,7 @@ impl PyByteArray {
     }
 
     #[pymethod(magic)]
-    fn iadd(zelf: PyRef<Self>, other: PyBytesLike, vm: &VirtualMachine) -> PyResult<PyRef<Self>> {
+    fn iadd(zelf: PyRef<Self>, other: ArgBytesLike, vm: &VirtualMachine) -> PyResult<PyRef<Self>> {
         zelf.try_resizable(vm)?
             .elements
             .extend(&*other.borrow_buf());
@@ -667,7 +665,7 @@ impl Comparable for PyByteArray {
     }
 }
 
-impl BufferProtocol for PyByteArray {
+impl AsBuffer for PyByteArray {
     fn get_buffer(zelf: &PyRef<Self>, _vm: &VirtualMachine) -> PyResult<Box<dyn PyBuffer>> {
         zelf.exports.fetch_add(1);
         let buf = ByteArrayBuffer {

vm/src/builtins/bytes.rs

@@ -1,31 +1,29 @@
-use bstr::ByteSlice;
-use crossbeam_utils::atomic::AtomicCell;
-use rustpython_common::borrow::{BorrowedValue, BorrowedValueMut};
-use std::mem::size_of;
-use std::ops::Deref;
-
 use super::dict::PyDictRef;
 use super::int::PyIntRef;
 use super::pystr::PyStrRef;
 use super::pytype::PyTypeRef;
 use crate::anystr::{self, AnyStr};
+use crate::buffer::{BufferOptions, PyBuffer};
 use crate::builtins::tuple::PyTupleRef;
 use crate::bytesinner::{
     bytes_decode, ByteInnerFindOptions, ByteInnerNewOptions, ByteInnerPaddingOptions,
     ByteInnerSplitOptions, ByteInnerTranslateOptions, DecodeArgs, PyBytesInner,
 };
-use crate::byteslike::PyBytesLike;
+use crate::byteslike::ArgBytesLike;
 use crate::common::hash::PyHash;
 use crate::function::{OptionalArg, OptionalOption};
-use crate::slots::{BufferProtocol, Comparable, Hashable, Iterable, PyComparisonOp, PyIter};
+use crate::slots::{AsBuffer, Comparable, Hashable, Iterable, PyComparisonOp, PyIter};
 use crate::utils::Either;
 use crate::vm::VirtualMachine;
 use crate::{
     IdProtocol, IntoPyObject, PyClassImpl, PyComparisonValue, PyContext, PyIterable, PyObjectRef,
-    PyRef, PyResult, PyValue, TryFromObject, TypeProtocol,
+    PyRef, PyResult, PyValue, TryFromBorrowedObject, TypeProtocol,
 };
-
-use crate::builtins::memory::{BufferOptions, PyBuffer};
+use bstr::ByteSlice;
+use crossbeam_utils::atomic::AtomicCell;
+use rustpython_common::borrow::{BorrowedValue, BorrowedValueMut};
+use std::mem::size_of;
+use std::ops::Deref;
 
 /// "bytes(iterable_of_ints) -> bytes\n\
 /// bytes(string, encoding[, errors]) -> bytes\n\
@@ -98,7 +96,7 @@ pub(crate) fn init(context: &PyContext) {
     PyBytesIterator::extend_class(context, &context.types.bytes_iterator_type);
 }
 
-#[pyimpl(flags(BASETYPE), with(Hashable, Comparable, BufferProtocol, Iterable))]
+#[pyimpl(flags(BASETYPE), with(Hashable, Comparable, AsBuffer, Iterable))]
 impl PyBytes {
     #[pyslot]
     fn tp_new(
@@ -136,7 +134,7 @@ impl PyBytes {
     }
 
     #[pymethod(magic)]
-    fn add(&self, other: PyBytesLike, vm: &VirtualMachine) -> PyObjectRef {
+    fn add(&self, other: ArgBytesLike, vm: &VirtualMachine) -> PyObjectRef {
         vm.ctx.new_bytes(self.inner.add(&*other.borrow_buf()))
     }
 
@@ -366,7 +364,7 @@ impl PyBytes {
 
     #[pymethod]
     fn partition(&self, sep: PyObjectRef, vm: &VirtualMachine) -> PyResult {
-        let sub = PyBytesInner::try_from_object(vm, sep.clone())?;
+        let sub = PyBytesInner::try_from_borrowed_object(vm, &sep)?;
         let (front, has_mid, back) = self.inner.partition(&sub, vm)?;
         Ok(vm.ctx.new_tuple(vec![
             vm.ctx.new_bytes(front),
@@ -381,7 +379,7 @@ impl PyBytes {
 
     #[pymethod]
     fn rpartition(&self, sep: PyObjectRef, vm: &VirtualMachine) -> PyResult {
-        let sub = PyBytesInner::try_from_object(vm, sep.clone())?;
+        let sub = PyBytesInner::try_from_borrowed_object(vm, &sep)?;
         let (back, has_mid, front) = self.inner.rpartition(&sub, vm)?;
         Ok(vm.ctx.new_tuple(vec![
             vm.ctx.new_bytes(front),
@@ -502,7 +500,7 @@ impl PyBytes {
     }
 }
 
-impl BufferProtocol for PyBytes {
+impl AsBuffer for PyBytes {
     fn get_buffer(zelf: &PyRef<Self>, _vm: &VirtualMachine) -> PyResult<Box<dyn PyBuffer>> {
         let buf = BytesBuffer {
             bytes: zelf.clone(),
@@ -604,8 +602,8 @@ impl PyIter for PyBytesIterator {
     }
 }
 
-impl TryFromObject for PyBytes {
-    fn try_from_object(vm: &VirtualMachine, obj: PyObjectRef) -> PyResult<Self> {
-        PyBytesInner::try_from_object(vm, obj).map(|x| x.into())
+impl TryFromBorrowedObject for PyBytes {
+    fn try_from_borrowed_object(vm: &VirtualMachine, obj: &PyObjectRef) -> PyResult<Self> {
+        PyBytesInner::try_from_borrowed_object(vm, obj).map(|x| x.into())
     }
 }

vm/src/builtins/make_module.rs

@@ -20,7 +20,7 @@ mod decl {
     use crate::builtins::pystr::{PyStr, PyStrRef};
     use crate::builtins::pytype::PyTypeRef;
     use crate::builtins::{PyByteArray, PyBytes};
-    use crate::byteslike::PyBytesLike;
+    use crate::byteslike::ArgBytesLike;
     use crate::common::{hash::PyHash, str::to_ascii};
     #[cfg(feature = "rustpython-compiler")]
     use crate::compile;
@@ -563,7 +563,7 @@ mod decl {
     }
 
     #[pyfunction]
-    fn ord(string: Either<PyBytesLike, PyStrRef>, vm: &VirtualMachine) -> PyResult<u32> {
+    fn ord(string: Either<ArgBytesLike, PyStrRef>, vm: &VirtualMachine) -> PyResult<u32> {
         match string {
             Either::A(bytes) => bytes.with_ref(|bytes| {
                 let bytes_len = bytes.len();