Merge pull request #2150 from davidhewitt/python-3.11-buffer

davidhewitt · web-flow · commit 3beb122a830c · 2022-02-06T23:28:04.000Z
python-3.11: support buffer API on abi3
diff --git a/guide/src/building_and_distribution.md b/guide/src/building_and_distribution.md
@@ -163,7 +163,7 @@ not work when compiling for `abi3`. These are:
 
 - `#[pyo3(text_signature = "...")]` does not work on classes until Python 3.10 or greater.
 - The `dict` and `weakref` options on classes are not supported until Python 3.9 or greater.
-- The buffer API is not supported.
+- The buffer API is not supported until Python 3.11 or greater.
 - Optimizations which rely on knowledge of the exact Python version compiled against.
 
 ## Embedding Python in Rust
diff --git a/pyo3-ffi/src/buffer.rs b/pyo3-ffi/src/buffer.rs
@@ -0,0 +1,133 @@
+use crate::object::PyObject;
+use crate::pyport::Py_ssize_t;
+use std::os::raw::{c_char, c_int, c_void};
+use std::ptr;
+
+#[cfg(PyPy)]
+const Py_MAX_NDIMS: usize = 36;
+
+#[repr(C)]
+#[derive(Copy, Clone)]
+pub struct Py_buffer {
+    pub buf: *mut c_void,
+    /// Owned reference
+    pub obj: *mut crate::PyObject,
+    pub len: Py_ssize_t,
+    pub itemsize: Py_ssize_t,
+    pub readonly: c_int,
+    pub ndim: c_int,
+    pub format: *mut c_char,
+    pub shape: *mut Py_ssize_t,
+    pub strides: *mut Py_ssize_t,
+    pub suboffsets: *mut Py_ssize_t,
+    pub internal: *mut c_void,
+    #[cfg(PyPy)]
+    pub flags: c_int,
+    #[cfg(PyPy)]
+    pub _strides: [Py_ssize_t; Py_MAX_NDIMS],
+    #[cfg(PyPy)]
+    pub _shape: [Py_ssize_t; Py_MAX_NDIMS],
+}
+
+impl Py_buffer {
+    pub const fn new() -> Self {
+        Py_buffer {
+            buf: ptr::null_mut(),
+            obj: ptr::null_mut(),
+            len: 0,
+            itemsize: 0,
+            readonly: 0,
+            ndim: 0,
+            format: ptr::null_mut(),
+            shape: ptr::null_mut(),
+            strides: ptr::null_mut(),
+            suboffsets: ptr::null_mut(),
+            internal: ptr::null_mut(),
+            #[cfg(PyPy)]
+            flags: 0,
+            #[cfg(PyPy)]
+            _strides: [0; Py_MAX_NDIMS],
+            #[cfg(PyPy)]
+            _shape: [0; Py_MAX_NDIMS],
+        }
+    }
+}
+
+/* Return 1 if the getbuffer function is available, otherwise return 0. */
+extern "C" {
+    #[cfg(not(PyPy))]
+    pub fn PyObject_CheckBuffer(obj: *mut PyObject) -> c_int;
+
+    #[cfg_attr(PyPy, link_name = "PyPyObject_GetBuffer")]
+    pub fn PyObject_GetBuffer(obj: *mut PyObject, view: *mut Py_buffer, flags: c_int) -> c_int;
+    #[cfg_attr(PyPy, link_name = "PyPyBuffer_GetPointer")]
+    pub fn PyBuffer_GetPointer(view: *const Py_buffer, indices: *const Py_ssize_t) -> *mut c_void;
+    #[cfg_attr(PyPy, link_name = "PyPyBuffer_SizeFromFormat")]
+    pub fn PyBuffer_SizeFromFormat(format: *const c_char) -> Py_ssize_t;
+    #[cfg_attr(PyPy, link_name = "PyPyBuffer_ToContiguous")]
+    pub fn PyBuffer_ToContiguous(
+        buf: *mut c_void,
+        view: *const Py_buffer,
+        len: Py_ssize_t,
+        order: c_char,
+    ) -> c_int;
+    #[cfg_attr(PyPy, link_name = "PyPyBuffer_FromContiguous")]
+    pub fn PyBuffer_FromContiguous(
+        view: *const Py_buffer,
+        buf: *const c_void,
+        len: Py_ssize_t,
+        order: c_char,
+    ) -> c_int;
+    pub fn PyObject_CopyData(dest: *mut PyObject, src: *mut PyObject) -> c_int;
+    #[cfg_attr(PyPy, link_name = "PyPyBuffer_IsContiguous")]
+    pub fn PyBuffer_IsContiguous(view: *const Py_buffer, fort: c_char) -> c_int;
+    pub fn PyBuffer_FillContiguousStrides(
+        ndims: c_int,
+        shape: *mut Py_ssize_t,
+        strides: *mut Py_ssize_t,
+        itemsize: c_int,
+        fort: c_char,
+    );
+    #[cfg_attr(PyPy, link_name = "PyPyBuffer_FillInfo")]
+    pub fn PyBuffer_FillInfo(
+        view: *mut Py_buffer,
+        o: *mut PyObject,
+        buf: *mut c_void,
+        len: Py_ssize_t,
+        readonly: c_int,
+        flags: c_int,
+    ) -> c_int;
+    #[cfg_attr(PyPy, link_name = "PyPyBuffer_Release")]
+    pub fn PyBuffer_Release(view: *mut Py_buffer);
+}
+
+/// Maximum number of dimensions
+pub const PyBUF_MAX_NDIM: c_int = 64;
+
+/* Flags for getting buffers */
+pub const PyBUF_SIMPLE: c_int = 0;
+pub const PyBUF_WRITABLE: c_int = 0x0001;
+/* we used to include an E, backwards compatible alias */
+pub const PyBUF_WRITEABLE: c_int = PyBUF_WRITABLE;
+pub const PyBUF_FORMAT: c_int = 0x0004;
+pub const PyBUF_ND: c_int = 0x0008;
+pub const PyBUF_STRIDES: c_int = 0x0010 | PyBUF_ND;
+pub const PyBUF_C_CONTIGUOUS: c_int = 0x0020 | PyBUF_STRIDES;
+pub const PyBUF_F_CONTIGUOUS: c_int = 0x0040 | PyBUF_STRIDES;
+pub const PyBUF_ANY_CONTIGUOUS: c_int = 0x0080 | PyBUF_STRIDES;
+pub const PyBUF_INDIRECT: c_int = 0x0100 | PyBUF_STRIDES;
+
+pub const PyBUF_CONTIG: c_int = PyBUF_ND | PyBUF_WRITABLE;
+pub const PyBUF_CONTIG_RO: c_int = PyBUF_ND;
+
+pub const PyBUF_STRIDED: c_int = PyBUF_STRIDES | PyBUF_WRITABLE;
+pub const PyBUF_STRIDED_RO: c_int = PyBUF_STRIDES;
+
+pub const PyBUF_RECORDS: c_int = PyBUF_STRIDES | PyBUF_WRITABLE | PyBUF_FORMAT;
+pub const PyBUF_RECORDS_RO: c_int = PyBUF_STRIDES | PyBUF_FORMAT;
+
+pub const PyBUF_FULL: c_int = PyBUF_INDIRECT | PyBUF_WRITABLE | PyBUF_FORMAT;
+pub const PyBUF_FULL_RO: c_int = PyBUF_INDIRECT | PyBUF_FORMAT;
+
+pub const PyBUF_READ: c_int = 0x100;
+pub const PyBUF_WRITE: c_int = 0x200;
diff --git a/pyo3-ffi/src/cpython/abstract_.rs b/pyo3-ffi/src/cpython/abstract_.rs
@@ -1,5 +1,8 @@
-use crate::{PyObject, Py_buffer, Py_ssize_t};
-use std::os::raw::{c_char, c_int, c_void};
+use crate::{PyObject, Py_ssize_t};
+use std::os::raw::{c_char, c_int};
+
+#[cfg(not(Py_3_11))]
+use crate::Py_buffer;
 
 #[cfg(all(Py_3_8, not(PyPy)))]
 use crate::{
@@ -211,6 +214,7 @@ extern "C" {
     #[cfg_attr(PyPy, link_name = "PyPyObject_LengthHint")]
     pub fn PyObject_LengthHint(o: *mut PyObject, arg1: Py_ssize_t) -> Py_ssize_t;
 
+    #[cfg(not(Py_3_11))] // moved to src/buffer.rs from 3.11
     #[cfg(all(Py_3_9, not(PyPy)))]
     pub fn PyObject_CheckBuffer(obj: *mut PyObject) -> c_int;
 }
@@ -222,24 +226,28 @@ pub unsafe fn PyObject_CheckBuffer(o: *mut PyObject) -> c_int {
     (!tp_as_buffer.is_null() && (*tp_as_buffer).bf_getbuffer.is_some()) as c_int
 }
 
+#[cfg(not(Py_3_11))] // moved to src/buffer.rs from 3.11
 extern "C" {
     #[cfg_attr(PyPy, link_name = "PyPyObject_GetBuffer")]
     pub fn PyObject_GetBuffer(obj: *mut PyObject, view: *mut Py_buffer, flags: c_int) -> c_int;
     #[cfg_attr(PyPy, link_name = "PyPyBuffer_GetPointer")]
-    pub fn PyBuffer_GetPointer(view: *mut Py_buffer, indices: *mut Py_ssize_t) -> *mut c_void;
+    pub fn PyBuffer_GetPointer(
+        view: *mut Py_buffer,
+        indices: *mut Py_ssize_t,
+    ) -> *mut std::os::raw::c_void;
     #[cfg_attr(PyPy, link_name = "PyPyBuffer_SizeFromFormat")]
     pub fn PyBuffer_SizeFromFormat(format: *const c_char) -> Py_ssize_t;
     #[cfg_attr(PyPy, link_name = "PyPyBuffer_ToContiguous")]
     pub fn PyBuffer_ToContiguous(
-        buf: *mut c_void,
+        buf: *mut std::os::raw::c_void,
         view: *mut Py_buffer,
         len: Py_ssize_t,
         order: c_char,
     ) -> c_int;
     #[cfg_attr(PyPy, link_name = "PyPyBuffer_FromContiguous")]
     pub fn PyBuffer_FromContiguous(
         view: *mut Py_buffer,
-        buf: *mut c_void,
+        buf: *mut std::os::raw::c_void,
         len: Py_ssize_t,
         order: c_char,
     ) -> c_int;
@@ -257,7 +265,7 @@ extern "C" {
     pub fn PyBuffer_FillInfo(
         view: *mut Py_buffer,
         o: *mut PyObject,
-        buf: *mut c_void,
+        buf: *mut std::os::raw::c_void,
         len: Py_ssize_t,
         readonly: c_int,
         flags: c_int,
diff --git a/pyo3-ffi/src/cpython/object.rs b/pyo3-ffi/src/cpython/object.rs
@@ -11,6 +11,7 @@ use std::os::raw::c_int;
 // skipped _Py_static_string
 // skipped _Py_IDENTIFIER
 
+#[cfg(not(Py_3_11))] // moved to src/buffer.rs from Python
 mod bufferinfo {
     use crate::Py_ssize_t;
     use std::os::raw::{c_char, c_int, c_void};
@@ -106,6 +107,7 @@ mod bufferinfo {
     pub const PyBUF_WRITE: c_int = 0x200;
 }
 
+#[cfg(not(Py_3_11))]
 pub use self::bufferinfo::*;
 
 #[cfg(Py_3_8)]
@@ -199,8 +201,8 @@ mod typeobject {
     #[repr(C)]
     #[derive(Clone, Default)]
     pub struct PyBufferProcs {
-        pub bf_getbuffer: Option<super::getbufferproc>,
-        pub bf_releasebuffer: Option<super::releasebufferproc>,
+        pub bf_getbuffer: Option<crate::getbufferproc>,
+        pub bf_releasebuffer: Option<crate::releasebufferproc>,
     }
 
     pub type printfunc =
diff --git a/pyo3-ffi/src/lib.rs b/pyo3-ffi/src/lib.rs
@@ -121,7 +121,7 @@
 //!         ml_flags: METH_FASTCALL,
 //!         ml_doc: "returns the sum of two integers as a string\0".as_ptr() as *const c_char,
 //!     };
-//!     
+//!
 //!     // PyModule_AddObject can technically fail.
 //!     // For more involved applications error checking may be necessary
 //!     PyModule_AddObject(
@@ -268,6 +268,8 @@ macro_rules! opaque_struct {
 pub use self::abstract_::*;
 pub use self::bltinmodule::*;
 pub use self::boolobject::*;
+#[cfg(Py_3_11)]
+pub use self::buffer::*;
 pub use self::bytearrayobject::*;
 pub use self::bytesobject::*;
 pub use self::ceval::*;
@@ -332,6 +334,8 @@ mod abstract_;
 // skipped ast.h
 mod bltinmodule;
 mod boolobject;
+#[cfg(Py_3_11)]
+mod buffer;
 mod bytearrayobject;
 mod bytesobject;
 // skipped cellobject.h
diff --git a/pyo3-ffi/src/object.rs b/pyo3-ffi/src/object.rs
@@ -162,6 +162,11 @@ pub type newfunc = unsafe extern "C" fn(
 ) -> *mut PyObject;
 pub type allocfunc =
     unsafe extern "C" fn(arg1: *mut PyTypeObject, arg2: Py_ssize_t) -> *mut PyObject;
+#[cfg(Py_3_11)]
+pub type getbufferproc =
+    unsafe extern "C" fn(arg1: *mut PyObject, arg2: *mut crate::Py_buffer, arg3: c_int) -> c_int;
+#[cfg(Py_3_11)]
+pub type releasebufferproc = unsafe extern "C" fn(arg1: *mut PyObject, arg2: *mut crate::Py_buffer);
 
 #[repr(C)]
 #[derive(Copy, Clone)]
diff --git a/pyo3-ffi/src/typeslots.rs b/pyo3-ffi/src/typeslots.rs
@@ -1,8 +1,6 @@
 use std::os::raw::c_int;
 
-#[cfg(not(Py_LIMITED_API))]
 pub const Py_bf_getbuffer: c_int = 1;
-#[cfg(not(Py_LIMITED_API))]
 pub const Py_bf_releasebuffer: c_int = 2;
 pub const Py_mp_ass_subscript: c_int = 3;
 pub const Py_mp_length: c_int = 4;
diff --git a/src/buffer.rs b/src/buffer.rs
@@ -1,4 +1,4 @@
-#![cfg(not(Py_LIMITED_API))]
+#![cfg(any(not(Py_LIMITED_API), Py_3_11))]
 // Copyright (c) 2017 Daniel Grunwald
 //
 // Permission is hereby granted, free of charge, to any person obtaining a copy of this
@@ -235,8 +235,20 @@ impl<T: Element> PyBuffer<T> {
         }
         unsafe {
             ffi::PyBuffer_GetPointer(
-                &*self.0 as *const ffi::Py_buffer as *mut ffi::Py_buffer,
-                indices.as_ptr() as *mut usize as *mut ffi::Py_ssize_t,
+                #[cfg(Py_3_11)]
+                &*self.0,
+                #[cfg(not(Py_3_11))]
+                {
+                    &*self.0 as *const ffi::Py_buffer as *mut ffi::Py_buffer
+                },
+                #[cfg(Py_3_11)]
+                {
+                    indices.as_ptr() as *const ffi::Py_ssize_t
+                },
+                #[cfg(not(Py_3_11))]
+                {
+                    indices.as_ptr() as *mut ffi::Py_ssize_t
+                },
             )
         }
     }
@@ -477,7 +489,12 @@ impl<T: Element> PyBuffer<T> {
                 py,
                 ffi::PyBuffer_ToContiguous(
                     target.as_ptr() as *mut raw::c_void,
-                    &*self.0 as *const ffi::Py_buffer as *mut ffi::Py_buffer,
+                    #[cfg(Py_3_11)]
+                    &*self.0,
+                    #[cfg(not(Py_3_11))]
+                    {
+                        &*self.0 as *const ffi::Py_buffer as *mut ffi::Py_buffer
+                    },
                     self.0.len,
                     fort as std::os::raw::c_char,
                 ),
@@ -510,8 +527,13 @@ impl<T: Element> PyBuffer<T> {
             err::error_on_minusone(
                 py,
                 ffi::PyBuffer_ToContiguous(
-                    vec.as_mut_ptr() as *mut raw::c_void,
-                    &*self.0 as *const ffi::Py_buffer as *mut ffi::Py_buffer,
+                    vec.as_ptr() as *mut raw::c_void,
+                    #[cfg(Py_3_11)]
+                    &*self.0,
+                    #[cfg(not(Py_3_11))]
+                    {
+                        &*self.0 as *const ffi::Py_buffer as *mut ffi::Py_buffer
+                    },
                     self.0.len,
                     fort as std::os::raw::c_char,
                 ),
@@ -564,8 +586,20 @@ impl<T: Element> PyBuffer<T> {
             err::error_on_minusone(
                 py,
                 ffi::PyBuffer_FromContiguous(
-                    &*self.0 as *const ffi::Py_buffer as *mut ffi::Py_buffer,
-                    source.as_ptr() as *mut raw::c_void,
+                    #[cfg(Py_3_11)]
+                    &*self.0,
+                    #[cfg(not(Py_3_11))]
+                    {
+                        &*self.0 as *const ffi::Py_buffer as *mut ffi::Py_buffer
+                    },
+                    #[cfg(Py_3_11)]
+                    {
+                        source.as_ptr() as *const raw::c_void
+                    },
+                    #[cfg(not(Py_3_11))]
+                    {
+                        source.as_ptr() as *mut raw::c_void
+                    },
                     self.0.len,
                     fort as std::os::raw::c_char,
                 ),
@@ -678,6 +712,8 @@ mod tests {
             assert_eq!(slice[0].get(), b'a');
             assert_eq!(slice[2].get(), b'c');
 
+            assert_eq!(unsafe { *(buffer.get_ptr(&[1]) as *mut u8) }, b'b');
+
             assert!(buffer.copy_to_slice(py, &mut [0u8]).is_err());
             let mut arr = [0; 5];
             buffer.copy_to_slice(py, &mut arr).unwrap();
diff --git a/src/class/mod.rs b/src/class/mod.rs
@@ -6,7 +6,7 @@
 mod macros;
 
 pub mod basic;
-#[cfg(not(Py_LIMITED_API))]
+#[cfg(any(not(Py_LIMITED_API), Py_3_11))]
 pub mod buffer;
 pub mod descr;
 pub mod gc;
@@ -19,7 +19,7 @@ pub mod pyasync;
 pub mod sequence;
 
 pub use self::basic::PyObjectProtocol;
-#[cfg(not(Py_LIMITED_API))]
+#[cfg(any(not(Py_LIMITED_API), Py_3_11))]
 pub use self::buffer::PyBufferProtocol;
 pub use self::descr::PyDescrProtocol;
 pub use self::gc::{PyGCProtocol, PyTraverseError, PyVisit};
diff --git a/tests/test_buffer.rs b/tests/test_buffer.rs
@@ -1,5 +1,5 @@
 #![cfg(feature = "macros")]
-#![cfg(not(Py_LIMITED_API))]
+#![cfg(any(not(Py_LIMITED_API), Py_3_11))]
 
 use pyo3::{buffer::PyBuffer, exceptions::PyBufferError, ffi, prelude::*, AsPyPointer};
 use std::{
diff --git a/tests/test_buffer_protocol.rs b/tests/test_buffer_protocol.rs
diff --git a/tests/test_buffer_protocol_pyproto.rs b/tests/test_buffer_protocol_pyproto.rs
