use super::{FileSet, ImportTarget, normalize_path};
use crate::types::ImportStatement;
pub fn resolve_py_import(
import: &ImportStatement,
from_file: usize,
files: &FileSet,
) -> ImportTarget {
let module_path = import.module_path.trim();
if module_path.is_empty() {
return ImportTarget::Unresolved(import.module_path.clone());
}
if module_path.starts_with('.') {
resolve_relative(module_path, from_file, files)
.map(ImportTarget::File)
.unwrap_or_else(|| ImportTarget::Unresolved(import.module_path.clone()))
} else if !module_path.contains('.') && !import.is_external {
resolve_relative(&format!(".{module_path}"), from_file, files)
.map(ImportTarget::File)
.unwrap_or_else(|| ImportTarget::Unresolved(import.module_path.clone()))
} else {
resolve_absolute(module_path, files)
}
}
fn resolve_relative(module_path: &str, from_file: usize, files: &FileSet) -> Option<usize> {
let level = module_path.chars().take_while(|&c| c == '.').count();
let remainder = &module_path[level..];
let dir = files.dir_of(from_file)?;
let ups = "../".repeat(level - 1);
let seg_path: String = remainder
.split('.')
.filter(|s| !s.is_empty())
.collect::<Vec<_>>()
.join("/");
let candidates: Vec<String> = if seg_path.is_empty() {
vec![format!("{ups}__init__.py")]
} else {
vec![
format!("{ups}{seg_path}.py"),
format!("{ups}{seg_path}/__init__.py"),
]
};
candidates.iter().find_map(|rel| files.probe(&dir, rel))
}
fn resolve_absolute(module_path: &str, files: &FileSet) -> ImportTarget {
let segs: Vec<&str> = module_path.split('.').filter(|s| !s.is_empty()).collect();
if segs.is_empty() {
return ImportTarget::External(module_path.to_string());
}
let mut n = segs.len();
while n >= 2 {
match match_dotted(&segs[..n], module_path, files) {
ImportTarget::External(_) => n -= 1,
hit => return hit,
}
}
match_dotted(&segs[..1], module_path, files)
}
fn match_dotted(segs: &[&str], module_path: &str, files: &FileSet) -> ImportTarget {
let seg_path = segs.join("/");
let module_suffix = format!("{seg_path}.py");
let package_suffix = format!("{seg_path}/__init__.py");
let mut matches: Vec<usize> = Vec::new();
for idx in 0..files.len() {
let path = match files.path(idx) {
Some(p) => normalize_path(p),
None => continue,
};
if path_ends_with_segment(&path, &module_suffix)
|| path_ends_with_segment(&path, &package_suffix)
{
matches.push(idx);
}
}
match matches.as_slice() {
[] => ImportTarget::External(module_path.to_string()),
[only] => ImportTarget::File(*only),
_ => ImportTarget::Unresolved(module_path.to_string()),
}
}
fn path_ends_with_segment(path: &str, suffix: &str) -> bool {
path == suffix || path.ends_with(&format!("/{suffix}"))
}
#[cfg(test)]
mod tests {
use super::*;
use crate::storage::graph::FileSet;
use crate::types::{ImportStatement, TreeNode};
fn py(path: &str) -> TreeNode {
TreeNode::new(path.to_string(), "python".to_string())
}
fn imp(module_path: &str, from_path: &str, is_external: bool) -> ImportStatement {
ImportStatement::new(module_path.to_string(), from_path.to_string())
.with_external(is_external)
}
#[test]
fn absolute_dotted_resolves_to_module_file() {
let files = vec![py("/repo/main.py"), py("/repo/pkg/mod.py")];
let fs = FileSet::new(&files);
let statement = imp("pkg.mod", "/repo/main.py", true);
assert_eq!(resolve_py_import(&statement, 0, &fs), ImportTarget::File(1));
}
#[test]
fn package_import_resolves_via_init() {
let files = vec![py("/repo/app.py"), py("/repo/pkg/__init__.py")];
let fs = FileSet::new(&files);
let statement = imp("pkg", "/repo/app.py", true);
assert_eq!(resolve_py_import(&statement, 0, &fs), ImportTarget::File(1));
}
#[test]
fn relative_level_one_resolves_sibling() {
let files = vec![py("/repo/pkg/sub/mod.py"), py("/repo/pkg/sub/x.py")];
let fs = FileSet::new(&files);
let statement = imp(".x", "/repo/pkg/sub/mod.py", false);
assert_eq!(resolve_py_import(&statement, 0, &fs), ImportTarget::File(1));
}
#[test]
fn relative_level_two_resolves_up() {
let files = vec![py("/repo/proj/sub/mod.py"), py("/repo/proj/pkg/y.py")];
let fs = FileSet::new(&files);
let statement = imp("..pkg.y", "/repo/proj/sub/mod.py", false);
assert_eq!(resolve_py_import(&statement, 0, &fs), ImportTarget::File(1));
}
#[test]
fn third_party_with_no_match_is_external() {
let files = vec![py("/repo/app.py")];
let fs = FileSet::new(&files);
let statement = imp("numpy", "/repo/app.py", true);
assert_eq!(
resolve_py_import(&statement, 0, &fs),
ImportTarget::External("numpy".to_string())
);
}
#[test]
fn unresolvable_relative_is_unresolved_never_wrong_file() {
let files = vec![py("/repo/proj/sub/mod.py"), py("/repo/proj/pkg/y.py")];
let fs = FileSet::new(&files);
let statement = imp("..missing.thing", "/repo/proj/sub/mod.py", false);
assert_eq!(
resolve_py_import(&statement, 0, &fs),
ImportTarget::Unresolved("..missing.thing".to_string())
);
}
#[test]
fn ambiguous_absolute_match_is_unresolved() {
let files = vec![
py("/repo/a/pkg/mod.py"),
py("/repo/b/pkg/mod.py"),
py("/repo/main.py"),
];
let fs = FileSet::new(&files);
let statement = imp("pkg.mod", "/repo/main.py", true);
assert_eq!(
resolve_py_import(&statement, 2, &fs),
ImportTarget::Unresolved("pkg.mod".to_string())
);
}
#[test]
fn trailing_symbol_segment_is_trimmed_to_the_module() {
let files = vec![py("/repo/pkg/mod_a.py"), py("/repo/pkg/mod_b.py")];
let fs = FileSet::new(&files);
let statement = imp("pkg.mod_b.helper", "/repo/pkg/mod_a.py", false);
assert_eq!(resolve_py_import(&statement, 0, &fs), ImportTarget::File(1));
}
#[test]
fn trimming_does_not_invent_a_file_for_stdlib() {
let files = vec![py("/repo/app.py")];
let fs = FileSet::new(&files);
let statement = imp("typing.List", "/repo/app.py", true);
assert_eq!(
resolve_py_import(&statement, 0, &fs),
ImportTarget::External("typing.List".to_string())
);
}
#[test]
fn dot_stripped_relative_never_matches_an_unrelated_file() {
let files = vec![
py("/repo/pkg/mod_a.py"),
py("/repo/pkg/mod_b.py"),
py("/repo/other/helper.py"),
];
let fs = FileSet::new(&files);
let statement = imp("helper", "/repo/pkg/mod_a.py", false);
assert_eq!(
resolve_py_import(&statement, 0, &fs),
ImportTarget::Unresolved("helper".to_string())
);
}
#[test]
fn dot_stripped_relative_resolves_against_own_package() {
let files = vec![py("/repo/pkg/mod_a.py"), py("/repo/pkg/sibling.py")];
let fs = FileSet::new(&files);
let statement = imp("sibling", "/repo/pkg/mod_a.py", false);
assert_eq!(resolve_py_import(&statement, 0, &fs), ImportTarget::File(1));
}
#[test]
fn suffix_match_is_whole_segment() {
let files = vec![py("/repo/mypkg/mod.py"), py("/repo/main.py")];
let fs = FileSet::new(&files);
let statement = imp("pkg.mod", "/repo/main.py", true);
assert_eq!(
resolve_py_import(&statement, 1, &fs),
ImportTarget::External("pkg.mod".to_string())
);
}
}