pyo3_introspection/
stubs.rs

1use crate::model::{Class, Function, Module};
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    format!("def {}(*args, **kwargs): ...", function.name)
52}
pyo3_introspection/stubs.rs

pyo3_introspection/
stubs.rs
