Implements basic method introspection (#5087)

* Implements basic method introspection

* Code review feedback

* Fixes CI

Author: Tpt

pyo3-introspection/src/introspection.rs

@@ -7,6 +7,7 @@ use goblin::mach::{Mach, MachO, SingleArch};
 use goblin::pe::PE;
 use goblin::Object;
 use serde::Deserialize;
+use std::cmp::Ordering;
 use std::collections::HashMap;
 use std::fs;
 use std::path::Path;
@@ -21,19 +22,23 @@ pub fn introspect_cdylib(library_path: impl AsRef<Path>, main_module_name: &str)
 
 /// Parses the introspection chunks found in the binary
 fn parse_chunks(chunks: &[Chunk], main_module_name: &str) -> Result<Module> {
-    let chunks_by_id = chunks
-        .iter()
-        .map(|c| {
-            (
-                match c {
-                    Chunk::Module { id, .. } => id,
-                    Chunk::Class { id, .. } => id,
-                    Chunk::Function { id, .. } => id,
-                },
-                c,
-            )
-        })
-        .collect::<HashMap<_, _>>();
+    let mut chunks_by_id = HashMap::<&str, &Chunk>::new();
+    let mut chunks_by_parent = HashMap::<&str, Vec<&Chunk>>::new();
+    for chunk in chunks {
+        if let Some(id) = match chunk {
+            Chunk::Module { id, .. } => Some(id),
+            Chunk::Class { id, .. } => Some(id),
+            Chunk::Function { id, .. } => id.as_ref(),
+        } {
+            chunks_by_id.insert(id, chunk);
+        }
+        if let Some(parent) = match chunk {
+            Chunk::Module { .. } | Chunk::Class { .. } => None,
+            Chunk::Function { parent, .. } => parent.as_ref(),
+        } {
+            chunks_by_parent.entry(parent).or_default().push(chunk);
+        }
+    }
     // We look for the root chunk
     for chunk in chunks {
         if let Chunk::Module {
@@ -43,7 +48,7 @@ fn parse_chunks(chunks: &[Chunk], main_module_name: &str) -> Result<Module> {
         } = chunk
         {
             if name == main_module_name {
-                return convert_module(name, members, &chunks_by_id);
+                return convert_module(name, members, &chunks_by_id, &chunks_by_parent);
             }
         }
     }
@@ -53,61 +58,126 @@ fn parse_chunks(chunks: &[Chunk], main_module_name: &str) -> Result<Module> {
 fn convert_module(
     name: &str,
     members: &[String],
-    chunks_by_id: &HashMap<&String, &Chunk>,
+    chunks_by_id: &HashMap<&str, &Chunk>,
+    chunks_by_parent: &HashMap<&str, Vec<&Chunk>>,
 ) -> Result<Module> {
+    let (modules, classes, functions) = convert_members(
+        &members
+            .iter()
+            .filter_map(|id| chunks_by_id.get(id.as_str()).copied())
+            .collect::<Vec<_>>(),
+        chunks_by_id,
+        chunks_by_parent,
+    )?;
+    Ok(Module {
+        name: name.into(),
+        modules,
+        classes,
+        functions,
+    })
+}
+
+/// Convert a list of members of a module or a class
+fn convert_members(
+    chunks: &[&Chunk],
+    chunks_by_id: &HashMap<&str, &Chunk>,
+    chunks_by_parent: &HashMap<&str, Vec<&Chunk>>,
+) -> Result<(Vec<Module>, Vec<Class>, Vec<Function>)> {
     let mut modules = Vec::new();
     let mut classes = Vec::new();
     let mut functions = Vec::new();
-    for member in members {
-        if let Some(chunk) = chunks_by_id.get(member) {
-            match chunk {
-                Chunk::Module {
+    for chunk in chunks {
+        match chunk {
+            Chunk::Module {
+                name,
+                members,
+                id: _,
+            } => {
+                modules.push(convert_module(
                     name,
                     members,
-                    id: _,
-                } => {
-                    modules.push(convert_module(name, members, chunks_by_id)?);
-                }
-                Chunk::Class { name, id: _ } => classes.push(Class { 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),
-                    },
-                }),
+                    chunks_by_id,
+                    chunks_by_parent,
+                )?);
             }
+            Chunk::Class { name, id } => {
+                classes.push(convert_class(id, name, chunks_by_id, chunks_by_parent)?)
+            }
+            Chunk::Function {
+                name,
+                id: _,
+                arguments,
+                parent: _,
+                decorators,
+            } => functions.push(convert_function(name, arguments, decorators)),
         }
     }
-    Ok(Module {
+    Ok((modules, classes, functions))
+}
+
+fn convert_class(
+    id: &str,
+    name: &str,
+    chunks_by_id: &HashMap<&str, &Chunk>,
+    chunks_by_parent: &HashMap<&str, Vec<&Chunk>>,
+) -> Result<Class> {
+    let (nested_modules, nested_classes, mut methods) = convert_members(
+        chunks_by_parent
+            .get(&id)
+            .map(Vec::as_slice)
+            .unwrap_or_default(),
+        chunks_by_id,
+        chunks_by_parent,
+    )?;
+    ensure!(
+        nested_modules.is_empty(),
+        "Classes cannot contain nested modules"
+    );
+    ensure!(
+        nested_classes.is_empty(),
+        "Nested classes are not supported yet"
+    );
+    // We sort methods to get a stable output
+    methods.sort_by(|l, r| match l.name.cmp(&r.name) {
+        Ordering::Equal => {
+            // We put the getter before the setter
+            if l.decorators.iter().any(|d| d == "property") {
+                Ordering::Less
+            } else if r.decorators.iter().any(|d| d == "property") {
+                Ordering::Greater
+            } else {
+                // We pick an ordering based on decorators
+                l.decorators.cmp(&r.decorators)
+            }
+        }
+        o => o,
+    });
+    Ok(Class {
         name: name.into(),
-        modules,
-        classes,
-        functions,
+        methods,
     })
 }
 
+fn convert_function(name: &str, arguments: &ChunkArguments, decorators: &[String]) -> Function {
+    Function {
+        name: name.into(),
+        decorators: decorators.to_vec(),
+        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),
+        },
+    }
+}
+
 fn convert_argument(arg: &ChunkArgument) -> Argument {
     Argument {
         name: arg.name.clone(),
@@ -290,9 +360,14 @@ enum Chunk {
         name: String,
     },
     Function {
-        id: String,
+        #[serde(default)]
+        id: Option<String>,
         name: String,
         arguments: ChunkArguments,
+        #[serde(default)]
+        parent: Option<String>,
+        #[serde(default)]
+        decorators: Vec<String>,
     },
 }
 

pyo3-introspection/src/model.rs

@@ -9,11 +9,14 @@ pub struct Module {
 #[derive(Debug, Eq, PartialEq, Clone, Hash)]
 pub struct Class {
     pub name: String,
+    pub methods: Vec<Function>,
 }
 
 #[derive(Debug, Eq, PartialEq, Clone, Hash)]
 pub struct Function {
     pub name: String,
+    /// decorator like 'property' or 'staticmethod'
+    pub decorators: Vec<String>,
     pub arguments: Arguments,
 }
 

pyo3-introspection/src/stubs.rs

@@ -39,12 +39,39 @@ fn module_stubs(module: &Module) -> String {
     for function in &module.functions {
         elements.push(function_stubs(function));
     }
-    elements.push(String::new()); // last line jump
-    elements.join("\n")
+
+    // We insert two line jumps (i.e. empty strings) only above and below multiple line elements (classes with methods, functions with decorators)
+    let mut output = String::new();
+    for element in elements {
+        let is_multiline = element.contains('\n');
+        if is_multiline && !output.is_empty() && !output.ends_with("\n\n") {
+            output.push('\n');
+        }
+        output.push_str(&element);
+        output.push('\n');
+        if is_multiline {
+            output.push('\n');
+        }
+    }
+    // We remove a line jump at the end if they are two
+    if output.ends_with("\n\n") {
+        output.pop();
+    }
+    output
 }
 
 fn class_stubs(class: &Class) -> String {
-    format!("class {}: ...", class.name)
+    let mut buffer = format!("class {}:", class.name);
+    if class.methods.is_empty() {
+        buffer.push_str(" ...");
+        return buffer;
+    }
+    for method in &class.methods {
+        // We do the indentation
+        buffer.push_str("\n    ");
+        buffer.push_str(&function_stubs(method).replace('\n', "\n    "));
+    }
+    buffer
 }
 
 fn function_stubs(function: &Function) -> String {
@@ -70,7 +97,18 @@ fn function_stubs(function: &Function) -> String {
     if let Some(argument) = &function.arguments.kwarg {
         parameters.push(format!("**{}", variable_length_argument_stub(argument)));
     }
-    format!("def {}({}): ...", function.name, parameters.join(", "))
+    let output = format!("def {}({}): ...", function.name, parameters.join(", "));
+    if function.decorators.is_empty() {
+        return output;
+    }
+    let mut buffer = String::new();
+    for decorator in &function.decorators {
+        buffer.push('@');
+        buffer.push_str(decorator);
+        buffer.push('\n');
+    }
+    buffer.push_str(&output);
+    buffer
 }
 
 fn argument_stub(argument: &Argument) -> String {
@@ -95,6 +133,7 @@ mod tests {
     fn function_stubs_with_variable_length() {
         let function = Function {
             name: "func".into(),
+            decorators: Vec::new(),
             arguments: Arguments {
                 positional_only_arguments: vec![Argument {
                     name: "posonly".into(),
@@ -126,6 +165,7 @@ mod tests {
     fn function_stubs_without_variable_length() {
         let function = Function {
             name: "afunc".into(),
+            decorators: Vec::new(),
             arguments: Arguments {
                 positional_only_arguments: vec![Argument {
                     name: "posonly".into(),

pyo3-introspection/tests/test.rs

@@ -42,9 +42,10 @@ fn pytests_stubs() -> Result<()> {
                 file_name.display()
             )
         });
+        // We normalize line jumps for compatibility with Windows
         assert_eq!(
-            &expected_file_content.replace('\r', ""), // Windows compatibility
-            actual_file_content,
+            expected_file_content.replace('\r', ""),
+            actual_file_content.replace('\r', ""),
             "The content of file {} is different",
             file_name.display()
         )