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

@@ -32,7 +32,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())
             }
@@ -129,6 +130,7 @@ fn convert_members<'a>(
                 id,
                 bases,
                 decorators,
+                parent: _,
             } => classes.push(convert_class(
                 id,
                 name,
@@ -202,10 +204,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
@@ -218,6 +216,7 @@ fn convert_class(
             .iter()
             .map(convert_python_identifier)
             .collect::<Result<_>>()?,
+        inner_classes: nested_classes,
     })
 }
 
@@ -472,6 +471,8 @@ enum Chunk {
         bases: Vec<ChunkTypeHint>,
         #[serde(default)]
         decorators: Vec<ChunkTypeHint>,
+        #[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<PythonIdentifier>,
+    pub inner_classes: Vec<Class>,
 }
 
 #[derive(Debug, Eq, PartialEq, Clone, Hash)]

pyo3-introspection/src/stubs.rs

@@ -138,7 +138,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;
     }
@@ -152,6 +152,11 @@ fn class_stubs(class: &Class, imports: &Imports) -> String {
         buffer.push_str("\n    ");
         buffer.push_str(&function_stubs(method, imports).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
 }
 
@@ -690,6 +695,7 @@ mod tests {
                         module: Some("typing".into()),
                         name: "final".into(),
                     }],
+                    inner_classes: Vec::new(),
                 }],
                 functions: vec![Function {
                     name: String::new(),

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)
 }