Refine arguments struct

Tpt · Tpt · commit 5b1f54ae6df7 · 2025-04-04T17:20:00.000+02:00
diff --git a/pyo3-introspection/src/introspection.rs b/pyo3-introspection/src/introspection.rs
@@ -1,4 +1,4 @@
-use crate::model::{Argument, Class, Function, Module, ParameterKind};
+use crate::model::{Argument, Arguments, Class, Function, Module};
 use anyhow::{bail, Context, Result};
 use goblin::elf::Elf;
 use goblin::mach::symbols::N_SECT;
@@ -42,14 +42,14 @@ fn parse_chunks(chunks: &[Chunk], main_module_name: &str) -> Result<Module> {
         } = chunk
         {
             if name == main_module_name {
-                return parse_module(name, members, &chunks_by_id);
+                return convert_module(name, members, &chunks_by_id);
             }
         }
     }
     bail!("No module named {main_module_name} found")
 }
 
-fn parse_module(
+fn convert_module(
     name: &str,
     members: &[String],
     chunks_by_id: &HashMap<&String, &Chunk>,
@@ -65,7 +65,7 @@ fn parse_module(
                     members,
                     id: _,
                 } => {
-                    modules.push(parse_module(name, members, chunks_by_id)?);
+                    modules.push(convert_module(name, members, chunks_by_id)?);
                 }
                 Chunk::Class { name, id: _ } => classes.push(Class { name: name.into() }),
                 Chunk::Function {
@@ -74,22 +74,21 @@ fn parse_module(
                     arguments,
                 } => functions.push(Function {
                     name: name.into(),
-                    arguments: arguments
-                        .iter()
-                        .map(|arg| Argument {
-                            name: arg.name.clone(),
-                            kind: match arg.kind {
-                                ChunkArgumentKind::PositionalOnly => ParameterKind::PositionalOnly,
-                                ChunkArgumentKind::PositionalOrKeyword => {
-                                    ParameterKind::PositionalOrKeyword
-                                }
-                                ChunkArgumentKind::VarPositional => ParameterKind::VarPositional,
-                                ChunkArgumentKind::KeywordOnly => ParameterKind::KeywordOnly,
-                                ChunkArgumentKind::VarKeyword => ParameterKind::VarKeyword,
-                            },
-                            default_value: arg.default_value.clone(),
-                        })
-                        .collect(),
+                    arguments: Arguments {
+                        positional_only_arguments: arguments
+                            .posonlyargs
+                            .iter()
+                            .map(convert_argument)
+                            .collect(),
+                        arguments: arguments.args.iter().map(convert_argument).collect(),
+                        vararg: arguments.vararg.as_ref().map(convert_argument),
+                        keyword_only_arguments: arguments
+                            .kwonlyargs
+                            .iter()
+                            .map(convert_argument)
+                            .collect(),
+                        kwarg: arguments.kwarg.as_ref().map(convert_argument),
+                    },
                 }),
             }
         }
@@ -102,6 +101,13 @@ fn parse_module(
     })
 }
 
+fn convert_argument(arg: &ChunkArgument) -> Argument {
+    Argument {
+        name: arg.name.clone(),
+        default_value: arg.default.clone(),
+    }
+}
+
 fn find_introspection_chunks_in_binary_object(path: &Path) -> Result<Vec<Chunk>> {
     let library_content =
         fs::read(path).with_context(|| format!("Failed to read {}", path.display()))?;
@@ -262,24 +268,27 @@ enum Chunk {
     Function {
         id: String,
         name: String,
-        arguments: Vec<ChunkArgument>,
+        arguments: ChunkArguments,
     },
 }
 
 #[derive(Deserialize)]
-struct ChunkArgument {
-    name: String,
-    kind: ChunkArgumentKind,
+struct ChunkArguments {
+    #[serde(default)]
+    posonlyargs: Vec<ChunkArgument>,
     #[serde(default)]
-    default_value: Option<String>,
+    args: Vec<ChunkArgument>,
+    #[serde(default)]
+    vararg: Option<ChunkArgument>,
+    #[serde(default)]
+    kwonlyargs: Vec<ChunkArgument>,
+    #[serde(default)]
+    kwarg: Option<ChunkArgument>,
 }
 
 #[derive(Deserialize)]
-#[serde(rename_all = "SCREAMING_SNAKE_CASE")]
-enum ChunkArgumentKind {
-    PositionalOnly,
-    PositionalOrKeyword,
-    VarPositional,
-    KeywordOnly,
-    VarKeyword,
+struct ChunkArgument {
+    name: String,
+    #[serde(default)]
+    default: Option<String>,
 }
diff --git a/pyo3-introspection/src/model.rs b/pyo3-introspection/src/model.rs
@@ -14,27 +14,26 @@ pub struct Class {
 #[derive(Debug, Eq, PartialEq, Clone, Hash)]
 pub struct Function {
     pub name: String,
+    pub arguments: Arguments,
+}
+
+#[derive(Debug, Eq, PartialEq, Clone, Hash)]
+pub struct Arguments {
+    /// Arguments before /
+    pub positional_only_arguments: Vec<Argument>,
+    /// Regular arguments (between / and *)
     pub arguments: Vec<Argument>,
+    /// *vararg
+    pub vararg: Option<Argument>,
+    /// Arguments after *
+    pub keyword_only_arguments: Vec<Argument>,
+    /// **kwarg
+    pub kwarg: Option<Argument>,
 }
 
 #[derive(Debug, Eq, PartialEq, Clone, Hash)]
 pub struct Argument {
     pub name: String,
-    pub kind: ParameterKind,
     /// Default value as a Python expression
     pub default_value: Option<String>,
 }
-
-#[derive(Debug, Eq, PartialEq, Clone, Copy, Hash)]
-pub enum ParameterKind {
-    /// Before /
-    PositionalOnly,
-    /// Between / and *
-    PositionalOrKeyword,
-    /// *args
-    VarPositional,
-    /// After *
-    KeywordOnly,
-    /// *kwargs
-    VarKeyword,
-}
diff --git a/pyo3-introspection/src/stubs.rs b/pyo3-introspection/src/stubs.rs
@@ -1,4 +1,4 @@
-use crate::model::{Class, Function, Module, ParameterKind};
+use crate::model::{Argument, Class, Function, Module};
 use std::collections::HashMap;
 use std::path::{Path, PathBuf};
 
@@ -49,33 +49,35 @@ fn class_stubs(class: &Class) -> String {
 
 fn function_stubs(function: &Function) -> String {
     // Signature
-    let mut positional_only = true;
-    let mut keyword_only = false;
     let mut parameters = Vec::new();
-    for argument in &function.arguments {
-        if positional_only && !matches!(argument.kind, ParameterKind::PositionalOnly) {
-            if !parameters.is_empty() {
-                parameters.push("/".into());
-            }
-            positional_only = false;
-        }
-        if !keyword_only && matches!(argument.kind, ParameterKind::KeywordOnly) {
-            parameters.push("*".into());
-            keyword_only = true;
-        }
-        let mut parameter_str = match argument.kind {
-            ParameterKind::VarPositional => {
-                keyword_only = true;
-                format!("*{}", argument.name)
-            }
-            ParameterKind::VarKeyword => format!("**{}", argument.name),
-            _ => argument.name.clone(),
-        };
-        if let Some(default_value) = &argument.default_value {
-            parameter_str.push('=');
-            parameter_str.push_str(default_value);
-        }
-        parameters.push(parameter_str);
+    for argument in &function.arguments.positional_only_arguments {
+        parameters.push(argument_stub(argument));
+    }
+    if !function.arguments.positional_only_arguments.is_empty() {
+        parameters.push("/".into());
+    }
+    for argument in &function.arguments.arguments {
+        parameters.push(argument_stub(argument));
+    }
+    if let Some(argument) = &function.arguments.vararg {
+        parameters.push(format!("*{}", argument_stub(argument)));
+    } else if !function.arguments.keyword_only_arguments.is_empty() {
+        parameters.push("*".into());
+    }
+    for argument in &function.arguments.keyword_only_arguments {
+        parameters.push(argument_stub(argument));
+    }
+    if let Some(argument) = &function.arguments.kwarg {
+        parameters.push(format!("**{}", argument_stub(argument)));
     }
     format!("def {}({}): ...", function.name, parameters.join(", "))
 }
+
+fn argument_stub(argument: &Argument) -> String {
+    let mut output = argument.name.clone();
+    if let Some(default_value) = &argument.default_value {
+        output.push('=');
+        output.push_str(default_value);
+    }
+    output
+}
diff --git a/pyo3-macros-backend/src/introspection.rs b/pyo3-macros-backend/src/introspection.rs
@@ -99,7 +99,11 @@ fn arguments_introspection_data<'a>(signature: &'a FunctionSignature<'a>) -> Int
         }
     });
 
-    let mut arguments = Vec::new();
+    let mut posonlyargs = Vec::new();
+    let mut args = Vec::new();
+    let mut vararg = None;
+    let mut kwonlyargs = Vec::new();
+    let mut kwarg = None;
 
     for (i, param) in signature
         .python_signature
@@ -112,24 +116,17 @@ fn arguments_introspection_data<'a>(signature: &'a FunctionSignature<'a>) -> Int
         } else {
             panic!("Less arguments than in python signature");
         };
-        arguments.push(argument_introspection_data(
-            param,
-            if i < signature.python_signature.positional_only_parameters {
-                "POSITIONAL_ONLY"
-            } else {
-                "POSITIONAL_OR_KEYWORD"
-            },
-            arg_desc,
-        ));
+        let arg = argument_introspection_data(param, arg_desc);
+        if i < signature.python_signature.positional_only_parameters {
+            posonlyargs.push(arg);
+        } else {
+            args.push(arg)
+        }
     }
 
     if let Some(param) = &signature.python_signature.varargs {
-        arguments.push(IntrospectionNode::Map(
-            [
-                ("name", IntrospectionNode::String(param.into())),
-                ("kind", IntrospectionNode::String("VAR_POSITIONAL".into())),
-            ]
-            .into(),
+        vararg = Some(IntrospectionNode::Map(
+            [("name", IntrospectionNode::String(param.into()))].into(),
         ));
     }
 
@@ -139,11 +136,11 @@ fn arguments_introspection_data<'a>(signature: &'a FunctionSignature<'a>) -> Int
         } else {
             panic!("Less arguments than in python signature");
         };
-        arguments.push(argument_introspection_data(param, "KEYWORD_ONLY", arg_desc));
+        kwonlyargs.push(argument_introspection_data(param, arg_desc));
     }
 
     if let Some(param) = &signature.python_signature.kwargs {
-        arguments.push(IntrospectionNode::Map(
+        kwarg = Some(IntrospectionNode::Map(
             [
                 ("name", IntrospectionNode::String(param.into())),
                 ("kind", IntrospectionNode::String("VAR_KEYWORD".into())),
@@ -152,22 +149,33 @@ fn arguments_introspection_data<'a>(signature: &'a FunctionSignature<'a>) -> Int
         ));
     }
 
-    IntrospectionNode::List(arguments)
+    let mut map = HashMap::new();
+    if !posonlyargs.is_empty() {
+        map.insert("posonlyargs", IntrospectionNode::List(posonlyargs));
+    }
+    if !args.is_empty() {
+        map.insert("args", IntrospectionNode::List(args));
+    }
+    if let Some(vararg) = vararg {
+        map.insert("vararg", vararg);
+    }
+    if !kwonlyargs.is_empty() {
+        map.insert("kwonlyargs", IntrospectionNode::List(kwonlyargs));
+    }
+    if let Some(kwarg) = kwarg {
+        map.insert("kwarg", kwarg);
+    }
+    IntrospectionNode::Map(map)
 }
 
 fn argument_introspection_data<'a>(
     name: &'a str,
-    kind: &'a str,
     desc: &'a RegularArg<'_>,
 ) -> IntrospectionNode<'a> {
-    let mut params: HashMap<_, _> = [
-        ("name", IntrospectionNode::String(name.into())),
-        ("kind", IntrospectionNode::String(kind.into())),
-    ]
-    .into();
+    let mut params: HashMap<_, _> = [("name", IntrospectionNode::String(name.into()))].into();
     if desc.default_value.is_some() {
         params.insert(
-            "default_value",
+            "default",
             IntrospectionNode::String(desc.default_value().into()),
         );
     }
