pyo3_introspection/
stubs.rs

1use crate::model::{Argument, Class, Const, Function, Module, VariableLengthArgument};
2use std::collections::{BTreeSet, HashMap};
3use std::path::{Path, PathBuf};
4
5/// Generates the [type stubs](https://typing.readthedocs.io/en/latest/source/stubs.html) of a given module.
6/// It returns a map between the file name and the file content.
7/// The root module stubs will be in the `__init__.pyi` file and the submodules directory
8/// in files with a relevant name.
9pub fn module_stub_files(module: &Module) -> HashMap<PathBuf, String> {
10    let mut output_files = HashMap::new();
11    add_module_stub_files(module, Path::new(""), &mut output_files);
12    output_files
13}
14
15fn add_module_stub_files(
16    module: &Module,
17    module_path: &Path,
18    output_files: &mut HashMap<PathBuf, String>,
19) {
20    output_files.insert(module_path.join("__init__.pyi"), module_stubs(module));
21    for submodule in &module.modules {
22        if submodule.modules.is_empty() {
23            output_files.insert(
24                module_path.join(format!("{}.pyi", submodule.name)),
25                module_stubs(submodule),
26            );
27        } else {
28            add_module_stub_files(submodule, &module_path.join(&submodule.name), output_files);
29        }
30    }
31}
32
33/// Generates the module stubs to a String, not including submodules
34fn module_stubs(module: &Module) -> String {
35    let mut modules_to_import = BTreeSet::new();
36    let mut elements = Vec::new();
37    for konst in &module.consts {
38        elements.push(const_stubs(konst, &mut modules_to_import));
39    }
40    for class in &module.classes {
41        elements.push(class_stubs(class, &mut modules_to_import));
42    }
43    for function in &module.functions {
44        elements.push(function_stubs(function, &mut modules_to_import));
45    }
46    let mut final_elements = Vec::new();
47    for module_to_import in &modules_to_import {
48        final_elements.push(format!("import {module_to_import}"));
49    }
50    final_elements.extend(elements);
51
52    let mut output = String::new();
53
54    // We insert two line jumps (i.e. empty strings) only above and below multiple line elements (classes with methods, functions with decorators)
55    for element in final_elements {
56        let is_multiline = element.contains('\n');
57        if is_multiline && !output.is_empty() && !output.ends_with("\n\n") {
58            output.push('\n');
59        }
60        output.push_str(&element);
61        output.push('\n');
62        if is_multiline {
63            output.push('\n');
64        }
65    }
66
67    // We remove a line jump at the end if they are two
68    if output.ends_with("\n\n") {
69        output.pop();
70    }
71    output
72}
73
74fn class_stubs(class: &Class, modules_to_import: &mut BTreeSet<String>) -> String {
75    let mut buffer = format!("class {}:", class.name);
76    if class.methods.is_empty() {
77        buffer.push_str(" ...");
78        return buffer;
79    }
80    for method in &class.methods {
81        // We do the indentation
82        buffer.push_str("\n    ");
83        buffer.push_str(&function_stubs(method, modules_to_import).replace('\n', "\n    "));
84    }
85    buffer
86}
87
88fn function_stubs(function: &Function, modules_to_import: &mut BTreeSet<String>) -> String {
89    // Signature
90    let mut parameters = Vec::new();
91    for argument in &function.arguments.positional_only_arguments {
92        parameters.push(argument_stub(argument, modules_to_import));
93    }
94    if !function.arguments.positional_only_arguments.is_empty() {
95        parameters.push("/".into());
96    }
97    for argument in &function.arguments.arguments {
98        parameters.push(argument_stub(argument, modules_to_import));
99    }
100    if let Some(argument) = &function.arguments.vararg {
101        parameters.push(format!("*{}", variable_length_argument_stub(argument)));
102    } else if !function.arguments.keyword_only_arguments.is_empty() {
103        parameters.push("*".into());
104    }
105    for argument in &function.arguments.keyword_only_arguments {
106        parameters.push(argument_stub(argument, modules_to_import));
107    }
108    if let Some(argument) = &function.arguments.kwarg {
109        parameters.push(format!("**{}", variable_length_argument_stub(argument)));
110    }
111    let mut buffer = String::new();
112    for decorator in &function.decorators {
113        buffer.push('@');
114        buffer.push_str(decorator);
115        buffer.push('\n');
116    }
117    buffer.push_str("def ");
118    buffer.push_str(&function.name);
119    buffer.push('(');
120    buffer.push_str(&parameters.join(", "));
121    buffer.push(')');
122    if let Some(returns) = &function.returns {
123        buffer.push_str(" -> ");
124        buffer.push_str(annotation_stub(returns, modules_to_import));
125    }
126    buffer.push_str(": ...");
127    buffer
128}
129
130fn const_stubs(konst: &Const, modules_to_import: &mut BTreeSet<String>) -> String {
131    modules_to_import.insert("typing".to_string());
132    let Const { name, value } = konst;
133    format!("{name}: typing.Final = {value}")
134}
135
136fn argument_stub(argument: &Argument, modules_to_import: &mut BTreeSet<String>) -> String {
137    let mut output = argument.name.clone();
138    if let Some(annotation) = &argument.annotation {
139        output.push_str(": ");
140        output.push_str(annotation_stub(annotation, modules_to_import));
141    }
142    if let Some(default_value) = &argument.default_value {
143        output.push_str(if argument.annotation.is_some() {
144            " = "
145        } else {
146            "="
147        });
148        output.push_str(default_value);
149    }
150    output
151}
152
153fn variable_length_argument_stub(argument: &VariableLengthArgument) -> String {
154    argument.name.clone()
155}
156
157fn annotation_stub<'a>(annotation: &'a str, modules_to_import: &mut BTreeSet<String>) -> &'a str {
158    if let Some((module, _)) = annotation.rsplit_once('.') {
159        // TODO: this is very naive
160        modules_to_import.insert(module.into());
161    }
162    annotation
163}
164
165#[cfg(test)]
166mod tests {
167    use super::*;
168    use crate::model::Arguments;
169
170    #[test]
171    fn function_stubs_with_variable_length() {
172        let function = Function {
173            name: "func".into(),
174            decorators: Vec::new(),
175            arguments: Arguments {
176                positional_only_arguments: vec![Argument {
177                    name: "posonly".into(),
178                    default_value: None,
179                    annotation: None,
180                }],
181                arguments: vec![Argument {
182                    name: "arg".into(),
183                    default_value: None,
184                    annotation: None,
185                }],
186                vararg: Some(VariableLengthArgument {
187                    name: "varargs".into(),
188                }),
189                keyword_only_arguments: vec![Argument {
190                    name: "karg".into(),
191                    default_value: None,
192                    annotation: Some("str".into()),
193                }],
194                kwarg: Some(VariableLengthArgument {
195                    name: "kwarg".into(),
196                }),
197            },
198            returns: Some("list[str]".into()),
199        };
200        assert_eq!(
201            "def func(posonly, /, arg, *varargs, karg: str, **kwarg) -> list[str]: ...",
202            function_stubs(&function, &mut BTreeSet::new())
203        )
204    }
205
206    #[test]
207    fn function_stubs_without_variable_length() {
208        let function = Function {
209            name: "afunc".into(),
210            decorators: Vec::new(),
211            arguments: Arguments {
212                positional_only_arguments: vec![Argument {
213                    name: "posonly".into(),
214                    default_value: Some("1".into()),
215                    annotation: None,
216                }],
217                arguments: vec![Argument {
218                    name: "arg".into(),
219                    default_value: Some("True".into()),
220                    annotation: None,
221                }],
222                vararg: None,
223                keyword_only_arguments: vec![Argument {
224                    name: "karg".into(),
225                    default_value: Some("\"foo\"".into()),
226                    annotation: Some("str".into()),
227                }],
228                kwarg: None,
229            },
230            returns: None,
231        };
232        assert_eq!(
233            "def afunc(posonly=1, /, arg=True, *, karg: str = \"foo\"): ...",
234            function_stubs(&function, &mut BTreeSet::new())
235        )
236    }
237}
pyo3_introspection/stubs.rs

pyo3_introspection/
stubs.rs
