pymethods: support buffer protocol

CHANGELOG.md

@@ -24,6 +24,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - All PyO3 proc-macros except the deprecated `#[pyproto]` now accept a supplemental attribute `#[pyo3(crate = "some::path")]` that specifies
   where to find the `pyo3` crate, in case it has been renamed or is re-exported and not found at the crate root. [#2022](https://github.com/PyO3/pyo3/pull/2022)
 - Expose `pyo3-build-config` APIs for cross-compiling and Python configuration discovery for use in other projects. [#1996](https://github.com/PyO3/pyo3/pull/1996)
+- Add buffer magic methods `__getbuffer__` and `__releasebuffer__` to `#[pymethods]`. [#2067](https://github.com/PyO3/pyo3/pull/2067)
 
 ### Changed
 
@@ -39,7 +40,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - `PyErr::new_type` now takes an optional docstring and now returns `PyResult<Py<PyType>>` rather than a `ffi::PyTypeObject` pointer.
 - The `create_exception!` macro can now take an optional docstring. This docstring, if supplied, is visible to users (with `.__doc__` and `help()`) and
   accompanies your error type in your crate's documentation.
-  
+
 ### Removed
 
 - Remove all functionality deprecated in PyO3 0.14. [#2007](https://github.com/PyO3/pyo3/pull/2007)

guide/src/class.md

@@ -899,12 +899,6 @@ impl pyo3::class::impl_::PyClassImpl for MyClass {
         visitor(collector.buffer_protocol_slots());
         visitor(collector.methods_protocol_slots());
     }
-
-    fn get_buffer() -> Option<&'static pyo3::class::impl_::PyBufferProcs> {
-        use pyo3::class::impl_::*;
-        let collector = PyClassImplCollector::<Self>::new();
-        collector.buffer_procs()
-    }
 }
 # Python::with_gil(|py| {
 #     let cls = py.get_type::<MyClass>();

guide/src/class/protocols.md

@@ -175,7 +175,8 @@ TODO; see [#1884](https://github.com/PyO3/pyo3/issues/1884)
 
 #### Buffer objects
 
-TODO; see [#1884](https://github.com/PyO3/pyo3/issues/1884)
+  - `__getbuffer__(<self>, *mut ffi::Py_buffer, flags) -> ()`
+  - `__releasebuffer__(<self>, *mut ffi::Py_buffer)` (no return value, not even `PyResult`)
 
 #### Garbage Collector Integration
 

pyo3-macros-backend/src/method.rs

@@ -229,6 +229,7 @@ pub struct FnSpec<'a> {
     pub convention: CallingConvention,
     pub text_signature: Option<TextSignatureAttribute>,
     pub krate: syn::Path,
+    pub unsafety: Option<syn::Token![unsafe]>,
 }
 
 pub fn get_return_info(output: &syn::ReturnType) -> syn::Type {
@@ -316,6 +317,7 @@ impl<'a> FnSpec<'a> {
             deprecations,
             text_signature,
             krate,
+            unsafety: sig.unsafety,
         })
     }
 

pyo3-macros-backend/src/pyclass.rs

@@ -829,12 +829,6 @@ impl<'a> PyClassImplsBuilder<'a> {
                     visitor(collector.buffer_protocol_slots());
                     #methods_protos
                 }
-
-                fn get_buffer() -> ::std::option::Option<&'static _pyo3::class::impl_::PyBufferProcs> {
-                    use _pyo3::class::impl_::*;
-                    let collector = PyClassImplCollector::<Self>::new();
-                    collector.buffer_procs()
-                }
             }
 
             #inventory_class

pyo3-macros-backend/src/pyfunction.rs

@@ -436,6 +436,7 @@ pub fn impl_wrap_pyfunction(
         deprecations: options.deprecations,
         text_signature: options.text_signature,
         krate: krate.clone(),
+        unsafety: func.sig.unsafety,
     };
 
     let wrapper_ident = format_ident!("__pyo3_raw_{}", spec.name);

pyo3-macros-backend/src/pymethod.rs

@@ -103,7 +103,7 @@ pub fn gen_py_method(
             ensure_no_forbidden_protocol_attributes(spec, &method.method_name)?;
             match proto_kind {
                 PyMethodProtoKind::Slot(slot_def) => {
-                    let slot = slot_def.generate_type_slot(cls, spec)?;
+                    let slot = slot_def.generate_type_slot(cls, spec, &method.method_name)?;
                     GeneratedPyMethod::Proto(slot)
                 }
                 PyMethodProtoKind::Call => {
@@ -556,6 +556,14 @@ const __IOR__: SlotDef = SlotDef::new("Py_nb_inplace_or", "binaryfunc")
     .arguments(&[Ty::Object])
     .extract_error_mode(ExtractErrorMode::NotImplemented)
     .return_self();
+const __GETBUFFER__: SlotDef = SlotDef::new("Py_bf_getbuffer", "getbufferproc")
+    .arguments(&[Ty::PyBuffer, Ty::Int])
+    .ret_ty(Ty::Int)
+    .require_unsafe();
+const __RELEASEBUFFER__: SlotDef = SlotDef::new("Py_bf_releasebuffer", "releasebufferproc")
+    .arguments(&[Ty::PyBuffer])
+    .ret_ty(Ty::Void)
+    .require_unsafe();
 
 fn pyproto(method_name: &str) -> Option<&'static SlotDef> {
     match method_name {
@@ -594,6 +602,8 @@ fn pyproto(method_name: &str) -> Option<&'static SlotDef> {
         "__iand__" => Some(&__IAND__),
         "__ixor__" => Some(&__IXOR__),
         "__ior__" => Some(&__IOR__),
+        "__getbuffer__" => Some(&__GETBUFFER__),
+        "__releasebuffer__" => Some(&__RELEASEBUFFER__),
         _ => None,
     }
 }
@@ -608,6 +618,7 @@ enum Ty {
     PyHashT,
     PySsizeT,
     Void,
+    PyBuffer,
 }
 
 impl Ty {
@@ -619,6 +630,7 @@ impl Ty {
             Ty::PyHashT => quote! { _pyo3::ffi::Py_hash_t },
             Ty::PySsizeT => quote! { _pyo3::ffi::Py_ssize_t },
             Ty::Void => quote! { () },
+            Ty::PyBuffer => quote! { *mut _pyo3::ffi::Py_buffer },
         }
     }
 
@@ -680,7 +692,8 @@ impl Ty {
                     let #ident = #extract;
                 }
             }
-            Ty::Int | Ty::PyHashT | Ty::PySsizeT | Ty::Void => todo!(),
+            // Just pass other types through unmodified
+            Ty::PyBuffer | Ty::Int | Ty::PyHashT | Ty::PySsizeT | Ty::Void => quote! {},
         }
     }
 }
@@ -752,6 +765,7 @@ struct SlotDef {
     before_call_method: Option<TokenGenerator>,
     extract_error_mode: ExtractErrorMode,
     return_mode: Option<ReturnMode>,
+    require_unsafe: bool,
 }
 
 const NO_ARGUMENTS: &[Ty] = &[];
@@ -766,6 +780,7 @@ impl SlotDef {
             before_call_method: None,
             extract_error_mode: ExtractErrorMode::Raise,
             return_mode: None,
+            require_unsafe: false,
         }
     }
 
@@ -799,7 +814,17 @@ impl SlotDef {
         self
     }
 
-    fn generate_type_slot(&self, cls: &syn::Type, spec: &FnSpec) -> Result<TokenStream> {
+    const fn require_unsafe(mut self) -> Self {
+        self.require_unsafe = true;
+        self
+    }
+
+    fn generate_type_slot(
+        &self,
+        cls: &syn::Type,
+        spec: &FnSpec,
+        method_name: &str,
+    ) -> Result<TokenStream> {
         let SlotDef {
             slot,
             func_ty,
@@ -808,7 +833,14 @@ impl SlotDef {
             extract_error_mode,
             ret_ty,
             return_mode,
+            require_unsafe,
         } = self;
+        if *require_unsafe {
+            ensure_spanned!(
+                spec.unsafety.is_some(),
+                spec.name.span() => format!("`{}` must be `unsafe fn`", method_name)
+            );
+        }
         let py = syn::Ident::new("_py", Span::call_site());
         let method_arguments = generate_method_arguments(arguments);
         let ret_ty = ret_ty.ffi_type();

pyo3-macros-backend/src/pyproto.rs

@@ -134,31 +134,6 @@ fn impl_proto_methods(
     let slots_trait = proto.slots_trait();
     let slots_trait_slots = proto.slots_trait_slots();
 
-    let mut maybe_buffer_methods = None;
-
-    let build_config = pyo3_build_config::get();
-    const PY39: pyo3_build_config::PythonVersion =
-        pyo3_build_config::PythonVersion { major: 3, minor: 9 };
-
-    if build_config.version <= PY39 && proto.name == "Buffer" {
-        maybe_buffer_methods = Some(quote! {
-            impl _pyo3::class::impl_::PyBufferProtocolProcs<#ty>
-                for _pyo3::class::impl_::PyClassImplCollector<#ty>
-            {
-                fn buffer_procs(
-                    self
-                ) -> ::std::option::Option<&'static _pyo3::class::impl_::PyBufferProcs> {
-                    static PROCS: _pyo3::class::impl_::PyBufferProcs
-                        = _pyo3::class::impl_::PyBufferProcs {
-                            bf_getbuffer: ::std::option::Option::Some(_pyo3::class::buffer::getbuffer::<#ty>),
-                            bf_releasebuffer: ::std::option::Option::Some(_pyo3::class::buffer::releasebuffer::<#ty>),
-                        };
-                    ::std::option::Option::Some(&PROCS)
-                }
-            }
-        });
-    }
-
     let mut tokens = proto
         .slot_defs(method_names)
         .map(|def| {
@@ -178,8 +153,6 @@ fn impl_proto_methods(
     }
 
     quote! {
-        #maybe_buffer_methods
-
         impl _pyo3::class::impl_::#slots_trait<#ty>
             for _pyo3::class::impl_::PyClassImplCollector<#ty>
         {

src/class/impl_.rs

@@ -81,9 +81,6 @@ pub trait PyClassImpl: Sized {
         None
     }
     fn for_each_proto_slot(_visitor: &mut dyn FnMut(&[ffi::PyType_Slot])) {}
-    fn get_buffer() -> Option<&'static PyBufferProcs> {
-        None
-    }
 }
 
 // Traits describing known special methods.
@@ -673,25 +670,6 @@ methods_trait!(PyDescrProtocolMethods, descr_protocol_methods);
 methods_trait!(PyMappingProtocolMethods, mapping_protocol_methods);
 methods_trait!(PyNumberProtocolMethods, number_protocol_methods);
 
-// On Python < 3.9 setting the buffer protocol using slots doesn't work, so these procs are used
-// on those versions to set the slots manually (on the limited API).
-
-#[cfg(not(Py_LIMITED_API))]
-pub use ffi::PyBufferProcs;
-
-#[cfg(Py_LIMITED_API)]
-pub struct PyBufferProcs;
-
-pub trait PyBufferProtocolProcs<T> {
-    fn buffer_procs(self) -> Option<&'static PyBufferProcs>;
-}
-
-impl<T> PyBufferProtocolProcs<T> for &'_ PyClassImplCollector<T> {
-    fn buffer_procs(self) -> Option<&'static PyBufferProcs> {
-        None
-    }
-}
-
 // Thread checkers
 
 #[doc(hidden)]

src/pyclass.rs

@@ -105,10 +105,26 @@ where
 
     // protocol methods
     let mut has_gc_methods = false;
+    // Before Python 3.9, need to patch in buffer methods manually (they don't work in slots)
+    #[cfg(all(not(Py_3_9), not(Py_LIMITED_API)))]
+    let mut buffer_procs: ffi::PyBufferProcs = Default::default();
+
     T::for_each_proto_slot(&mut |proto_slots| {
-        has_gc_methods |= proto_slots
-            .iter()
-            .any(|slot| slot.slot == ffi::Py_tp_clear || slot.slot == ffi::Py_tp_traverse);
+        for slot in proto_slots {
+            has_gc_methods |= slot.slot == ffi::Py_tp_clear || slot.slot == ffi::Py_tp_traverse;
+
+            #[cfg(all(not(Py_3_9), not(Py_LIMITED_API)))]
+            if slot.slot == ffi::Py_bf_getbuffer {
+                // Safety: slot.pfunc is a valid function pointer
+                buffer_procs.bf_getbuffer = Some(unsafe { std::mem::transmute(slot.pfunc) });
+            }
+
+            #[cfg(all(not(Py_3_9), not(Py_LIMITED_API)))]
+            if slot.slot == ffi::Py_bf_releasebuffer {
+                // Safety: slot.pfunc is a valid function pointer
+                buffer_procs.bf_releasebuffer = Some(unsafe { std::mem::transmute(slot.pfunc) });
+            }
+        }
         slots.0.extend_from_slice(proto_slots);
     });
 
@@ -125,14 +141,21 @@ where
     if type_object.is_null() {
         Err(PyErr::fetch(py))
     } else {
-        tp_init_additional::<T>(type_object as _);
+        tp_init_additional::<T>(
+            type_object as _,
+            #[cfg(all(not(Py_3_9), not(Py_LIMITED_API)))]
+            &buffer_procs,
+        );
         Ok(type_object as _)
     }
 }
 
 /// Additional type initializations necessary before Python 3.10
 #[cfg(all(not(Py_LIMITED_API), not(Py_3_10)))]
-fn tp_init_additional<T: PyClass>(type_object: *mut ffi::PyTypeObject) {
+fn tp_init_additional<T: PyClass>(
+    type_object: *mut ffi::PyTypeObject,
+    #[cfg(all(not(Py_3_9), not(Py_LIMITED_API)))] buffer_procs: &ffi::PyBufferProcs,
+) {
     // Just patch the type objects for the things there's no
     // PyType_FromSpec API for... there's no reason this should work,
     // except for that it does and we have tests.
@@ -154,22 +177,15 @@ fn tp_init_additional<T: PyClass>(type_object: *mut ffi::PyTypeObject) {
         }
     }
 
-    // Setting buffer protocols via slots doesn't work until Python 3.9, so on older versions we
-    // must manually fixup the type object.
+    // Setting buffer protocols, tp_dictoffset and tp_weaklistoffset via slots doesn't work until
+    // Python 3.9, so on older versions we must manually fixup the type object.
     #[cfg(not(Py_3_9))]
     {
-        if let Some(buffer) = T::get_buffer() {
-            unsafe {
-                (*(*type_object).tp_as_buffer).bf_getbuffer = buffer.bf_getbuffer;
-                (*(*type_object).tp_as_buffer).bf_releasebuffer = buffer.bf_releasebuffer;
-            }
+        unsafe {
+            (*(*type_object).tp_as_buffer).bf_getbuffer = buffer_procs.bf_getbuffer;
+            (*(*type_object).tp_as_buffer).bf_releasebuffer = buffer_procs.bf_releasebuffer;
         }
-    }
 
-    // Setting tp_dictoffset and tp_weaklistoffset via slots doesn't work until Python 3.9, so on
-    // older versions again we must fixup the type object.
-    #[cfg(not(Py_3_9))]
-    {
         // __dict__ support
         if let Some(dict_offset) = PyCell::<T>::dict_offset() {
             unsafe {