Skip to main content

argyph_graph/resolve/
python.rs

1use camino::Utf8Path;
2
3use super::{ImportResolver, ModuleTarget};
4use crate::edge::Confidence;
5
6pub struct PythonResolver;
7
8impl ImportResolver for PythonResolver {
9    fn resolve_import(
10        &self,
11        source_file: &Utf8Path,
12        _module_path: &[String],
13        raw: &str,
14    ) -> Option<ModuleTarget> {
15        let trimmed = raw.trim();
16        if !trimmed.starts_with("from .") {
17            return None;
18        }
19
20        let after_from = trimmed.strip_prefix("from ")?;
21        let module_str = after_from.split(" import ").next()?.trim();
22
23        let source_dir = source_file.parent()?;
24        let mut file = source_dir.to_path_buf();
25
26        let depth = module_str.chars().take_while(|c| *c == '.').count();
27        for _ in 0..depth.saturating_sub(1) {
28            if !file.pop() {}
29        }
30
31        let remaining = &module_str[depth..];
32        if remaining.is_empty() {
33            file.push("__init__");
34        } else {
35            for seg in remaining.split('.') {
36                if !seg.is_empty() {
37                    file.push(seg);
38                }
39            }
40        }
41
42        let candidate = format!("{file}.py");
43        Some(ModuleTarget {
44            file_path: super::normalize_path(&candidate),
45            confidence: Confidence::Heuristic,
46        })
47    }
48}
49
50#[cfg(test)]
51#[allow(clippy::unwrap_used, clippy::expect_used)]
52mod tests {
53    use super::*;
54    use camino::Utf8Path;
55
56    #[test]
57    fn resolve_relative_dot_import() {
58        let resolver = PythonResolver;
59        let path: Vec<String> = ["math"].iter().map(|s| s.to_string()).collect();
60        let tgt = resolver
61            .resolve_import(Utf8Path::new("src/main.py"), &path, "from .math import add")
62            .expect("should resolve from .math import add");
63
64        assert!(tgt.file_path.contains("src/math"), "got {}", tgt.file_path);
65        assert!(tgt.file_path.ends_with(".py"));
66    }
67
68    #[test]
69    fn resolve_parent_import() {
70        let resolver = PythonResolver;
71        let tgt = resolver
72            .resolve_import(
73                Utf8Path::new("src/subpkg/module.py"),
74                &[],
75                "from ..types import User",
76            )
77            .expect("should resolve from ..types import User");
78
79        assert!(tgt.file_path.contains("src/types"), "got {}", tgt.file_path);
80    }
81
82    #[test]
83    fn bare_import_returns_none() {
84        let resolver = PythonResolver;
85        let path: Vec<String> = ["os"].iter().map(|s| s.to_string()).collect();
86        assert!(resolver
87            .resolve_import(Utf8Path::new("src/main.py"), &path, "import os",)
88            .is_none());
89    }
90}