Skip to content

Commit f7fe032

Browse files
TptTpt
committed
Implements basic method introspection
1 parent 0062d5e commit f7fe032

File tree

10 files changed

+404
-100
lines changed

10 files changed

+404
-100
lines changed

‎pyo3-introspection/src/introspection.rs

+120-58
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@ use goblin::mach::{Mach, MachO, SingleArch};
77
use goblin::pe::PE;
88
use goblin::Object;
99
use serde::Deserialize;
10+
use std::cmp::Ordering;
1011
use std::collections::HashMap;
1112
use std::fs;
1213
use std::path::Path;
@@ -21,19 +22,23 @@ pub fn introspect_cdylib(library_path: impl AsRef<Path>, main_module_name: &str)
2122

2223
/// Parses the introspection chunks found in the binary
2324
fn parse_chunks(chunks: &[Chunk], main_module_name: &str) -> Result<Module> {
24-
let chunks_by_id = chunks
25-
.iter()
26-
.map(|c| {
27-
(
28-
match c {
29-
Chunk::Module { id, .. } => id,
30-
Chunk::Class { id, .. } => id,
31-
Chunk::Function { id, .. } => id,
32-
},
33-
c,
34-
)
35-
})
36-
.collect::<HashMap<_, _>>();
25+
let mut chunks_by_id = HashMap::<&str, &Chunk>::new();
26+
let mut chunks_by_parent = HashMap::<&str, Vec<&Chunk>>::new();
27+
for chunk in chunks {
28+
if let Some(id) = match chunk {
29+
Chunk::Module { id, .. } => Some(id),
30+
Chunk::Class { id, .. } => Some(id),
31+
Chunk::Function { id, .. } => id.as_ref(),
32+
} {
33+
chunks_by_id.insert(id, chunk);
34+
}
35+
if let Some(parent) = match chunk {
36+
Chunk::Module { .. } | Chunk::Class { .. } => None,
37+
Chunk::Function { parent, .. } => parent.as_ref(),
38+
} {
39+
chunks_by_parent.entry(parent).or_default().push(chunk);
40+
}
41+
}
3742
// We look for the root chunk
3843
for chunk in chunks {
3944
if let Chunk::Module {
@@ -43,7 +48,7 @@ fn parse_chunks(chunks: &[Chunk], main_module_name: &str) -> Result<Module> {
4348
} = chunk
4449
{
4550
if name == main_module_name {
46-
return convert_module(name, members, &chunks_by_id);
51+
return convert_module(name, members, &chunks_by_id, &chunks_by_parent);
4752
}
4853
}
4954
}
@@ -53,59 +58,111 @@ fn parse_chunks(chunks: &[Chunk], main_module_name: &str) -> Result<Module> {
5358
fn convert_module(
5459
name: &str,
5560
members: &[String],
56-
chunks_by_id: &HashMap<&String, &Chunk>,
61+
chunks_by_id: &HashMap<&str, &Chunk>,
62+
chunks_by_parent: &HashMap<&str, Vec<&Chunk>>,
5763
) -> Result<Module> {
64+
let (modules, classes, functions) = convert_members(
65+
&members
66+
.iter()
67+
.filter_map(|id| chunks_by_id.get(id.as_str()).copied())
68+
.collect::<Vec<_>>(),
69+
chunks_by_id,
70+
chunks_by_parent,
71+
)?;
72+
Ok(Module {
73+
name: name.into(),
74+
modules,
75+
classes,
76+
functions,
77+
})
78+
}
79+
80+
/// Convert a list of members of a module or a class
81+
fn convert_members(
82+
chunks: &[&Chunk],
83+
chunks_by_id: &HashMap<&str, &Chunk>,
84+
chunks_by_parent: &HashMap<&str, Vec<&Chunk>>,
85+
) -> Result<(Vec<Module>, Vec<Class>, Vec<Function>)> {
5886
let mut modules = Vec::new();
5987
let mut classes = Vec::new();
6088
let mut functions = Vec::new();
61-
for member in members {
62-
if let Some(chunk) = chunks_by_id.get(member) {
63-
match chunk {
64-
Chunk::Module {
89+
for chunk in chunks {
90+
match chunk {
91+
Chunk::Module {
92+
name,
93+
members,
94+
id: _,
95+
} => {
96+
modules.push(convert_module(
6597
name,
6698
members,
67-
id: _,
68-
} => {
69-
modules.push(convert_module(name, members, chunks_by_id)?);
70-
}
71-
Chunk::Class { name, id: _ } => classes.push(Class { name: name.into() }),
72-
Chunk::Function {
73-
name,
74-
id: _,
75-
arguments,
76-
} => functions.push(Function {
99+
chunks_by_id,
100+
chunks_by_parent,
101+
)?);
102+
}
103+
Chunk::Class { name, id } => {
104+
let (_, _, mut methods) = convert_members(
105+
chunks_by_parent
106+
.get(&id.as_str())
107+
.map(Vec::as_slice)
108+
.unwrap_or_default(),
109+
chunks_by_id,
110+
chunks_by_parent,
111+
)?;
112+
// We sort methods to get a stable output
113+
methods.sort_by(|l, r| match l.name.cmp(&r.name) {
114+
Ordering::Equal => {
115+
// We put the getter before the setter
116+
if l.decorators.iter().any(|d| d == "property") {
117+
Ordering::Less
118+
} else if r.decorators.iter().any(|d| d == "property") {
119+
Ordering::Greater
120+
} else {
121+
// We pick an ordering based on decorators
122+
l.decorators.cmp(&r.decorators)
123+
}
124+
}
125+
o => o,
126+
});
127+
classes.push(Class {
77128
name: name.into(),
78-
arguments: Arguments {
79-
positional_only_arguments: arguments
80-
.posonlyargs
81-
.iter()
82-
.map(convert_argument)
83-
.collect(),
84-
arguments: arguments.args.iter().map(convert_argument).collect(),
85-
vararg: arguments
86-
.vararg
87-
.as_ref()
88-
.map(convert_variable_length_argument),
89-
keyword_only_arguments: arguments
90-
.kwonlyargs
91-
.iter()
92-
.map(convert_argument)
93-
.collect(),
94-
kwarg: arguments
95-
.kwarg
96-
.as_ref()
97-
.map(convert_variable_length_argument),
98-
},
99-
}),
129+
methods,
130+
})
100131
}
132+
Chunk::Function {
133+
name,
134+
id: _,
135+
arguments,
136+
parent: _,
137+
decorators,
138+
} => functions.push(Function {
139+
name: name.into(),
140+
decorators: decorators.clone(),
141+
arguments: Arguments {
142+
positional_only_arguments: arguments
143+
.posonlyargs
144+
.iter()
145+
.map(convert_argument)
146+
.collect(),
147+
arguments: arguments.args.iter().map(convert_argument).collect(),
148+
vararg: arguments
149+
.vararg
150+
.as_ref()
151+
.map(convert_variable_length_argument),
152+
keyword_only_arguments: arguments
153+
.kwonlyargs
154+
.iter()
155+
.map(convert_argument)
156+
.collect(),
157+
kwarg: arguments
158+
.kwarg
159+
.as_ref()
160+
.map(convert_variable_length_argument),
161+
},
162+
}),
101163
}
102164
}
103-
Ok(Module {
104-
name: name.into(),
105-
modules,
106-
classes,
107-
functions,
108-
})
165+
Ok((modules, classes, functions))
109166
}
110167

111168
fn convert_argument(arg: &ChunkArgument) -> Argument {
@@ -290,9 +347,14 @@ enum Chunk {
290347
name: String,
291348
},
292349
Function {
293-
id: String,
350+
#[serde(default)]
351+
id: Option<String>,
294352
name: String,
295353
arguments: ChunkArguments,
354+
#[serde(default)]
355+
parent: Option<String>,
356+
#[serde(default)]
357+
decorators: Vec<String>,
296358
},
297359
}
298360

‎pyo3-introspection/src/model.rs

+3
Original file line numberDiff line numberDiff line change
@@ -9,11 +9,14 @@ pub struct Module {
99
#[derive(Debug, Eq, PartialEq, Clone, Hash)]
1010
pub struct Class {
1111
pub name: String,
12+
pub methods: Vec<Function>,
1213
}
1314

1415
#[derive(Debug, Eq, PartialEq, Clone, Hash)]
1516
pub struct Function {
1617
pub name: String,
18+
/// decorator like 'property' or 'staticmethod'
19+
pub decorators: Vec<String>,
1720
pub arguments: Arguments,
1821
}
1922

‎pyo3-introspection/src/stubs.rs

+44-4
Original file line numberDiff line numberDiff line change
@@ -39,12 +39,39 @@ fn module_stubs(module: &Module) -> String {
3939
for function in &module.functions {
4040
elements.push(function_stubs(function));
4141
}
42-
elements.push(String::new()); // last line jump
43-
elements.join("\n")
42+
43+
// We insert two line jumps (i.e. empty strings) only above and below multiple line elements (classes with methods, functions with decorators)
44+
let mut output = String::new();
45+
for element in elements {
46+
let is_multiline = element.contains('\n');
47+
if is_multiline && !output.is_empty() && !output.ends_with("\n\n") {
48+
output.push('\n');
49+
}
50+
output.push_str(&element);
51+
output.push('\n');
52+
if is_multiline {
53+
output.push('\n');
54+
}
55+
}
56+
// We remove a line jump at the end if they are two
57+
if output.ends_with("\n\n") {
58+
output.pop();
59+
}
60+
output
4461
}
4562

4663
fn class_stubs(class: &Class) -> String {
47-
format!("class {}: ...", class.name)
64+
let mut buffer = format!("class {}:", class.name);
65+
if class.methods.is_empty() {
66+
buffer.push_str(" ...");
67+
return buffer;
68+
}
69+
for method in &class.methods {
70+
// We do the indentation
71+
buffer.push_str("\n ");
72+
buffer.push_str(&function_stubs(method).replace('\n', "\n "));
73+
}
74+
buffer
4875
}
4976

5077
fn function_stubs(function: &Function) -> String {
@@ -70,7 +97,18 @@ fn function_stubs(function: &Function) -> String {
7097
if let Some(argument) = &function.arguments.kwarg {
7198
parameters.push(format!("**{}", variable_length_argument_stub(argument)));
7299
}
73-
format!("def {}({}): ...", function.name, parameters.join(", "))
100+
let output = format!("def {}({}): ...", function.name, parameters.join(", "));
101+
if function.decorators.is_empty() {
102+
return output;
103+
}
104+
let mut buffer = String::new();
105+
for decorator in &function.decorators {
106+
buffer.push('@');
107+
buffer.push_str(decorator);
108+
buffer.push('\n');
109+
}
110+
buffer.push_str(&output);
111+
buffer
74112
}
75113

76114
fn argument_stub(argument: &Argument) -> String {
@@ -95,6 +133,7 @@ mod tests {
95133
fn function_stubs_with_variable_length() {
96134
let function = Function {
97135
name: "func".into(),
136+
decorators: Vec::new(),
98137
arguments: Arguments {
99138
positional_only_arguments: vec![Argument {
100139
name: "posonly".into(),
@@ -126,6 +165,7 @@ mod tests {
126165
fn function_stubs_without_variable_length() {
127166
let function = Function {
128167
name: "afunc".into(),
168+
decorators: Vec::new(),
129169
arguments: Arguments {
130170
positional_only_arguments: vec![Argument {
131171
name: "posonly".into(),

0 commit comments

Comments
 (0)