Skip to content

Implements basic method introspection #5087

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 1 commit into
base: main
Choose a base branch
from
Open

Implements basic method introspection #5087

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
178 changes: 120 additions & 58 deletions pyo3-introspection/src/introspection.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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 {
Expand All @@ -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);
}
}
}
Expand All @@ -53,59 +58,111 @@ 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 {
chunks_by_id,
chunks_by_parent,
)?);
}
Chunk::Class { name, id } => {
let (_, _, mut methods) = convert_members(
chunks_by_parent
.get(&id.as_str())
.map(Vec::as_slice)
.unwrap_or_default(),
chunks_by_id,
chunks_by_parent,
)?;
// 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,
});
classes.push(Class {
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),
},
}),
methods,
})
}
Chunk::Function {
name,
id: _,
arguments,
parent: _,
decorators,
} => functions.push(Function {
name: name.into(),
decorators: decorators.clone(),
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),
},
}),
}
}
Ok(Module {
name: name.into(),
modules,
classes,
functions,
})
Ok((modules, classes, functions))
}

fn convert_argument(arg: &ChunkArgument) -> Argument {
Expand Down Expand Up @@ -290,9 +347,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>,
},
}

Expand Down
3 changes: 3 additions & 0 deletions pyo3-introspection/src/model.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
}

Expand Down
48 changes: 44 additions & 4 deletions pyo3-introspection/src/stubs.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand All @@ -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 {
Expand All @@ -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(),
Expand Down Expand Up @@ -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(),
Expand Down
Loading
Loading