pyo3_introspection/
stubs.rs

1use crate::model::{Argument, Class, Function, Module, VariableLengthArgument};
2use std::collections::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 elements = Vec::new();
36    for class in &module.classes {
37        elements.push(class_stubs(class));
38    }
39    for function in &module.functions {
40        elements.push(function_stubs(function));
41    }
42    elements.push(String::new()); // last line jump
43    elements.join("\n")
44}
45
46fn class_stubs(class: &Class) -> String {
47    format!("class {}: ...", class.name)
48}
49
50fn function_stubs(function: &Function) -> String {
51    // Signature
52    let mut parameters = Vec::new();
53    for argument in &function.arguments.positional_only_arguments {
54        parameters.push(argument_stub(argument));
55    }
56    if !function.arguments.positional_only_arguments.is_empty() {
57        parameters.push("/".into());
58    }
59    for argument in &function.arguments.arguments {
60        parameters.push(argument_stub(argument));
61    }
62    if let Some(argument) = &function.arguments.vararg {
63        parameters.push(format!("*{}", variable_length_argument_stub(argument)));
64    } else if !function.arguments.keyword_only_arguments.is_empty() {
65        parameters.push("*".into());
66    }
67    for argument in &function.arguments.keyword_only_arguments {
68        parameters.push(argument_stub(argument));
69    }
70    if let Some(argument) = &function.arguments.kwarg {
71        parameters.push(format!("**{}", variable_length_argument_stub(argument)));
72    }
73    format!("def {}({}): ...", function.name, parameters.join(", "))
74}
75
76fn argument_stub(argument: &Argument) -> String {
77    let mut output = argument.name.clone();
78    if let Some(default_value) = &argument.default_value {
79        output.push('=');
80        output.push_str(default_value);
81    }
82    output
83}
84
85fn variable_length_argument_stub(argument: &VariableLengthArgument) -> String {
86    argument.name.clone()
87}
88
89#[cfg(test)]
90mod tests {
91    use super::*;
92    use crate::model::Arguments;
93
94    #[test]
95    fn function_stubs_with_variable_length() {
96        let function = Function {
97            name: "func".into(),
98            arguments: Arguments {
99                positional_only_arguments: vec![Argument {
100                    name: "posonly".into(),
101                    default_value: None,
102                }],
103                arguments: vec![Argument {
104                    name: "arg".into(),
105                    default_value: None,
106                }],
107                vararg: Some(VariableLengthArgument {
108                    name: "varargs".into(),
109                }),
110                keyword_only_arguments: vec![Argument {
111                    name: "karg".into(),
112                    default_value: None,
113                }],
114                kwarg: Some(VariableLengthArgument {
115                    name: "kwarg".into(),
116                }),
117            },
118        };
119        assert_eq!(
120            "def func(posonly, /, arg, *varargs, karg, **kwarg): ...",
121            function_stubs(&function)
122        )
123    }
124
125    #[test]
126    fn function_stubs_without_variable_length() {
127        let function = Function {
128            name: "afunc".into(),
129            arguments: Arguments {
130                positional_only_arguments: vec![Argument {
131                    name: "posonly".into(),
132                    default_value: Some("1".into()),
133                }],
134                arguments: vec![Argument {
135                    name: "arg".into(),
136                    default_value: Some("True".into()),
137                }],
138                vararg: None,
139                keyword_only_arguments: vec![Argument {
140                    name: "karg".into(),
141                    default_value: Some("\"foo\"".into()),
142                }],
143                kwarg: None,
144            },
145        };
146        assert_eq!(
147            "def afunc(posonly=1, /, arg=True, *, karg=\"foo\"): ...",
148            function_stubs(&function)
149        )
150    }
151}
pyo3_introspection/stubs.rs

pyo3_introspection/
stubs.rs
