1
use std::collections::HashSet;
2
 
3
use ruff_python_ast::Stmt;
4
use ruff_python_ast::StmtImportFrom;
5
use ruff_python_parser::parse_module;
6
 
7
#[cxx::bridge(namespace = "edgeworker::rust::python_parser")]
8
mod ffi {
9
 
10
    extern "Rust" {
11
        fn get_imports(sources: &[&str]) -> Vec<String>;
12
    }
13
}
14
 
15
#[must_use]
16
pub fn get_imports(sources: &[&str]) -> Vec<String> {
17
    let mut names: HashSet<String> = HashSet::new();
18
    for src in sources {
19
        // Just skip it if it doesn't parse.
20
        let Ok(module) = parse_module(src) else {
21
            continue;
22
        };
23
        for stmt in &module.syntax().body {
24
            match stmt {
25
                Stmt::Import(s) => {
26
                    names.extend(s.names.iter().map(|x| x.name.id.as_str().to_owned()));
27
                }
28
                Stmt::ImportFrom(StmtImportFrom {
29
                    module: Some(module),
30
                    level: 0,
31
                    ..
32
                }) => {
33
                    names.insert(module.id.as_str().to_owned());
34
                }
35
                _ => {}
36
            }
37
        }
38
    }
39
    let mut result: Vec<_> = names.drain().collect();
40
    result.sort();
41
    result
42
}