monty_type_checking/
imports.rs

1//! Scan snippet imports for later type checks and host stub requests.
2
3use std::iter::once;
4
5use ruff_python_ast::{Alias, ExceptHandler, Stmt};
6use ruff_python_parser::parse_module;
7use ruff_text_size::{Ranged, TextRange};
8
9/// Imports collected by [`scan_imports`].
10#[derive(Debug, Default, PartialEq, Eq)]
11pub struct SnippetImports {
12    /// Import statements with aliases, one per line, injected before later checks' `.pyi` star import.
13    /// This preserves bindings such as `import math`, which `.pyi` files do not re-export.
14    pub prelude: String,
15    /// Each top-level module imported, once, in source order.
16    pub modules: Vec<ImportedModule>,
17}
18
19/// A top-level module a snippet imports (`a` for `import a.b` and `from a.b import c`).
20#[derive(Debug, PartialEq, Eq)]
21pub struct ImportedModule {
22    pub name: String,
23    /// The first import statement naming it.
24    pub range: TextRange,
25}
26
27/// Collects imports outside function and class bodies, excluding `__future__` and relative imports.
28/// Invalid syntax yields no imports. All module-level branches are included because the checker
29/// cannot tell which ran, and any names bound at runtime must resolve.
30#[must_use]
31pub fn scan_imports(source: &str) -> SnippetImports {
32    let Ok(parsed) = parse_module(source) else {
33        return SnippetImports::default();
34    };
35    let mut imports = SnippetImports::default();
36    collect(&parsed.syntax().body, &mut imports);
37    imports
38}
39
40impl SnippetImports {
41    /// Appends a prelude line and records each module's first import.
42    fn add(&mut self, line: &str, modules: impl Iterator<Item = (String, TextRange)>) {
43        self.prelude.push_str(line);
44        self.prelude.push('\n');
45        for (name, range) in modules {
46            if !self.modules.iter().any(|module| module.name == name) {
47                self.modules.push(ImportedModule { name, range });
48            }
49        }
50    }
51}
52
53/// Descends into compound statements, excluding function and class bodies whose imports bind locally.
54fn collect(body: &[Stmt], imports: &mut SnippetImports) {
55    for statement in body {
56        match statement {
57            Stmt::Import(import) => {
58                let range = import.range();
59                imports.add(
60                    &format!("import {}", aliases(&import.names)),
61                    import.names.iter().map(|alias| (top_level(alias.name.as_str()), range)),
62                );
63            }
64            Stmt::ImportFrom(import) if import.level == 0 => {
65                if let Some(module) = &import.module
66                    && module.as_str() != "__future__"
67                {
68                    imports.add(
69                        &format!("from {} import {}", module.as_str(), aliases(&import.names)),
70                        once((top_level(module.as_str()), import.range())),
71                    );
72                }
73            }
74            Stmt::If(node) => {
75                collect(&node.body, imports);
76                for clause in &node.elif_else_clauses {
77                    collect(&clause.body, imports);
78                }
79            }
80            Stmt::For(node) => {
81                collect(&node.body, imports);
82                collect(&node.orelse, imports);
83            }
84            Stmt::While(node) => {
85                collect(&node.body, imports);
86                collect(&node.orelse, imports);
87            }
88            Stmt::With(node) => collect(&node.body, imports),
89            Stmt::Try(node) => {
90                collect(&node.body, imports);
91                for ExceptHandler::ExceptHandler(handler) in &node.handlers {
92                    collect(&handler.body, imports);
93                }
94                collect(&node.orelse, imports);
95                collect(&node.finalbody, imports);
96            }
97            Stmt::Match(node) => {
98                for case in &node.cases {
99                    collect(&case.body, imports);
100                }
101            }
102            _ => {}
103        }
104    }
105}
106
107fn top_level(dotted: &str) -> String {
108    dotted.split('.').next().unwrap_or_default().to_owned()
109}
110
111fn aliases(names: &[Alias]) -> String {
112    names
113        .iter()
114        .map(|alias| match &alias.asname {
115            Some(asname) => format!("{} as {}", alias.name.as_str(), asname.as_str()),
116            None => alias.name.as_str().to_owned(),
117        })
118        .collect::<Vec<_>>()
119        .join(", ")
120}
monty_type_checking/imports.rs

monty_type_checking/
imports.rs
