fix type confusion when different #[pyclass] types returned from #[new] (#6062)

* fix type confusion when different `#[pyclass]` types returned from `#[new]`

* fixup

* newsfragment

* don't quote `unsafe` at user code

* review feedback

Author: davidhewitt

newsfragments/6062.fixed.md

@@ -0,0 +1 @@
+Fix type confusion when returning a `#[pyclass]` from a different pyclass' `#[new]` method.

pyo3-macros-backend/src/method.rs

@@ -449,6 +449,7 @@ pub struct FnSpec<'a> {
     pub asyncness: Option<syn::Token![async]>,
     pub unsafety: Option<syn::Token![unsafe]>,
     pub warnings: Vec<PyFunctionWarning>,
+    #[cfg_attr(not(feature = "experimental-inspect"), expect(dead_code))]
     pub output: syn::ReturnType,
 }
 

pyo3-macros-backend/src/pymethod.rs

@@ -16,8 +16,8 @@ use crate::{
 use crate::{quotes, utils};
 use proc_macro2::{Span, TokenStream};
 use quote::{format_ident, quote, quote_spanned, ToTokens};
+use syn::LitCStr;
 use syn::{ext::IdentExt, spanned::Spanned, Field, Ident, Result};
-use syn::{parse_quote, LitCStr};
 
 /// Generated code for a single pymethod item.
 pub struct MethodAndMethodDef {
@@ -1473,25 +1473,25 @@ fn generate_method_body(
                 }
             });
 
-            let output = if let syn::ReturnType::Type(_, ty) = &spec.output {
-                let mut ty = ty.clone();
-                utils::elide_lifetimes(&mut ty);
-                ty
-            } else {
-                parse_quote!(())
+            let py = syn::Ident::new("py", Span::call_site());
+            let initializer = syn::Ident::new("initializer", Span::call_site());
+            let slf = syn::Ident::new("_slf", Span::call_site());
+
+            // Having just this call emitted at the span of the return value helps surface errors
+            // if the user passed an invalid return type.
+            let conversion = quote_spanned! { *output_span =>
+                #pyo3_path::impl_::pymethods::tp_new_impl::<_, #cls>(#py, #initializer, #slf)
             };
+
             let body = quote! {
                 #text_signature_impl
-
-                use #pyo3_path::impl_::pyclass::Probe as _;
                 #warnings
                 #arg_convert
+
                 let result = #call;
-                #pyo3_path::impl_::pymethods::tp_new_impl::<
-                    _,
-                    { #pyo3_path::impl_::pyclass::IsPyClass::<#output>::VALUE },
-                    { #pyo3_path::impl_::pyclass::IsInitializerTuple::<#output>::VALUE }
-                >(py, result, _slf)
+                let value = #pyo3_path::impl_::wrap::OkWrapper::new(&result).ok_wrap(result)?;
+                let #initializer = #pyo3_path::impl_::pymethods::tp_new_resolver::<#cls, _>(&value).resolve(value);
+                unsafe { #conversion }
             };
             (arg_idents, arg_types, body)
         }

pyo3-macros-backend/src/utils.rs

@@ -380,6 +380,7 @@ pub(crate) fn locate_tokens_at(tokens: TokenStream, span: Span) -> TokenStream {
 ///
 /// This is useful if `Self` is used in `const` context, where explicit
 /// lifetimes are not allowed (yet).
+#[cfg(feature = "experimental-inspect")]
 pub(crate) fn elide_lifetimes(ty: &mut syn::Type) {
     struct ElideLifetimesVisitor;
 

src/impl_/pyclass.rs

@@ -5,9 +5,9 @@ use crate::{
     impl_::{
         freelist::PyObjectFreeList,
         pycell::{GetBorrowChecker, PyClassMutability, PyClassObjectBaseLayout},
-        pyclass_init::PyObjectInit,
         pymethods::{PyGetterDef, PyMethodDefType},
     },
+    internal::pyclass_init::PyObjectInit,
     pycell::{impl_::PyClassObjectLayout, PyBorrowError},
     types::{any::PyAnyMethods, PyBool},
     Borrowed, IntoPyObject, IntoPyObjectExt, Py, PyAny, PyClass, PyClassGuard, PyErr, PyResult,
@@ -1087,6 +1087,10 @@ impl<T> PyClassThreadChecker<T> for ThreadCheckerImpl {
         note = "subclassing native types requires Python >= 3.12 when using the `abi3` feature",
     )
 )]
+#[expect(
+    private_bounds,
+    reason = "`PyObjectInit` is an internal trait implementation"
+)]
 pub trait PyClassBaseType: Sized {
     type LayoutAsBase: PyClassObjectBaseLayout<Self>;
     type BaseNativeType;

src/impl_/pyclass_init.rs

@@ -1,172 +1,2 @@
-//! Contains initialization utilities for `#[pyclass]`.
-use crate::exceptions::PyTypeError;
-use crate::ffi_ptr_ext::FfiPtrExt;
-use crate::impl_::pyclass::PyClassBaseType;
-use crate::internal::get_slot::TP_NEW;
-use crate::types::{PyTuple, PyType};
-use crate::{ffi, PyClass, PyClassInitializer, PyErr, PyResult, Python};
-use crate::{ffi::PyTypeObject, sealed::Sealed, type_object::PyTypeInfo};
-use core::marker::PhantomData;
-
-/// Initializer for Python types.
-///
-/// This trait is intended to use internally for distinguishing `#[pyclass]` and
-/// Python native types.
-pub trait PyObjectInit<T>: Sized + Sealed {
-    /// # Safety
-    /// - `subtype` must be a valid pointer to a type object of T or a subclass.
-    unsafe fn into_new_object(
-        self,
-        py: Python<'_>,
-        subtype: *mut PyTypeObject,
-    ) -> PyResult<*mut ffi::PyObject>;
-}
-
-/// Initializer for Python native types, like `PyDict`.
-pub struct PyNativeTypeInitializer<T: PyTypeInfo>(pub PhantomData<T>);
-
-impl<T: PyTypeInfo> PyObjectInit<T> for PyNativeTypeInitializer<T> {
-    unsafe fn into_new_object(
-        self,
-        py: Python<'_>,
-        subtype: *mut PyTypeObject,
-    ) -> PyResult<*mut ffi::PyObject> {
-        unsafe fn inner(
-            py: Python<'_>,
-            type_ptr: *mut PyTypeObject,
-            subtype: *mut PyTypeObject,
-        ) -> PyResult<*mut ffi::PyObject> {
-            let tp_new = unsafe {
-                type_ptr
-                    .cast::<ffi::PyObject>()
-                    .assume_borrowed_unchecked(py)
-                    .cast_unchecked::<PyType>()
-                    .get_slot(TP_NEW)
-                    .ok_or_else(|| PyTypeError::new_err("base type without tp_new"))?
-            };
-
-            // TODO: make it possible to provide real arguments to the base tp_new
-            let obj =
-                unsafe { tp_new(subtype, PyTuple::empty(py).as_ptr(), core::ptr::null_mut()) };
-            if obj.is_null() {
-                Err(PyErr::fetch(py))
-            } else {
-                Ok(obj)
-            }
-        }
-        unsafe { inner(py, T::type_object_raw(py), subtype) }
-    }
-}
-
-pub trait PyClassInit<'py, const IS_PYCLASS: bool, const IS_INITIALIZER_TUPLE: bool>:
-    seal_pyclass_init::Sealed<'py, IS_PYCLASS, IS_INITIALIZER_TUPLE>
-{
-    fn init(
-        self,
-        cls: crate::Borrowed<'_, 'py, crate::types::PyType>,
-    ) -> PyResult<crate::Bound<'py, crate::PyAny>>;
-}
-
-mod seal_pyclass_init {
-    use crate::impl_::pyclass::{self, PyClassBaseType};
-    use crate::impl_::pyclass_init::{PyClassInit, PyNativeTypeInitializer};
-    use crate::{PyClass, PyClassInitializer, PyErr};
-
-    pub trait Sealed<'py, const IS_PYCLASS: bool, const IS_INITIALIZER_TUPLE: bool> {}
-
-    impl<'py, T> Sealed<'py, false, false> for T where T: crate::IntoPyObject<'py> {}
-    impl<'py, T> Sealed<'py, true, false> for T
-    where
-        T: crate::PyClass,
-        T::BaseType: pyclass::PyClassBaseType<Initializer = PyNativeTypeInitializer<T::BaseType>>,
-    {
-    }
-    impl<'py, T, E, const IS_PYCLASS: bool, const IS_INITIALIZER_TUPLE: bool>
-        Sealed<'py, IS_PYCLASS, IS_INITIALIZER_TUPLE> for Result<T, E>
-    where
-        T: PyClassInit<'py, IS_PYCLASS, IS_INITIALIZER_TUPLE>,
-        E: Into<PyErr>,
-    {
-    }
-    impl<'py, T> Sealed<'py, false, false> for PyClassInitializer<T> where T: PyClass {}
-    impl<'py, S, B> Sealed<'py, false, true> for (S, B)
-    where
-        S: PyClass<BaseType = B>,
-        B: PyClass + PyClassBaseType<Initializer = PyClassInitializer<B>>,
-        B::BaseType: PyClassBaseType<Initializer = PyNativeTypeInitializer<B::BaseType>>,
-    {
-    }
-}
-
-impl<'py, T> PyClassInit<'py, false, false> for T
-where
-    T: crate::IntoPyObject<'py>,
-{
-    fn init(
-        self,
-        cls: crate::Borrowed<'_, 'py, crate::types::PyType>,
-    ) -> PyResult<crate::Bound<'py, crate::PyAny>> {
-        self.into_pyobject(cls.py())
-            .map(crate::BoundObject::into_any)
-            .map(crate::BoundObject::into_bound)
-            .map_err(Into::into)
-    }
-}
-
-impl<'py, T> PyClassInit<'py, true, false> for T
-where
-    T: crate::PyClass,
-    T::BaseType:
-        super::pyclass::PyClassBaseType<Initializer = PyNativeTypeInitializer<T::BaseType>>,
-{
-    fn init(
-        self,
-        cls: crate::Borrowed<'_, 'py, crate::types::PyType>,
-    ) -> PyResult<crate::Bound<'py, crate::PyAny>> {
-        PyClassInitializer::from(self).init(cls)
-    }
-}
-
-impl<'py, T, E, const IS_PYCLASS: bool, const IS_INITIALIZER_TUPLE: bool>
-    PyClassInit<'py, IS_PYCLASS, IS_INITIALIZER_TUPLE> for Result<T, E>
-where
-    T: PyClassInit<'py, IS_PYCLASS, IS_INITIALIZER_TUPLE>,
-    E: Into<PyErr>,
-{
-    fn init(
-        self,
-        cls: crate::Borrowed<'_, 'py, crate::types::PyType>,
-    ) -> PyResult<crate::Bound<'py, crate::PyAny>> {
-        self.map_err(Into::into)?.init(cls)
-    }
-}
-
-impl<'py, T> PyClassInit<'py, false, false> for PyClassInitializer<T>
-where
-    T: PyClass,
-{
-    fn init(
-        self,
-        cls: crate::Borrowed<'_, 'py, crate::types::PyType>,
-    ) -> PyResult<crate::Bound<'py, crate::PyAny>> {
-        unsafe {
-            self.create_class_object_of_type(cls.py(), cls.as_ptr().cast())
-                .map(crate::Bound::into_any)
-        }
-    }
-}
-
-impl<'py, S, B> PyClassInit<'py, false, true> for (S, B)
-where
-    S: PyClass<BaseType = B>,
-    B: PyClass + PyClassBaseType<Initializer = PyClassInitializer<B>>,
-    B::BaseType: PyClassBaseType<Initializer = PyNativeTypeInitializer<B::BaseType>>,
-{
-    fn init(
-        self,
-        cls: crate::Borrowed<'_, 'py, crate::types::PyType>,
-    ) -> PyResult<crate::Bound<'py, crate::PyAny>> {
-        let (sub, base) = self;
-        PyClassInitializer::from(base).add_subclass(sub).init(cls)
-    }
-}
+/// Re-exported so that macros can name this internal type.
+pub use crate::internal::pyclass_init::PyNativeTypeInitializer;

src/impl_/pymethods.rs

@@ -3,10 +3,13 @@ use crate::impl_::callback::IntoPyCallbackOutput;
 use crate::impl_::panic::PanicTrap;
 use crate::impl_::pycell::PyClassObjectBaseLayout;
 use crate::internal::get_slot::{get_slot, TP_BASE, TP_CLEAR, TP_TRAVERSE};
+use crate::internal::pyclass_init::PyClassInit;
 use crate::internal::state::ForbidAttaching;
 use crate::pycell::impl_::{PyClassBorrowChecker as _, PyClassObjectLayout};
 use crate::types::PyType;
-use crate::{ffi, Bound, Py, PyAny, PyClass, PyErr, PyResult, PyTraverseError, PyVisit, Python};
+use crate::{
+    ffi, Borrowed, Bound, Py, PyAny, PyClass, PyErr, PyResult, PyTraverseError, PyVisit, Python,
+};
 use core::ffi::CStr;
 use core::ffi::{c_int, c_void};
 use core::fmt;
@@ -688,18 +691,26 @@ pub trait AsyncIterResultOptionKind {
 
 impl<Value, Error> AsyncIterResultOptionKind for Result<Option<Value>, Error> {}
 
-pub unsafe fn tp_new_impl<'py, T, const IS_PYCLASS: bool, const IS_INITIALIZER_TUPLE: bool>(
+/// Re-exported so that `#[new]` generated code can resolve the type tag for `tp_new_impl`
+pub use crate::internal::pyclass_init::tp_new_resolver;
+
+#[expect(
+    private_bounds,
+    reason = "`PyClassInit` is not a public trait, bound exist for diagnostics"
+)]
+/// # SAFETY
+/// - `cls` must be the type object for `ClassT` (or a subclass)
+pub unsafe fn tp_new_impl<'py, InitializerT, ClassT>(
     py: Python<'py>,
-    obj: T,
+    initializer: InitializerT,
     cls: *mut ffi::PyTypeObject,
 ) -> PyResult<*mut ffi::PyObject>
 where
-    T: super::pyclass_init::PyClassInit<'py, IS_PYCLASS, IS_INITIALIZER_TUPLE>,
+    InitializerT: PyClassInit<'py, ClassT>,
 {
-    unsafe {
-        obj.init(crate::Borrowed::from_ptr_unchecked(py, cls.cast()).cast_unchecked())
-            .map(Bound::into_ptr)
-    }
+    // SAFETY: caller has guaranteed `cls` is the correct object
+    unsafe { initializer.init(Borrowed::from_ptr_unchecked(py, cls.cast()).cast_unchecked()) }
+        .map(Bound::into_ptr)
 }
 
 #[cfg(test)]

src/impl_/wrap.rs

@@ -23,6 +23,34 @@ impl<T> SomeWrap<T> for Option<T> {
     }
 }
 
+pub struct OkWrapper<T>(OkWrapperInner<T>);
+pub struct OkWrapperInner<T>(PhantomData<T>);
+
+impl<T> OkWrapper<T> {
+    pub fn new(_: &T) -> Self {
+        Self(OkWrapperInner(PhantomData))
+    }
+}
+
+impl<T> Deref for OkWrapper<T> {
+    type Target = OkWrapperInner<T>;
+    fn deref(&self) -> &Self::Target {
+        &self.0
+    }
+}
+
+impl<T, E> OkWrapper<Result<T, E>> {
+    pub fn ok_wrap(&self, value: Result<T, E>) -> Result<T, E> {
+        value
+    }
+}
+
+impl<T> OkWrapperInner<T> {
+    pub fn ok_wrap(&self, value: T) -> Result<T, Infallible> {
+        Ok(value)
+    }
+}
+
 // Hierarchy of conversions used in the function return type machinery
 pub struct Converter<T>(EmptyTupleConverter<T>);
 pub struct EmptyTupleConverter<T>(IntoPyObjectConverter<T>);
@@ -155,4 +183,15 @@ mod tests {
         let b: Option<u8> = SomeWrap::wrap(None);
         assert_eq!(b, None);
     }
+
+    #[test]
+    fn wrap_result() {
+        let a = 42;
+        let Ok(a) = OkWrapper::new(&a).ok_wrap(a);
+        assert_eq!(a, 42);
+
+        let b = Result::<_, String>::Ok(42);
+        let b = OkWrapper::new(&b).ok_wrap(b);
+        assert_eq!(b, Ok(42));
+    }
 }

src/internal.rs

@@ -1,4 +1,5 @@
 //! Holding place for code which is not intended to be reachable from outside of PyO3.
 
 pub(crate) mod get_slot;
+pub(crate) mod pyclass_init;
 pub(crate) mod state;