FIx method resolution in number slots when a single method is a subclass

newsfragments/6228.fixed.md

@@ -0,0 +1 @@
+Fix reversible arithmetic operators don't respect subclassing if only one arm implemented Closes #6024

pyo3-macros-backend/src/pyimpl.rs

@@ -319,6 +319,16 @@ fn add_shared_proto_slots(
         }};
     }
 
+    macro_rules! try_add_binop_slot {
+        ($slot:ident, $forward:literal, $reflected:literal) => {{
+            let has_forward = implemented_proto_fragments.remove($forward);
+            let has_reflected = implemented_proto_fragments.remove($reflected);
+            if has_forward || has_reflected {
+                proto_impls.push(quote! { #pyo3_path::impl_::pyclass::$slot!(#ty, #has_forward, #has_reflected) });
+            }
+        }};
+    }
+
     try_add_shared_slot!(
         generate_pyclass_getattro_slot,
         "__getattribute__",
@@ -327,24 +337,24 @@ fn add_shared_proto_slots(
     try_add_shared_slot!(generate_pyclass_setattr_slot, "__setattr__", "__delattr__");
     try_add_shared_slot!(generate_pyclass_setdescr_slot, "__set__", "__delete__");
     try_add_shared_slot!(generate_pyclass_setitem_slot, "__setitem__", "__delitem__");
-    try_add_shared_slot!(generate_pyclass_add_slot, "__add__", "__radd__");
-    try_add_shared_slot!(generate_pyclass_sub_slot, "__sub__", "__rsub__");
-    try_add_shared_slot!(generate_pyclass_mul_slot, "__mul__", "__rmul__");
-    try_add_shared_slot!(generate_pyclass_mod_slot, "__mod__", "__rmod__");
-    try_add_shared_slot!(generate_pyclass_divmod_slot, "__divmod__", "__rdivmod__");
-    try_add_shared_slot!(generate_pyclass_lshift_slot, "__lshift__", "__rlshift__");
-    try_add_shared_slot!(generate_pyclass_rshift_slot, "__rshift__", "__rrshift__");
-    try_add_shared_slot!(generate_pyclass_and_slot, "__and__", "__rand__");
-    try_add_shared_slot!(generate_pyclass_or_slot, "__or__", "__ror__");
-    try_add_shared_slot!(generate_pyclass_xor_slot, "__xor__", "__rxor__");
-    try_add_shared_slot!(generate_pyclass_matmul_slot, "__matmul__", "__rmatmul__");
-    try_add_shared_slot!(generate_pyclass_truediv_slot, "__truediv__", "__rtruediv__");
-    try_add_shared_slot!(
+    try_add_binop_slot!(generate_pyclass_add_slot, "__add__", "__radd__");
+    try_add_binop_slot!(generate_pyclass_sub_slot, "__sub__", "__rsub__");
+    try_add_binop_slot!(generate_pyclass_mul_slot, "__mul__", "__rmul__");
+    try_add_binop_slot!(generate_pyclass_mod_slot, "__mod__", "__rmod__");
+    try_add_binop_slot!(generate_pyclass_divmod_slot, "__divmod__", "__rdivmod__");
+    try_add_binop_slot!(generate_pyclass_lshift_slot, "__lshift__", "__rlshift__");
+    try_add_binop_slot!(generate_pyclass_rshift_slot, "__rshift__", "__rrshift__");
+    try_add_binop_slot!(generate_pyclass_and_slot, "__and__", "__rand__");
+    try_add_binop_slot!(generate_pyclass_or_slot, "__or__", "__ror__");
+    try_add_binop_slot!(generate_pyclass_xor_slot, "__xor__", "__rxor__");
+    try_add_binop_slot!(generate_pyclass_matmul_slot, "__matmul__", "__rmatmul__");
+    try_add_binop_slot!(generate_pyclass_truediv_slot, "__truediv__", "__rtruediv__");
+    try_add_binop_slot!(
         generate_pyclass_floordiv_slot,
         "__floordiv__",
         "__rfloordiv__"
     );
-    try_add_shared_slot!(generate_pyclass_pow_slot, "__pow__", "__rpow__");
+    try_add_binop_slot!(generate_pyclass_pow_slot, "__pow__", "__rpow__");
     try_add_shared_slot!(
         generate_pyclass_richcompare_slot,
         "__lt__",

pyo3-macros-backend/src/pymethod.rs

@@ -1655,20 +1655,18 @@ impl SlotFragmentDef {
 
     const fn binary_operator(fragment: &'static str) -> Self {
         // Binary operator fragments (`__add__`, `__radd__`, etc.) are combined
-        // into a shared slot (e.g. `nb_add`) that may call the forward fragment
-        // with a non-class receiver (e.g. `1 + MyClass()` → `nb_add(1, c)`).
-        // The runtime helper then tries the reflected fragment with the operands
-        // swapped, which can also produce a non-class `_slf`.  Both cases require
-        // a checked type conversion so that a mismatch gracefully returns
-        // `NotImplemented` rather than causing undefined behaviour.
-        // TODO: addressing #6024 could allow to simplify this by using type
-        // checks to directly dispatch to the right fragment.
+        // into a shared slot (e.g. `nb_add`) that dispatch to the appropriate
+        // fragment based on the receiver type.
+        // Direct calls to the fragment through method wrappers are guaranteed
+        // to have a receiver of the correct type.
         SlotFragmentDef {
             fragment,
             arguments: &[Ty::Object],
             extract_error_mode: ExtractErrorMode::NotImplemented,
             ret_ty: Ty::Object,
-            self_conversion: SelfConversionPolicy::checked(),
+            // SAFETY: Combined slot dispatch based on type which guarantees
+            // the receiver is of the correct type for direct calls to the fragment.
+            self_conversion: unsafe { SelfConversionPolicy::trusted() },
         }
     }
 

src/impl_/pyclass.rs

@@ -23,6 +23,7 @@ use core::{
     ffi::{c_int, c_void},
     marker::PhantomData,
     ptr::{self, NonNull},
+    result::Result,
 };
 use std::{sync::Mutex, thread};
 
@@ -517,7 +518,6 @@ define_pyclass_setattr_slot! {
     objobjargproc,
 }
 
-/// Macro which expands to three items
 /// - Trait for a lhs dunder e.g. __add__
 /// - Trait for the corresponding rhs e.g. __radd__
 /// - A macro which will use dtolnay specialisation to generate the shared slot for the two dunders
@@ -563,21 +563,126 @@ macro_rules! define_pyclass_binary_operator_slot {
         #[doc(hidden)]
         #[macro_export]
         macro_rules! $generate_macro {
-            ($cls:ty) => {{
+            // Both forward and reflected available
+            ($cls:ty, true, true) => {{
                 unsafe fn slot_impl(
                     py: $crate::Python<'_>,
                     _slf: *mut $crate::ffi::PyObject,
                     _other: *mut $crate::ffi::PyObject,
                 ) -> $crate::PyResult<*mut $crate::ffi::PyObject> {
                     use $crate::impl_::pyclass::*;
                     let collector = PyClassImplCollector::<$cls>::new();
-                    let lhs_result = unsafe { collector.$lhs(py, _slf, _other) }?;
-                    if lhs_result == unsafe { $crate::ffi::Py_NotImplemented() } {
-                        unsafe { $crate::ffi::Py_DECREF(lhs_result) };
-                        unsafe { collector.$rhs(py, _other, _slf) }
-                    } else {
-                        ::core::result::Result::Ok(lhs_result)
+
+                    let lhs_obj = unsafe { $crate::Bound::from_borrowed_ptr(py, _slf) };
+                    if lhs_obj.is_instance_of::<$cls>() {
+                        return unsafe { collector.$lhs(py, _slf, _other) };
+                    }
+
+                    let rhs_obj = unsafe { $crate::Bound::from_borrowed_ptr(py, _other) };
+                    if rhs_obj.is_instance_of::<$cls>() {
+                        return unsafe { collector.$rhs(py, _other, _slf) };
+                    }
+
+                    Result::Ok(py.NotImplemented().into_ptr())
+                }
+
+                $crate::ffi::PyType_Slot {
+                    slot: $crate::ffi::$slot,
+                    pfunc: $crate::impl_::trampoline::get_trampoline_function!(
+                        binaryfunc, slot_impl
+                    ) as $crate::ffi::binaryfunc as _,
+                }
+            }};
+
+            // Only forward available
+            ($cls:ty, true, false) => {{
+                unsafe fn slot_impl(
+                    py: $crate::Python<'_>,
+                    _slf: *mut $crate::ffi::PyObject,
+                    _other: *mut $crate::ffi::PyObject,
+                ) -> $crate::PyResult<*mut $crate::ffi::PyObject> {
+                    use $crate::impl_::pyclass::*;
+                    let lhs_obj = unsafe { $crate::Bound::from_borrowed_ptr(py, _slf) };
+
+                    if lhs_obj.is_instance_of::<$cls>() {
+                        let collector = PyClassImplCollector::<$cls>::new();
+                        return unsafe { collector.$lhs(py, _slf, _other) };
+                    }
+
+                    let rhs_obj = unsafe { $crate::Bound::from_borrowed_ptr(py, _other) };
+                    if rhs_obj.is_instance_of::<$cls>() {
+                        // RHS is our class, try parent's reflected method via super()
+                        // This ensures parent knows it's being called for reflection
+                        let py_type = <$cls as $crate::PyTypeInfo>::type_object(py);
+                        let reflected_name = $crate::intern!(py, concat!(stringify!($rhs)));
+                        let super_obj = match $crate::types::PySuper::new(&py_type, &rhs_obj) {
+                            Ok(s) => s,
+                            Err(e) => return Err(e),
+                        };
+                        match super_obj.call_method1(reflected_name, (&lhs_obj,)) {
+                            Ok(result) => return Ok(result.into_ptr()),
+                            Err(e) => {
+                                // If parent doesn't implement the method, return NotImplemented instead of AttributeError
+                                if e.is_instance_of::<$crate::exceptions::PyAttributeError>(py) {
+                                    return Ok(py.NotImplemented().into_ptr());
+                                } else {
+                                    return Err(e);
+                                }
+                            }
+                        }
+                    }
+                    Result::Ok(py.NotImplemented().into_ptr())
+                }
+
+                $crate::ffi::PyType_Slot {
+                    slot: $crate::ffi::$slot,
+                    pfunc: $crate::impl_::trampoline::get_trampoline_function!(
+                        binaryfunc, slot_impl
+                    ) as $crate::ffi::binaryfunc as _,
+                }
+            }};
+
+            // Only reflected available
+            ($cls:ty, false, true) => {{
+                unsafe fn slot_impl(
+                    py: $crate::Python<'_>,
+                    _slf: *mut $crate::ffi::PyObject,
+                    _other: *mut $crate::ffi::PyObject,
+                ) -> $crate::PyResult<*mut $crate::ffi::PyObject> {
+                    use $crate::impl_::pyclass::*;
+
+                    let lhs_obj = unsafe { $crate::Bound::from_borrowed_ptr(py, _slf) };
+                    if lhs_obj.is_instance_of::<$cls>() {
+                        // LHS is our class, try parent's forward method via super()
+                        // This ensures parent knows it's being called for forward dispatch
+                        let py_type = <$cls as $crate::PyTypeInfo>::type_object(py);
+                        let forward_name = $crate::intern!(py, concat!(stringify!($lhs)));
+                        let other_bound = unsafe { $crate::Bound::from_borrowed_ptr(py, _other) };
+                        let super_obj = match $crate::types::PySuper::new(&py_type, &lhs_obj) {
+                            Ok(s) => s,
+                            Err(e) => return Err(e),
+                        };
+                        match super_obj.call_method1(forward_name, (&other_bound,)) {
+                            Ok(result) => return Ok(result.into_ptr()),
+                            Err(e) => {
+                                // If parent doesn't implement the method, return NotImplemented instead of AttributeError
+                                if e.is_instance_of::<$crate::exceptions::PyAttributeError>(py) {
+                                    return Ok(py.NotImplemented().into_ptr());
+                                } else {
+                                    return Err(e);
+                                }
+                            }
+                        }
                     }
+
+                    let rhs_obj = unsafe { $crate::Bound::from_borrowed_ptr(py, _other) };
+                    let rhs_is_instance = rhs_obj.is_instance_of::<$cls>();
+                    if rhs_is_instance {
+                        let collector = PyClassImplCollector::<$cls>::new();
+                        return unsafe { collector.$rhs(py, _other, _slf) };
+                    }
+
+                    Result::Ok(py.NotImplemented().into_ptr())
                 }
 
                 $crate::ffi::PyType_Slot {
@@ -744,22 +849,135 @@ slot_fragment_trait! {
 #[doc(hidden)]
 #[macro_export]
 macro_rules! generate_pyclass_pow_slot {
-    ($cls:ty) => {{
-        fn slot_impl(
+    // Both forward and reflected available
+    ($cls:ty, true, true) => {{
+        unsafe fn slot_impl(
             py: $crate::Python<'_>,
             _slf: *mut $crate::ffi::PyObject,
             _other: *mut $crate::ffi::PyObject,
             _mod: *mut $crate::ffi::PyObject,
         ) -> $crate::PyResult<*mut $crate::ffi::PyObject> {
             use $crate::impl_::pyclass::*;
             let collector = PyClassImplCollector::<$cls>::new();
-            let lhs_result = unsafe { collector.__pow__(py, _slf, _other, _mod) }?;
-            if lhs_result == unsafe { $crate::ffi::Py_NotImplemented() } {
-                unsafe { $crate::ffi::Py_DECREF(lhs_result) };
-                unsafe { collector.__rpow__(py, _other, _slf, _mod) }
-            } else {
-                ::core::result::Result::Ok(lhs_result)
+
+            let lhs_obj = unsafe { $crate::Bound::from_borrowed_ptr(py, _slf) };
+            if lhs_obj.is_instance_of::<$cls>() {
+                return unsafe { collector.__pow__(py, _slf, _other, _mod) };
             }
+
+            let rhs_obj = unsafe { $crate::Bound::from_borrowed_ptr(py, _other) };
+            if rhs_obj.is_instance_of::<$cls>() {
+                return unsafe { collector.__rpow__(py, _other, _slf, _mod) };
+            }
+
+            Result::Ok(py.NotImplemented().into_ptr())
+        }
+
+        $crate::ffi::PyType_Slot {
+            slot: $crate::ffi::Py_nb_power,
+            pfunc: $crate::impl_::trampoline::get_trampoline_function!(ternaryfunc, slot_impl)
+                as $crate::ffi::ternaryfunc as _,
+        }
+    }};
+
+    // Only forward available
+    ($cls:ty, true, false) => {{
+        unsafe fn slot_impl(
+            py: $crate::Python<'_>,
+            _slf: *mut $crate::ffi::PyObject,
+            _other: *mut $crate::ffi::PyObject,
+            _mod: *mut $crate::ffi::PyObject,
+        ) -> $crate::PyResult<*mut $crate::ffi::PyObject> {
+            use $crate::impl_::pyclass::*;
+            let lhs_obj = unsafe { $crate::Bound::from_borrowed_ptr(py, _slf) };
+
+            if lhs_obj.is_instance_of::<$cls>() {
+                let collector = PyClassImplCollector::<$cls>::new();
+                return unsafe { collector.__pow__(py, _slf, _other, _mod) };
+            }
+
+            let rhs_obj = unsafe { $crate::Bound::from_borrowed_ptr(py, _other) };
+            if rhs_obj.is_instance_of::<$cls>() {
+                // RHS is our class, try parent's reflected method via super()
+                // This ensures parent knows it's being called for reflection
+                let py_type = <$cls as $crate::PyTypeInfo>::type_object(py);
+                let mod_obj = if _mod.is_null() {
+                    py.None().into_bound(py)
+                } else {
+                    unsafe { $crate::Bound::from_borrowed_ptr(py, _mod) }
+                };
+                let super_obj = match $crate::types::PySuper::new(&py_type, &rhs_obj) {
+                    Ok(s) => s,
+                    Err(e) => return Err(e),
+                };
+                match super_obj.call_method1($crate::intern!(py, "__rpow__"), (&lhs_obj, mod_obj)) {
+                    Ok(result) => return Ok(result.into_ptr()),
+                    Err(e) => {
+                        // If parent doesn't implement the method, return NotImplemented instead of AttributeError
+                        if e.is_instance_of::<$crate::exceptions::PyAttributeError>(py) {
+                            return Ok(py.NotImplemented().into_ptr());
+                        } else {
+                            return Err(e);
+                        }
+                    }
+                }
+            }
+            Result::Ok(py.NotImplemented().into_ptr())
+        }
+
+        $crate::ffi::PyType_Slot {
+            slot: $crate::ffi::Py_nb_power,
+            pfunc: $crate::impl_::trampoline::get_trampoline_function!(ternaryfunc, slot_impl)
+                as $crate::ffi::ternaryfunc as _,
+        }
+    }};
+
+    // Only reflected available
+    ($cls:ty, false, true) => {{
+        unsafe fn slot_impl(
+            py: $crate::Python<'_>,
+            _slf: *mut $crate::ffi::PyObject,
+            _other: *mut $crate::ffi::PyObject,
+            _mod: *mut $crate::ffi::PyObject,
+        ) -> $crate::PyResult<*mut $crate::ffi::PyObject> {
+            use $crate::impl_::pyclass::*;
+
+            let lhs_obj = unsafe { $crate::Bound::from_borrowed_ptr(py, _slf) };
+            if lhs_obj.is_instance_of::<$cls>() {
+                // LHS is our class, try parent's forward method via super()
+                // This ensures parent knows it's being called for forward dispatch
+                let py_type = <$cls as $crate::PyTypeInfo>::type_object(py);
+                let rhs_obj = unsafe { $crate::Bound::from_borrowed_ptr(py, _other) };
+                let mod_obj = if _mod.is_null() {
+                    py.None().into_bound(py)
+                } else {
+                    unsafe { $crate::Bound::from_borrowed_ptr(py, _mod) }
+                };
+                let super_obj = match $crate::types::PySuper::new(&py_type, &lhs_obj) {
+                    Ok(s) => s,
+                    Err(e) => return Err(e),
+                };
+                match super_obj.call_method1($crate::intern!(py, "__pow__"), (&rhs_obj, mod_obj)) {
+                    Ok(result) => return Ok(result.into_ptr()),
+                    Err(e) => {
+                        // If parent doesn't implement the method, return NotImplemented instead of AttributeError
+                        if e.is_instance_of::<$crate::exceptions::PyAttributeError>(py) {
+                            return Ok(py.NotImplemented().into_ptr());
+                        } else {
+                            return Err(e);
+                        }
+                    }
+                }
+            }
+
+            let rhs_obj = unsafe { $crate::Bound::from_borrowed_ptr(py, _other) };
+            let rhs_is_instance = rhs_obj.is_instance_of::<$cls>();
+            if rhs_is_instance {
+                let collector = PyClassImplCollector::<$cls>::new();
+                return unsafe { collector.__rpow__(py, _other, _slf, _mod) };
+            }
+
+            Result::Ok(py.NotImplemented().into_ptr())
         }
 
         $crate::ffi::PyType_Slot {