Introspection: Basic pyclass(extends) support

Introduces PyClassBaseType::BASE_NAME to store the module name + class name of the base type

Author: Tpt

newsfragments/5331.added.md

@@ -0,0 +1 @@
+Introspection: emit base classes.

pyo3-introspection/src/introspection.rs

@@ -120,9 +120,13 @@ fn convert_members<'a>(
                     chunks_by_parent,
                 )?);
             }
-            Chunk::Class { name, id } => {
-                classes.push(convert_class(id, name, chunks_by_id, chunks_by_parent)?)
-            }
+            Chunk::Class { name, id, bases } => classes.push(convert_class(
+                id,
+                name,
+                bases,
+                chunks_by_id,
+                chunks_by_parent,
+            )?),
             Chunk::Function {
                 name,
                 id: _,
@@ -164,6 +168,7 @@ fn convert_members<'a>(
 fn convert_class(
     id: &str,
     name: &str,
+    bases: &[String],
     chunks_by_id: &HashMap<&str, &Chunk>,
     chunks_by_parent: &HashMap<&str, Vec<&Chunk>>,
 ) -> Result<Class> {
@@ -182,6 +187,7 @@ fn convert_class(
     );
     Ok(Class {
         name: name.into(),
+        bases: bases.into(),
         methods,
         attributes,
     })
@@ -404,6 +410,8 @@ enum Chunk {
     Class {
         id: String,
         name: String,
+        #[serde(default)]
+        bases: Vec<String>,
     },
     Function {
         #[serde(default)]

pyo3-introspection/src/model.rs

@@ -11,6 +11,7 @@ pub struct Module {
 #[derive(Debug, Eq, PartialEq, Clone, Hash)]
 pub struct Class {
     pub name: String,
+    pub bases: Vec<String>,
     pub methods: Vec<Function>,
     pub attributes: Vec<Attribute>,
 }

pyo3-introspection/src/stubs.rs

@@ -100,7 +100,18 @@ fn module_stubs(module: &Module) -> String {
 }
 
 fn class_stubs(class: &Class, modules_to_import: &mut BTreeSet<String>) -> String {
-    let mut buffer = format!("class {}:", class.name);
+    let mut buffer = format!("class {}", class.name);
+    if !class.bases.is_empty() {
+        buffer.push('(');
+        for (i, base) in class.bases.iter().enumerate() {
+            if i > 0 {
+                buffer.push_str(", ");
+            }
+            buffer.push_str(annotation_stub(base, modules_to_import));
+        }
+        buffer.push(')');
+    }
+    buffer.push(':');
     if class.methods.is_empty() && class.attributes.is_empty() {
         buffer.push_str(" ...");
         return buffer;

pyo3-macros-backend/src/introspection.rs

@@ -20,7 +20,7 @@ use std::hash::{Hash, Hasher};
 use std::mem::take;
 use std::sync::atomic::{AtomicUsize, Ordering};
 use syn::visit_mut::{visit_type_mut, VisitMut};
-use syn::{Attribute, Ident, ReturnType, Type, TypePath};
+use syn::{Attribute, Ident, Path, ReturnType, Type, TypePath};
 
 static GLOBAL_COUNTER_FOR_UNIQUE_NAMES: AtomicUsize = AtomicUsize::new(0);
 
@@ -61,19 +61,23 @@ pub fn class_introspection_code(
     pyo3_crate_path: &PyO3CratePath,
     ident: &Ident,
     name: &str,
+    extends: Option<&Path>,
 ) -> TokenStream {
-    IntrospectionNode::Map(
-        [
-            ("type", IntrospectionNode::String("class".into())),
-            (
-                "id",
-                IntrospectionNode::IntrospectionId(Some(ident_to_type(ident))),
-            ),
-            ("name", IntrospectionNode::String(name.into())),
-        ]
-        .into(),
-    )
-    .emit(pyo3_crate_path)
+    let mut desc = HashMap::from([
+        ("type", IntrospectionNode::String("class".into())),
+        (
+            "id",
+            IntrospectionNode::IntrospectionId(Some(ident_to_type(ident))),
+        ),
+        ("name", IntrospectionNode::String(name.into())),
+    ]);
+    if let Some(extends) = extends {
+        desc.insert(
+            "bases",
+            IntrospectionNode::List(vec![IntrospectionNode::BaseType(extends).into()]),
+        );
+    }
+    IntrospectionNode::Map(desc).emit(pyo3_crate_path)
 }
 
 #[allow(clippy::too_many_arguments)]
@@ -345,6 +349,7 @@ enum IntrospectionNode<'a> {
     IntrospectionId(Option<Cow<'a, Type>>),
     InputType { rust_type: Type, nullable: bool },
     OutputType { rust_type: Type, is_final: bool },
+    BaseType(&'a Path),
     Map(HashMap<&'static str, IntrospectionNode<'a>>),
     List(Vec<AttributedIntrospectionNode<'a>>),
 }
@@ -410,6 +415,13 @@ impl IntrospectionNode<'_> {
                 }
                 content.push_str("\"");
             }
+            Self::BaseType(path) => {
+                content.push_str("\"");
+                content.push_tokens(
+                    quote! { <#path as #pyo3_crate_path::impl_::pyclass::PyClassBaseType>::BASE_NAME.as_bytes() },
+                );
+                content.push_str("\"");
+            }
             Self::Map(map) => {
                 content.push_str("{");
                 for (i, (key, value)) in map.into_iter().enumerate() {

pyo3-macros-backend/src/pyclass.rs

@@ -399,7 +399,7 @@ impl FieldPyO3Options {
     }
 }
 
-fn get_class_python_name<'a>(cls: &'a syn::Ident, args: &'a PyClassArgs) -> Cow<'a, syn::Ident> {
+fn get_class_python_name<'a>(cls: &'a Ident, args: &'a PyClassArgs) -> Cow<'a, Ident> {
     args.options
         .name
         .as_ref()
@@ -2352,12 +2352,14 @@ impl<'a> PyClassImplsBuilder<'a> {
         };
 
         let pyclass_base_type_impl = attr.options.subclass.map(|subclass| {
+            let name_with_module = get_class_python_module_and_name(cls, self.attr);
             quote_spanned! { subclass.span() =>
                 impl #pyo3_path::impl_::pyclass::PyClassBaseType for #cls {
                     type LayoutAsBase = #pyo3_path::impl_::pycell::PyClassObject<Self>;
                     type BaseNativeType = <Self as #pyo3_path::impl_::pyclass::PyClassImpl>::BaseNativeType;
                     type Initializer = #pyo3_path::pyclass_init::PyClassInitializer<Self>;
                     type PyClassMutability = <Self as #pyo3_path::impl_::pyclass::PyClassImpl>::PyClassMutability;
+                    const BASE_NAME: &'static str = #name_with_module;
                 }
             }
         });
@@ -2498,9 +2500,14 @@ impl<'a> PyClassImplsBuilder<'a> {
     #[cfg(feature = "experimental-inspect")]
     fn impl_introspection(&self, ctx: &Ctx) -> TokenStream {
         let Ctx { pyo3_path, .. } = ctx;
-        let name = get_class_python_name(self.cls, self.attr).to_string();
         let ident = self.cls;
-        let static_introspection = class_introspection_code(pyo3_path, ident, &name);
+        let name = get_class_python_name(ident, self.attr).to_string();
+        let static_introspection = class_introspection_code(
+            pyo3_path,
+            ident,
+            &name,
+            self.attr.options.extends.as_ref().map(|attr| &attr.value),
+        );
         let introspection_id = introspection_id_const();
         quote! {
             #static_introspection

pytests/src/lib.rs

@@ -25,6 +25,7 @@ mod pyo3_pytests {
     #[pymodule_export]
     use {
         comparisons::comparisons, consts::consts, pyclasses::pyclasses, pyfunctions::pyfunctions,
+        subclassing::subclassing,
     };
 
     // Inserting to sys.modules allows importing submodules nicely from Python
@@ -43,7 +44,6 @@ mod pyo3_pytests {
         m.add_wrapped(wrap_pymodule!(othermod::othermod))?;
         m.add_wrapped(wrap_pymodule!(path::path))?;
         m.add_wrapped(wrap_pymodule!(sequence::sequence))?;
-        m.add_wrapped(wrap_pymodule!(subclassing::subclassing))?;
 
         // Inserting to sys.modules allows importing submodules nicely from Python
         // e.g. import pyo3_pytests.buf_and_str as bas

pytests/src/subclassing.rs

@@ -1,6 +1,8 @@
 //! Test for [#220](https://github.com/PyO3/pyo3/issues/220)
 
 use pyo3::prelude::*;
+#[cfg(not(Py_LIMITED_API))]
+use pyo3::types::PyDict;
 
 #[pyclass(subclass)]
 pub struct Subclassable {}
@@ -17,8 +19,43 @@ impl Subclassable {
     }
 }
 
+#[pyclass(extends = Subclassable)]
+pub struct Subclass {}
+
+#[pymethods]
+impl Subclass {
+    #[new]
+    fn new() -> (Self, Subclassable) {
+        (Subclass {}, Subclassable::new())
+    }
+
+    fn __str__(&self) -> &'static str {
+        "Subclass"
+    }
+}
+
+#[cfg(not(Py_LIMITED_API))]
+#[pyclass(extends = PyDict)]
+pub struct SubDict {}
+
+#[cfg(not(Py_LIMITED_API))]
+#[pymethods]
+impl SubDict {
+    #[new]
+    fn new() -> Self {
+        Self {}
+    }
+
+    fn __str__(&self) -> &'static str {
+        "SubDict"
+    }
+}
+
 #[pymodule(gil_used = false)]
-pub fn subclassing(m: &Bound<'_, PyModule>) -> PyResult<()> {
-    m.add_class::<Subclassable>()?;
-    Ok(())
+pub mod subclassing {
+    #[cfg(not(Py_LIMITED_API))]
+    #[pymodule_export]
+    use super::SubDict;
+    #[pymodule_export]
+    use super::{Subclass, Subclassable};
 }

pytests/stubs/subclassing.pyi

@@ -0,0 +1,11 @@
+class SubDict(dict):
+    def __new__(cls, /) -> None: ...
+    def __str__(self, /) -> str: ...
+
+class Subclass(Subclassable):
+    def __new__(cls, /) -> None: ...
+    def __str__(self, /) -> str: ...
+
+class Subclassable:
+    def __new__(cls, /) -> None: ...
+    def __str__(self, /) -> str: ...

pytests/tests/test_subclassing.py

@@ -1,12 +1,18 @@
-from pyo3_pytests.subclassing import Subclassable
+from pyo3_pytests.subclassing import Subclassable, Subclass
 
 
 class SomeSubClass(Subclassable):
     def __str__(self):
         return "SomeSubclass"
 
 
-def test_subclassing():
+def test_python_subclassing():
     a = SomeSubClass()
     assert str(a) == "SomeSubclass"
     assert type(a) is SomeSubClass
+
+
+def test_rust_subclassing():
+    a = Subclass()
+    assert str(a) == "Subclass"
+    assert type(a) is Subclass