Introspection: add function signatures

pyo3-introspection/src/introspection.rs

@@ -1,4 +1,4 @@
-use crate::model::{Class, Function, Module};
+use crate::model::{Argument, Arguments, Class, Function, Module, VariableLengthArgument};
 use anyhow::{bail, ensure, Context, Result};
 use goblin::elf::Elf;
 use goblin::mach::load_command::CommandVariant;
@@ -43,14 +43,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>,
@@ -66,10 +66,37 @@ 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 { name, id: _ } => functions.push(Function { name: name.into() }),
+                Chunk::Function {
+                    name,
+                    id: _,
+                    arguments,
+                } => functions.push(Function {
+                    name: name.into(),
+                    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_variable_length_argument),
+                        keyword_only_arguments: arguments
+                            .kwonlyargs
+                            .iter()
+                            .map(convert_argument)
+                            .collect(),
+                        kwarg: arguments
+                            .kwarg
+                            .as_ref()
+                            .map(convert_variable_length_argument),
+                    },
+                }),
             }
         }
     }
@@ -81,6 +108,19 @@ fn parse_module(
     })
 }
 
+fn convert_argument(arg: &ChunkArgument) -> Argument {
+    Argument {
+        name: arg.name.clone(),
+        default_value: arg.default.clone(),
+    }
+}
+
+fn convert_variable_length_argument(arg: &ChunkArgument) -> VariableLengthArgument {
+    VariableLengthArgument {
+        name: arg.name.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()))?;
@@ -252,5 +292,27 @@ enum Chunk {
     Function {
         id: String,
         name: String,
+        arguments: ChunkArguments,
     },
 }
+
+#[derive(Deserialize)]
+struct ChunkArguments {
+    #[serde(default)]
+    posonlyargs: Vec<ChunkArgument>,
+    #[serde(default)]
+    args: Vec<ChunkArgument>,
+    #[serde(default)]
+    vararg: Option<ChunkArgument>,
+    #[serde(default)]
+    kwonlyargs: Vec<ChunkArgument>,
+    #[serde(default)]
+    kwarg: Option<ChunkArgument>,
+}
+
+#[derive(Deserialize)]
+struct ChunkArgument {
+    name: String,
+    #[serde(default)]
+    default: Option<String>,
+}

pyo3-introspection/src/model.rs

@@ -14,4 +14,32 @@ 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<VariableLengthArgument>,
+    /// Arguments after *
+    pub keyword_only_arguments: Vec<Argument>,
+    /// **kwarg
+    pub kwarg: Option<VariableLengthArgument>,
+}
+
+#[derive(Debug, Eq, PartialEq, Clone, Hash)]
+pub struct Argument {
+    pub name: String,
+    /// Default value as a Python expression
+    pub default_value: Option<String>,
+}
+
+/// A variable length argument ie. *vararg or **kwarg
+#[derive(Debug, Eq, PartialEq, Clone, Hash)]
+pub struct VariableLengthArgument {
+    pub name: String,
 }

pyo3-introspection/src/stubs.rs

@@ -1,4 +1,4 @@
-use crate::model::{Class, Function, Module};
+use crate::model::{Argument, Class, Function, Module, VariableLengthArgument};
 use std::collections::HashMap;
 use std::path::{Path, PathBuf};
 
@@ -48,5 +48,104 @@ fn class_stubs(class: &Class) -> String {
 }
 
 fn function_stubs(function: &Function) -> String {
-    format!("def {}(*args, **kwargs): ...", function.name)
+    // Signature
+    let mut parameters = Vec::new();
+    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!("*{}", variable_length_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!("**{}", variable_length_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
+}
+
+fn variable_length_argument_stub(argument: &VariableLengthArgument) -> String {
+    argument.name.clone()
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::model::Arguments;
+
+    #[test]
+    fn function_stubs_with_variable_length() {
+        let function = Function {
+            name: "func".into(),
+            arguments: Arguments {
+                positional_only_arguments: vec![Argument {
+                    name: "posonly".into(),
+                    default_value: None,
+                }],
+                arguments: vec![Argument {
+                    name: "arg".into(),
+                    default_value: None,
+                }],
+                vararg: Some(VariableLengthArgument {
+                    name: "varargs".into(),
+                }),
+                keyword_only_arguments: vec![Argument {
+                    name: "karg".into(),
+                    default_value: None,
+                }],
+                kwarg: Some(VariableLengthArgument {
+                    name: "kwarg".into(),
+                }),
+            },
+        };
+        assert_eq!(
+            "def func(posonly, /, arg, *varargs, karg, **kwarg): ...",
+            function_stubs(&function)
+        )
+    }
+
+    #[test]
+    fn function_stubs_without_variable_length() {
+        let function = Function {
+            name: "afunc".into(),
+            arguments: Arguments {
+                positional_only_arguments: vec![Argument {
+                    name: "posonly".into(),
+                    default_value: Some("1".into()),
+                }],
+                arguments: vec![Argument {
+                    name: "arg".into(),
+                    default_value: Some("True".into()),
+                }],
+                vararg: None,
+                keyword_only_arguments: vec![Argument {
+                    name: "karg".into(),
+                    default_value: Some("\"foo\"".into()),
+                }],
+                kwarg: None,
+            },
+        };
+        assert_eq!(
+            "def afunc(posonly=1, /, arg=True, *, karg=\"foo\"): ...",
+            function_stubs(&function)
+        )
+    }
 }