Implement #[init] method attribute in #[pymethods] (#4951)

* Implement #[init] method attribute in #[pymethods]

This allows to control objects initialization flow in the Rust code
in case of inheritance from native Python types.

* Apply suggestions from code review

Co-authored-by: Bruno Kolenbrander <[email protected]>

* review feedback

* expose `PySuper` on PyPy, GraalPy and ABI3

* fix graalpy issue

---------

Co-authored-by: David Hewitt <[email protected]>
Co-authored-by: Bruno Kolenbrander <[email protected]>
Co-authored-by: Icxolu <[email protected]>

Author: 4 people

guide/src/class.md

@@ -145,7 +145,7 @@ There is a [detailed discussion on thread-safety](./class/thread-safety.md) late
 
 By default, it is not possible to create an instance of a custom class from Python code.
 To declare a constructor, you need to define a method and annotate it with the `#[new]` attribute.
-Only Python's `__new__` method can be specified, `__init__` is not available.
+A constructor is accessible as Python's `__new__` method.
 
 ```rust
 # #![allow(dead_code)]
@@ -192,6 +192,76 @@ If no method marked with `#[new]` is declared, object instances can only be crea
 
 For arguments, see the [`Method arguments`](#method-arguments) section below.
 
+## Initializer
+
+An initializer implements Python's `__init__` method.
+
+It may be required when it's needed to control an object initalization flow on the Rust code.
+If possible handling this in `__new__` should be preferred, but in some cases, like subclassing native types, overwriting `__init__` might be necessary.
+For example, you define a class that extends `PyDict` and don't want that the original `__init__` method of `PyDict` been called.
+In this case by defining an own `__init__` method it's possible to stop initialization flow.
+
+If you declare an `__init__` method you may need to call a super class' `__init__` method explicitly like in Python code.
+
+To declare an initializer, you need to define the `__init__` method.
+Like in Python `__init__` must have the `self` receiver as the first argument, followed by the same arguments as the constructor.
+It can either return `()` or `PyResult<()>`.
+
+```rust
+# #![allow(dead_code)]
+# use pyo3::prelude::*;
+# #[cfg(not(any(Py_LIMITED_API, GraalPy)))]
+use pyo3::types::{PyDict, PyTuple, PySuper};
+# #[cfg(not(any(Py_LIMITED_API, GraalPy)))]
+use crate::pyo3::PyTypeInfo;
+
+# #[cfg(not(any(Py_LIMITED_API, GraalPy)))]
+#[pyclass(extends = PyDict)]
+struct MyDict;
+
+# #[cfg(not(any(Py_LIMITED_API, GraalPy)))]
+#[pymethods]
+impl MyDict {
+#   #[allow(unused_variables)]
+    #[new]
+    #[pyo3(signature = (*args, **kwargs))]
+    fn __new__(
+        args: &Bound<'_, PyTuple>,
+        kwargs: Option<&Bound<'_, PyDict>>,
+    ) -> PyResult<Self> {
+        Ok(Self)
+    }
+
+    #[pyo3(signature = (*args, **kwargs))]
+    fn __init__(
+        slf: &Bound<'_, Self>,
+        args: &Bound<'_, PyTuple>,
+        kwargs: Option<&Bound<'_, PyDict>>,
+    ) -> PyResult<()> {
+        // call the super types __init__
+        PySuper::new(&PyDict::type_object(slf.py()), slf)?
+            .call_method("__init__", args.to_owned(), kwargs)?;
+        // Note: if `MyDict` allows further subclassing, and this is called from such a subclass,
+        // then this will not that any overrides into account that such a subclass may have defined.
+        // In such a case it may be preferred to just call `slf.set_item` and let Python figure it out.
+        slf.as_super().set_item("my_key", "always insert this key")?;
+        Ok(())
+    }
+}
+
+# #[cfg(not(any(Py_LIMITED_API, GraalPy)))]
+# fn main() {
+#     Python::attach(|py| {
+#         let typeobj = py.get_type::<MyDict>();
+#         let obj = typeobj.call((), None).unwrap().cast_into::<MyDict>().unwrap();
+#         // check __init__ was called
+#         assert_eq!(obj.get_item("my_key").unwrap().extract::<&str>().unwrap(), "always insert this key");
+#     });
+# }
+# #[cfg(any(Py_LIMITED_API, GraalPy))]
+# fn main() {}
+```
+
 ## Adding the class to a module
 
 The next step is to create the Python module and add our class to it:

guide/src/class/protocols.md

@@ -6,7 +6,7 @@ Because of the double-underscores surrounding their name, these are also known a
 
 PyO3 makes it possible for every magic method to be implemented in `#[pymethods]` just as they would be done in a regular Python class, with a few notable differences:
 
-- `__new__` and `__init__` are replaced by the [`#[new]` attribute](../class.md#constructor).
+- `__new__` is replaced by the [`#[new]` attribute](../class.md#constructor).
 - `__del__` is not yet supported, but may be in the future.
 - `__buffer__` and `__release_buffer__` are currently not supported and instead PyO3 supports [`__getbuffer__` and `__releasebuffer__`](#buffer-objects) methods (these predate [PEP 688](https://peps.python.org/pep-0688/#python-level-buffer-protocol)), again this may change in the future.
 - PyO3 adds [`__traverse__` and `__clear__`](#garbage-collector-integration) methods for controlling garbage collection.

newsfragments/4951.added.md

@@ -0,0 +1,2 @@
+Added `__init__` support in `#[pymethods]`.
+expose `PySuper` on PyPy, GraalPy and ABI3

pyo3-macros-backend/src/pymethod.rs

@@ -81,6 +81,7 @@ impl PyMethodKind {
         match name {
             // Protocol implemented through slots
             "__new__" => PyMethodKind::Proto(PyMethodProtoKind::Slot(&__NEW__)),
+            "__init__" => PyMethodKind::Proto(PyMethodProtoKind::Slot(&__INIT__)),
             "__str__" => PyMethodKind::Proto(PyMethodProtoKind::Slot(&__STR__)),
             "__repr__" => PyMethodKind::Proto(PyMethodProtoKind::Slot(&__REPR__)),
             "__hash__" => PyMethodKind::Proto(PyMethodProtoKind::Slot(&__HASH__)),
@@ -312,11 +313,11 @@ fn ensure_no_forbidden_protocol_attributes(
     method_name: &str,
 ) -> syn::Result<()> {
     if let Some(signature) = &spec.signature.attribute {
-        // __new__ and __call__ are allowed to have a signature, but nothing else is.
+        // __new__, __init__ and __call__ are allowed to have a signature, but nothing else is.
         if !matches!(
             proto_kind,
             PyMethodProtoKind::Slot(SlotDef {
-                calling_convention: SlotCallingConvention::TpNew,
+                calling_convention: SlotCallingConvention::TpNew | SlotCallingConvention::TpInit,
                 ..
             })
         ) && !matches!(proto_kind, PyMethodProtoKind::Call)
@@ -367,7 +368,6 @@ pub fn impl_py_method_def(
 
 fn impl_call_slot(cls: &syn::Type, spec: &FnSpec<'_>, ctx: &Ctx) -> Result<MethodAndSlotDef> {
     let Ctx { pyo3_path, .. } = ctx;
-
     let wrapper_ident = syn::Ident::new("__pymethod___call____", Span::call_site());
     let associated_method =
         spec.get_wrapper_function(&wrapper_ident, Some(cls), CallingConvention::Varargs, ctx)?;
@@ -904,6 +904,7 @@ impl PropertyType<'_> {
 }
 
 pub const __NEW__: SlotDef = SlotDef::new("Py_tp_new", "newfunc");
+pub const __INIT__: SlotDef = SlotDef::new("Py_tp_init", "initproc");
 pub const __STR__: SlotDef = SlotDef::new("Py_tp_str", "reprfunc");
 pub const __REPR__: SlotDef = SlotDef::new("Py_tp_repr", "reprfunc");
 pub const __HASH__: SlotDef =
@@ -1164,13 +1165,15 @@ enum SlotCallingConvention {
     FixedArguments(&'static [Ty]),
     /// Arbitrary arguments for `__new__` from the signature (extracted from args / kwargs)
     TpNew,
+    TpInit,
 }
 
 impl SlotDef {
     const fn new(slot: &'static str, func_ty: &'static str) -> Self {
         // The FFI function pointer type determines the arguments and return type
         let (calling_convention, ret_ty) = match func_ty.as_bytes() {
             b"newfunc" => (SlotCallingConvention::TpNew, Ty::Object),
+            b"initproc" => (SlotCallingConvention::TpInit, Ty::Int),
             b"reprfunc" => (SlotCallingConvention::FixedArguments(&[]), Ty::Object),
             b"hashfunc" => (SlotCallingConvention::FixedArguments(&[]), Ty::PyHashT),
             b"richcmpfunc" => (
@@ -1384,6 +1387,35 @@ fn generate_method_body(
             };
             (arg_idents, arg_types, body)
         }
+        SlotCallingConvention::TpInit => {
+            let arg_idents = vec![
+                format_ident!("_slf"),
+                format_ident!("_args"),
+                format_ident!("_kwargs"),
+            ];
+            let arg_types = vec![
+                quote! { *mut #pyo3_path::ffi::PyObject },
+                quote! { *mut #pyo3_path::ffi::PyObject },
+                quote! { *mut #pyo3_path::ffi::PyObject },
+            ];
+            let (arg_convert, args) = impl_arg_params(spec, Some(cls), false, holders, ctx);
+            let call = quote! {{
+                let r = #cls::#rust_name(#self_arg #(#args),*);
+                #pyo3_path::impl_::wrap::converter(&r)
+                    .wrap(r)
+                    .map_err(::core::convert::Into::<#pyo3_path::PyErr>::into)?
+            }};
+            let output = quote_spanned! { *output_span => result.convert(py) };
+
+            let body = quote! {
+                use #pyo3_path::impl_::callback::IntoPyCallbackOutput;
+                #warnings
+                #arg_convert
+                let result = #call;
+                #output
+            };
+            (arg_idents, arg_types, body)
+        }
         SlotCallingConvention::FixedArguments(arguments) => {
             let arg_idents: Vec<_> = std::iter::once(format_ident!("_slf"))
                 .chain((0..arguments.len()).map(|i| format_ident!("arg{}", i)))

pytests/src/pyclasses.rs

@@ -3,6 +3,8 @@ use std::{thread, time};
 use pyo3::exceptions::{PyStopIteration, PyValueError};
 use pyo3::prelude::*;
 use pyo3::types::PyType;
+#[cfg(not(any(Py_LIMITED_API, GraalPy)))]
+use pyo3::types::{PyDict, PyTuple};
 
 #[pyclass(from_py_object)]
 #[derive(Clone, Default)]
@@ -104,6 +106,34 @@ impl ClassWithDict {
     }
 }
 
+#[cfg(not(any(Py_LIMITED_API, GraalPy)))] // Can't subclass native types on abi3 yet
+#[pyclass(extends = PyDict)]
+struct SubClassWithInit;
+
+#[cfg(not(any(Py_LIMITED_API, GraalPy)))]
+#[pymethods]
+impl SubClassWithInit {
+    #[new]
+    #[pyo3(signature = (*args, **kwargs))]
+    #[allow(unused_variables)]
+    fn __new__(args: &Bound<'_, PyTuple>, kwargs: Option<&Bound<'_, PyDict>>) -> Self {
+        Self
+    }
+
+    #[pyo3(signature = (*args, **kwargs))]
+    fn __init__(
+        self_: &Bound<'_, Self>,
+        args: &Bound<'_, PyTuple>,
+        kwargs: Option<&Bound<'_, PyDict>>,
+    ) -> PyResult<()> {
+        self_
+            .py_super()?
+            .call_method("__init__", args.to_owned(), kwargs)?;
+        self_.as_super().set_item("__init__", true)?;
+        Ok(())
+    }
+}
+
 #[pyclass(skip_from_py_object)]
 #[derive(Clone)]
 struct ClassWithDecorators {
@@ -173,6 +203,9 @@ pub mod pyclasses {
     #[cfg(any(Py_3_10, not(Py_LIMITED_API)))]
     #[pymodule_export]
     use super::ClassWithDict;
+    #[cfg(not(any(Py_LIMITED_API, GraalPy)))]
+    #[pymodule_export]
+    use super::SubClassWithInit;
     #[pymodule_export]
     use super::{
         map_a_class, AssertingBaseClass, ClassWithDecorators, ClassWithoutConstructor, EmptyClass,

pytests/stubs/pyclasses.pyi

@@ -1,5 +1,5 @@
 from _typeshed import Incomplete
-from typing import final
+from typing import Any, final
 
 class AssertingBaseClass:
     def __new__(cls, /, expected_type: type) -> AssertingBaseClass: ...
@@ -53,6 +53,11 @@ class PyClassThreadIter:
     def __new__(cls, /) -> PyClassThreadIter: ...
     def __next__(self, /) -> int: ...
 
+@final
+class SubClassWithInit(dict):
+    def __init__(self, /, *args, **kwargs) -> Any: ...
+    def __new__(cls, /, *args, **kwargs) -> SubClassWithInit: ...
+
 def map_a_class(
     cls: EmptyClass | tuple[EmptyClass, EmptyClass] | Incomplete,
 ) -> EmptyClass | tuple[EmptyClass, EmptyClass] | Incomplete: ...

pytests/tests/test_pyclasses.py

@@ -164,3 +164,16 @@ def test_class_method(benchmark):
 def test_static_method(benchmark):
     cls = pyclasses.ClassWithDecorators
     benchmark(lambda: cls.static_method())
+
+
+def test_class_init_method():
+    try:
+        SubClassWithInit = pyclasses.SubClassWithInit
+    except AttributeError:
+        pytest.skip("not defined using abi3")
+
+    d = SubClassWithInit()
+    assert d == {"__init__": True}
+
+    d = SubClassWithInit({"a": 1}, b=2)
+    assert d == {"__init__": True, "a": 1, "b": 2}

src/impl_/trampoline.rs

@@ -162,6 +162,12 @@ trampolines!(
         kwargs: *mut ffi::PyObject,
     ) -> *mut ffi::PyObject;
 
+    pub fn initproc(
+        slf: *mut ffi::PyObject,
+        args: *mut ffi::PyObject,
+        kwargs: *mut ffi::PyObject,
+    ) -> c_int;
+
     pub fn objobjproc(slf: *mut ffi::PyObject, arg1: *mut ffi::PyObject) -> c_int;
 
     pub fn reprfunc(slf: *mut ffi::PyObject) -> *mut ffi::PyObject;

src/types/any.rs

@@ -8,7 +8,6 @@ use crate::instance::Bound;
 use crate::internal::get_slot::TP_DESCR_GET;
 use crate::py_result_ext::PyResultExt;
 use crate::type_object::{PyTypeCheck, PyTypeInfo};
-#[cfg(not(any(PyPy, GraalPy)))]
 use crate::types::PySuper;
 use crate::types::{PyDict, PyIterator, PyList, PyString, PyType};
 use crate::{err, ffi, Borrowed, BoundObject, IntoPyObjectExt, Py, Python};
@@ -939,7 +938,6 @@ pub trait PyAnyMethods<'py>: crate::sealed::Sealed {
     /// Return a proxy object that delegates method calls to a parent or sibling class of type.
     ///
     /// This is equivalent to the Python expression `super()`
-    #[cfg(not(any(PyPy, GraalPy)))]
     fn py_super(&self) -> PyResult<Bound<'py, PySuper>>;
 }
 
@@ -1611,7 +1609,6 @@ impl<'py> PyAnyMethods<'py> for Bound<'py, PyAny> {
         )
     }
 
-    #[cfg(not(any(PyPy, GraalPy)))]
     fn py_super(&self) -> PyResult<Bound<'py, PySuper>> {
         PySuper::new(&self.get_type(), self)
     }

src/types/mod.rs

@@ -39,7 +39,6 @@ pub use self::mutex::{PyMutex, PyMutexGuard};
 pub use self::none::PyNone;
 pub use self::notimplemented::PyNotImplemented;
 pub use self::num::PyInt;
-#[cfg(not(any(PyPy, GraalPy)))]
 pub use self::pysuper::PySuper;
 pub use self::range::{PyRange, PyRangeMethods};
 pub use self::sequence::{PySequence, PySequenceMethods};
@@ -278,7 +277,6 @@ mod mutex;
 mod none;
 mod notimplemented;
 mod num;
-#[cfg(not(any(PyPy, GraalPy)))]
 mod pysuper;
 pub(crate) mod range;
 pub(crate) mod sequence;