Introspection: generate nested classes for complex enums

newsfragments/5708.added.md

@@ -0,0 +1 @@
+Introspection: generate nested classes for complex enums

pyo3-introspection/src/introspection.rs

@@ -31,7 +31,8 @@ fn parse_chunks(chunks: &[Chunk], main_module_name: &str) -> Result<Module> {
     let mut chunks_by_parent = HashMap::<&str, Vec<&Chunk>>::new();
     for chunk in chunks {
         let (id, parent) = match chunk {
-            Chunk::Module { id, .. } | Chunk::Class { id, .. } => (Some(id.as_str()), None),
+            Chunk::Module { id, .. } => (Some(id.as_str()), None),
+            Chunk::Class { id, parent, .. } => (Some(id.as_str()), parent.as_deref()),
             Chunk::Function { id, parent, .. } | Chunk::Attribute { id, parent, .. } => {
                 (id.as_deref(), parent.as_deref())
             }
@@ -148,6 +149,7 @@ fn convert_members<'a>(
                 id,
                 bases,
                 decorators,
+                parent: _,
             } => classes.push(convert_class(
                 id,
                 name,
@@ -234,10 +236,6 @@ fn convert_class(
         nested_modules.is_empty(),
         "Classes cannot contain nested modules"
     );
-    ensure!(
-        nested_classes.is_empty(),
-        "Nested classes are not supported yet"
-    );
     Ok(Class {
         name: name.into(),
         bases: bases
@@ -250,6 +248,7 @@ fn convert_class(
             .iter()
             .map(|e| convert_expr(e, type_hint_for_annotation_id))
             .collect(),
+        inner_classes: nested_classes,
     })
 }
 
@@ -408,7 +407,7 @@ fn introspection_id_to_type_hint_for_root_module(
     chunks_by_id: &HashMap<&str, &Chunk>,
     chunks_by_parent: &HashMap<&str, Vec<&Chunk>>,
 ) -> HashMap<String, Expr> {
-    fn add_introspection_id_to_type_hint_for_module(
+    fn add_introspection_id_to_type_hint_for_module_members(
         module_id: &str,
         module_full_name: &str,
         module_members: &[String],
@@ -426,13 +425,12 @@ fn introspection_id_to_type_hint_for_root_module(
                     .filter_map(|id| chunks_by_id.get(id.as_str())),
             )
             .copied()
-            .collect::<Vec<_>>()
         {
             match member {
                 Chunk::Module {
                     name, id, members, ..
                 } => {
-                    add_introspection_id_to_type_hint_for_module(
+                    add_introspection_id_to_type_hint_for_module_members(
                         id,
                         &format!("{}.{}", module_full_name, name),
                         members,
@@ -451,20 +449,57 @@ fn introspection_id_to_type_hint_for_root_module(
                             attr: name.clone(),
                         },
                     );
+                    add_introspection_id_to_type_hint_for_class_subclasses(
+                        id,
+                        name,
+                        module_full_name,
+                        chunks_by_parent,
+                        output,
+                    );
                 }
                 _ => (),
             }
         }
     }
 
+    fn add_introspection_id_to_type_hint_for_class_subclasses(
+        class_id: &str,
+        class_name: &str,
+        class_module: &str,
+        chunks_by_parent: &HashMap<&str, Vec<&Chunk>>,
+        output: &mut HashMap<String, Expr>,
+    ) {
+        for member in chunks_by_parent.get(&class_id).into_iter().flatten() {
+            if let Chunk::Class { id, name, .. } = member {
+                let class_name = format!("{}.{}", class_name, name);
+                add_introspection_id_to_type_hint_for_class_subclasses(
+                    id,
+                    &class_name,
+                    class_module,
+                    chunks_by_parent,
+                    output,
+                );
+                output.insert(
+                    id.clone(),
+                    Expr::Attribute {
+                        value: Box::new(Expr::Name {
+                            id: class_module.into(),
+                        }),
+                        attr: class_name,
+                    },
+                );
+            }
+        }
+    }
+
     let mut output = HashMap::new();
     let Chunk::Module {
         id, name, members, ..
     } = module_chunk
     else {
         unreachable!("The chunk must be a module")
     };
-    add_introspection_id_to_type_hint_for_module(
+    add_introspection_id_to_type_hint_for_module_members(
         id,
         name,
         members,
@@ -620,6 +655,8 @@ enum Chunk {
         bases: Vec<ChunkExpr>,
         #[serde(default)]
         decorators: Vec<ChunkExpr>,
+        #[serde(default)]
+        parent: Option<String>,
     },
     Function {
         #[serde(default)]

pyo3-introspection/src/model.rs

@@ -16,6 +16,7 @@ pub struct Class {
     pub attributes: Vec<Attribute>,
     /// decorator like 'typing.final'
     pub decorators: Vec<Expr>,
+    pub inner_classes: Vec<Class>,
 }
 
 #[derive(Debug, Eq, PartialEq, Clone, Hash)]

pyo3-introspection/src/stubs.rs

@@ -135,7 +135,7 @@ fn class_stubs(class: &Class, imports: &Imports) -> String {
         buffer.push(')');
     }
     buffer.push(':');
-    if class.methods.is_empty() && class.attributes.is_empty() {
+    if class.methods.is_empty() && class.attributes.is_empty() && class.inner_classes.is_empty() {
         buffer.push_str(" ...");
         return buffer;
     }
@@ -150,6 +150,11 @@ fn class_stubs(class: &Class, imports: &Imports) -> String {
         buffer
             .push_str(&function_stubs(method, imports, Some(&class.name)).replace('\n', "\n    "));
     }
+    for inner_class in &class.inner_classes {
+        // We do the indentation
+        buffer.push_str("\n    ");
+        buffer.push_str(&class_stubs(inner_class, imports).replace('\n', "\n    "));
+    }
     buffer
 }
 
@@ -509,6 +514,9 @@ impl ElementsUsedInAnnotations {
         for attr in &class.attributes {
             self.walk_attribute(attr);
         }
+        for class in &class.inner_classes {
+            self.walk_class(class);
+        }
     }
 
     fn walk_attribute(&mut self, attribute: &Attribute) {
@@ -758,13 +766,15 @@ mod tests {
                             }),
                             attr: "final".into(),
                         }],
+                        inner_classes: Vec::new(),
                     },
                     Class {
                         name: "int".into(),
                         bases: Vec::new(),
                         methods: Vec::new(),
                         attributes: Vec::new(),
                         decorators: Vec::new(),
+                        inner_classes: Vec::new(),
                     },
                 ],
                 functions: vec![Function {

pyo3-macros-backend/src/introspection.rs

@@ -62,6 +62,7 @@ pub fn class_introspection_code(
     name: &str,
     extends: Option<PythonTypeHint>,
     is_final: bool,
+    parent: Option<&Type>,
 ) -> TokenStream {
     let mut desc = HashMap::from([
         ("type", IntrospectionNode::String("class".into())),
@@ -80,6 +81,12 @@ pub fn class_introspection_code(
             IntrospectionNode::List(vec![PythonTypeHint::module_attr("typing", "final").into()]),
         );
     }
+    if let Some(parent) = parent {
+        desc.insert(
+            "parent",
+            IntrospectionNode::IntrospectionId(Some(Cow::Borrowed(parent))),
+        );
+    }
     IntrospectionNode::Map(desc).emit(pyo3_crate_path)
 }
 
@@ -98,7 +105,6 @@ pub fn function_introspection_code(
     let mut desc = HashMap::from([
         ("type", IntrospectionNode::String("function".into())),
         ("name", IntrospectionNode::String(name.into())),
-        ("async", IntrospectionNode::Bool(is_async)),
         (
             "arguments",
             arguments_introspection_data(signature, first_argument, parent),
@@ -120,6 +126,9 @@ pub fn function_introspection_code(
             },
         ),
     ]);
+    if is_async {
+        desc.insert("async", IntrospectionNode::Bool(true));
+    }
     if let Some(ident) = ident {
         desc.insert(
             "id",

pyo3-macros-backend/src/konst.rs

@@ -14,6 +14,10 @@ use syn::{
 pub struct ConstSpec {
     pub rust_ident: syn::Ident,
     pub attributes: ConstAttributes,
+    #[cfg(feature = "experimental-inspect")]
+    pub expr: Option<syn::Expr>,
+    #[cfg(feature = "experimental-inspect")]
+    pub ty: syn::Type,
 }
 
 impl ConstSpec {