Added as_super methods to PyRef and PyRefMut.

newsfragments/4219.added.md

@@ -0,0 +1 @@
+Added `as_super` methods to `PyRef` and `PyRefMut`.

src/pycell.rs

@@ -611,6 +611,7 @@ impl<T: PyClass + fmt::Debug> fmt::Debug for PyCell<T> {
 /// ```
 ///
 /// See the [module-level documentation](self) for more information.
+#[repr(transparent)]
 pub struct PyRef<'p, T: PyClass> {
     // TODO: once the GIL Ref API is removed, consider adding a lifetime parameter to `PyRef` to
     // store `Borrowed` here instead, avoiding reference counting overhead.
@@ -630,7 +631,7 @@ where
     U: PyClass,
 {
     fn as_ref(&self) -> &T::BaseType {
-        unsafe { &*self.inner.get_class_object().ob_base.get_ptr() }
+        self.as_super()
     }
 }
 
@@ -742,6 +743,57 @@ where
             },
         }
     }
+    /// Borrows a shared reference to `PyRef<T::BaseType>`.
+    ///
+    /// With the help of this method, you can access attributes and call methods
+    /// on the superclass without consuming the `PyRef<T>`. This method can also
+    /// be chained to access the super-superclass (and so on).
+    ///
+    /// # Examples
+    /// ```
+    /// # use pyo3::prelude::*;
+    /// #[pyclass(subclass)]
+    /// struct Base {
+    ///     base_name: &'static str,
+    /// }
+    /// #[pymethods]
+    /// impl Base {
+    ///     fn base_name_len(&self) -> usize {
+    ///         self.base_name.len()
+    ///     }
+    /// }
+    ///
+    /// #[pyclass(extends=Base)]
+    /// struct Sub {
+    ///     sub_name: &'static str,
+    /// }
+    ///
+    /// #[pymethods]
+    /// impl Sub {
+    ///     #[new]
+    ///     fn new() -> (Self, Base) {
+    ///         (Self { sub_name: "sub_name" }, Base { base_name: "base_name" })
+    ///     }
+    ///     fn sub_name_len(&self) -> usize {
+    ///         self.sub_name.len()
+    ///     }
+    ///     fn format_name_lengths(slf: PyRef<'_, Self>) -> String {
+    ///         format!("{} {}", slf.as_super().base_name_len(), slf.sub_name_len())
+    ///     }
+    /// }
+    /// # Python::with_gil(|py| {
+    /// #     let sub = Py::new(py, Sub::new()).unwrap();
+    /// #     pyo3::py_run!(py, sub, "assert sub.format_name_lengths() == '9 8'")
+    /// # });
+    /// ```
+    pub fn as_super(&self) -> &PyRef<'p, U> {
+        unsafe {
+            // trait bound guarantees compatibility
+            let inner: &Bound<'p, U> = self.inner.as_any().downcast_unchecked();
+            // `repr(transparent)` guarantees `Bound<U>` has the same layout as `PyRef<U>`
+            &*(inner as *const Bound<'p, U> as *const PyRef<'p, U>)
+        }
+    }
 }
 
 impl<'p, T: PyClass> Deref for PyRef<'p, T> {
@@ -798,6 +850,7 @@ impl<T: PyClass + fmt::Debug> fmt::Debug for PyRef<'_, T> {
 /// A wrapper type for a mutably borrowed value from a [`Bound<'py, T>`].
 ///
 /// See the [module-level documentation](self) for more information.
+#[repr(transparent)]
 pub struct PyRefMut<'p, T: PyClass<Frozen = False>> {
     // TODO: once the GIL Ref API is removed, consider adding a lifetime parameter to `PyRef` to
     // store `Borrowed` here instead, avoiding reference counting overhead.
@@ -817,7 +870,7 @@ where
     U: PyClass<Frozen = False>,
 {
     fn as_ref(&self) -> &T::BaseType {
-        unsafe { &*self.inner.get_class_object().ob_base.get_ptr() }
+        PyRefMut::downgrade(self).as_super()
     }
 }
 
@@ -827,7 +880,7 @@ where
     U: PyClass<Frozen = False>,
 {
     fn as_mut(&mut self) -> &mut T::BaseType {
-        unsafe { &mut *self.inner.get_class_object().ob_base.get_ptr() }
+        self.as_super()
     }
 }
 
@@ -869,6 +922,12 @@ impl<'py, T: PyClass<Frozen = False>> PyRefMut<'py, T> {
             .try_borrow_mut()
             .map(|_| Self { inner: obj.clone() })
     }
+
+    pub(crate) fn downgrade(slf: &Self) -> &PyRef<'py, T> {
+        // `PyRef<T>` and `PyRefMut<T>` have the same layout, and they are behind
+        // a reference so the difference in `Drop` behavior is irrelevant.
+        unsafe { &*(slf as *const Self as *const PyRef<'py, T>) }
+    }
 }
 
 impl<'p, T, U> PyRefMut<'p, T>
@@ -890,6 +949,22 @@ where
             },
         }
     }
+    /// Borrows a mutable reference to `PyRefMut<T::BaseType>`.
+    ///
+    /// With the help of this method, you can mutate attributes and call mutating
+    /// methods on the superclass without consuming the `PyRefMet<T>`. This method
+    /// can also be chained to access the super-superclass (and so on).
+    ///
+    /// See [`PyRef::as_super`] for more.
+    pub fn as_super(&mut self) -> &mut PyRefMut<'p, U> {
+        unsafe {
+            // trait bound guarantees compatibility
+            let inner = &mut self.inner as *mut Bound<'p, T> as *mut Bound<'p, U>;
+            // `repr(transparent)` guarantees `Bound<U>` has the same layout as
+            // `PyRefMut<U>`, and the mutable borrow on `self` prevents aliasing
+            &mut *(inner as *mut PyRefMut<'p, U>)
+        }
+    }
 }
 
 impl<'p, T: PyClass<Frozen = False>> Deref for PyRefMut<'p, T> {
@@ -1139,4 +1214,67 @@ mod tests {
             unsafe { ffi::Py_DECREF(ptr) };
         })
     }
+
+    #[crate::pyclass]
+    #[pyo3(crate = "crate", subclass)]
+    struct Base {
+        base_name: &'static str,
+    }
+
+    #[crate::pyclass]
+    #[pyo3(crate = "crate", extends=Base, subclass)]
+    struct Sub {
+        sub_name: &'static str,
+    }
+
+    #[crate::pyclass]
+    #[pyo3(crate = "crate", extends=Sub)]
+    struct SubSub {
+        subsub_name: &'static str,
+    }
+
+    #[crate::pymethods]
+    #[pyo3(crate = "crate")]
+    impl SubSub {
+        #[new]
+        #[rustfmt::skip]
+        fn new(py: Python<'_>) -> crate::PyResult<crate::Py<SubSub>> {
+            let base = Base { base_name: "base_name" };
+            let sub = Sub { sub_name: "sub_name" };
+            let subsub = SubSub { subsub_name: "subsub_name"};
+            let init = crate::PyClassInitializer::from(base)
+                .add_subclass(sub).add_subclass(subsub);
+            crate::Py::new(py, init)
+        }
+    }
+
+    #[test]
+    fn test_pyref_as_super() {
+        Python::with_gil(|py| {
+            let subsub = SubSub::new(py).unwrap().into_bound(py);
+            let pyref = subsub.borrow();
+            assert_eq!(pyref.as_super().as_super().base_name, "base_name");
+            assert_eq!(pyref.as_super().sub_name, "sub_name");
+            assert_eq!(pyref.subsub_name, "subsub_name");
+            // `as_ref` still works the same
+            assert_eq!(pyref.as_ref().sub_name, "sub_name");
+        });
+    }
+
+    #[test]
+    fn test_pyrefmut_as_super() {
+        Python::with_gil(|py| {
+            let subsub = SubSub::new(py).unwrap().into_bound(py);
+            let mut pyrefmut = subsub.borrow_mut();
+            pyrefmut.as_super().as_super().base_name = "base_name2";
+            pyrefmut.as_super().sub_name = "sub_name2";
+            pyrefmut.subsub_name = "subsub_name2";
+            assert_eq!(pyrefmut.as_super().as_super().base_name, "base_name2");
+            assert_eq!(pyrefmut.as_super().sub_name, "sub_name2");
+            assert_eq!(pyrefmut.subsub_name, "subsub_name2");
+            // `as_ref`/`as_mut` still work the same
+            pyrefmut.as_mut().sub_name = "sub_name3";
+            assert_eq!(pyrefmut.as_ref().sub_name, "sub_name3");
+        });
+    }
 }