fxrank_lang_python/
module_map.rs

1//! Python module map: dotted module keys via `__init__.py` package roots, and
2//! absolute/relative import resolution against the in-batch set, by path
3//! convention (spec 025-3e §5.3). No disk, no sys.path, no libcst.
4
5use std::collections::HashSet;
6
7use fxrank_core::frontend::SourceFile;
8
9pub struct PyModuleMap {
10    keys: HashSet<Vec<String>>,
11    // dir paths (with trailing '/') that contain an __init__.py in the batch.
12    pkg_dirs: HashSet<String>,
13}
14
15impl PyModuleMap {
16    pub fn build(files: &[SourceFile]) -> Self {
17        let mut pkg_dirs = HashSet::new();
18        for f in files {
19            if f.path.ends_with("/__init__.py") || f.path == "__init__.py" {
20                pkg_dirs.insert(dir_of(&f.path));
21            }
22        }
23        let mut keys = HashSet::new();
24        for f in files {
25            if !f.path.ends_with(".py") {
26                continue;
27            }
28            if let Some(k) = dotted_key(&f.path, &pkg_dirs) {
29                keys.insert(k);
30            }
31        }
32        Self { keys, pkg_dirs }
33    }
34
35    pub fn module_of(&self, file_path: &str) -> Option<Vec<String>> {
36        if !file_path.ends_with(".py") {
37            return None;
38        }
39        dotted_key(file_path, &self.pkg_dirs)
40    }
41
42    /// True when the file is a package `__init__.py` (its module key IS its package).
43    pub fn is_package(&self, file_path: &str) -> bool {
44        file_path.ends_with("/__init__.py") || file_path == "__init__.py"
45    }
46
47    pub fn resolve_absolute(&self, dotted: &str) -> Option<Vec<String>> {
48        let segs: Vec<String> = dotted.split('.').map(|s| s.to_string()).collect();
49        if self.keys.contains(&segs) {
50            Some(segs)
51        } else {
52            None
53        }
54    }
55
56    /// Resolve a relative import. The relative anchor is the importing module's
57    /// PACKAGE: the key itself when the importer is a package `__init__.py`
58    /// (`is_package`), else the key minus its module stem. `level` dots then walk
59    /// up `level-1` more packages from that anchor (Python: level 1 = the package
60    /// containing the importer). This `is_package` distinction is REQUIRED — a key
61    /// like `["pkg","sub"]` is ambiguous (regular module `pkg/sub.py` vs package
62    /// `pkg/sub/__init__.py`) and the two anchor differently.
63    pub fn resolve_relative(
64        &self,
65        referencing: &[String],
66        is_package: bool,
67        level: usize,
68        suffix: &str,
69    ) -> Option<Vec<String>> {
70        if level == 0 {
71            return None; // not a relative import
72        }
73        let anchor: Vec<String> = if is_package {
74            referencing.to_vec()
75        } else if referencing.is_empty() {
76            return None;
77        } else {
78            referencing[..referencing.len() - 1].to_vec()
79        };
80        // A relative import REQUIRES a containing package. An empty anchor means
81        // the referencing module has no parent package (a top-level `top.py`, or a
82        // file under a non-`__init__` dir) — Python errors here ("no known parent
83        // package"), so we must NOT resolve against a root-level module. (P2, round 3)
84        if anchor.is_empty() {
85            return None;
86        }
87        let up = level - 1; // level 1 = the anchor package itself
88        if up > anchor.len() {
89            return None; // escaped above the top package
90        }
91        let mut target: Vec<String> = anchor[..anchor.len() - up].to_vec();
92        if !suffix.is_empty() {
93            target.extend(suffix.split('.').map(|s| s.to_string()));
94        }
95        if self.keys.contains(&target) {
96            Some(target)
97        } else {
98            None
99        }
100    }
101}
102
103/// Directory of a path, WITH trailing '/'. `"pkg/sub/mod.py"` → `"pkg/sub/"`.
104fn dir_of(path: &str) -> String {
105    match path.rfind('/') {
106        Some(i) => path[..=i].to_string(),
107        None => String::new(),
108    }
109}
110
111/// Dotted module key for a `.py` file: walk up from its dir while each dir is a
112/// package (has `__init__.py` in the batch); the outermost non-package dir is the
113/// root (excluded). An `__init__.py` keys to its package (no `__init__` segment).
114fn dotted_key(path: &str, pkg_dirs: &HashSet<String>) -> Option<Vec<String>> {
115    let stem = path.strip_suffix(".py")?;
116    // Split into directory segments + file stem.
117    let (dir_part, file_stem) = match stem.rfind('/') {
118        Some(i) => (&stem[..i], &stem[i + 1..]),
119        None => ("", stem),
120    };
121    // Collect the package segments: starting at the file's dir, walk up while the
122    // dir is a package. Build the dir prefix incrementally to test membership.
123    let dir_segs: Vec<&str> = if dir_part.is_empty() {
124        Vec::new()
125    } else {
126        dir_part.split('/').collect()
127    };
128    // Find the deepest ancestor index that is NOT a package → everything below it is the module path.
129    let mut first_pkg = dir_segs.len(); // index of the first package dir from the left
130    for i in (0..dir_segs.len()).rev() {
131        let prefix = format!("{}/", dir_segs[..=i].join("/"));
132        if pkg_dirs.contains(&prefix) {
133            first_pkg = i;
134        } else {
135            break;
136        }
137    }
138    let mut segs: Vec<String> = dir_segs[first_pkg..]
139        .iter()
140        .map(|s| s.to_string())
141        .collect();
142    if file_stem != "__init__" {
143        segs.push(file_stem.to_string());
144    }
145    Some(segs)
146}
147
148#[cfg(test)]
149mod tests {
150    use super::*;
151    use fxrank_core::frontend::SourceFile;
152    fn sf(p: &str) -> SourceFile {
153        SourceFile {
154            path: p.into(),
155            text: String::new(),
156        }
157    }
158
159    fn batch() -> Vec<SourceFile> {
160        vec![
161            sf("pkg/__init__.py"),
162            sf("pkg/sub/__init__.py"),
163            sf("pkg/sub/mod.py"),
164            sf("pkg/util.py"),
165            sf("top.py"), // no __init__.py sibling → top-level module
166        ]
167    }
168
169    #[test]
170    fn module_key_via_init_packages() {
171        let m = PyModuleMap::build(&batch());
172        assert_eq!(
173            m.module_of("pkg/sub/mod.py"),
174            Some(vec!["pkg".into(), "sub".into(), "mod".into()])
175        );
176        assert_eq!(
177            m.module_of("pkg/util.py"),
178            Some(vec!["pkg".into(), "util".into()])
179        );
180        assert_eq!(
181            m.module_of("pkg/sub/__init__.py"),
182            Some(vec!["pkg".into(), "sub".into()])
183        );
184        assert_eq!(m.module_of("top.py"), Some(vec!["top".into()]));
185    }
186
187    #[test]
188    fn resolve_absolute_in_batch_only() {
189        let m = PyModuleMap::build(&batch());
190        assert_eq!(
191            m.resolve_absolute("pkg.sub.mod"),
192            Some(vec!["pkg".into(), "sub".into(), "mod".into()])
193        );
194        assert_eq!(
195            m.resolve_absolute("pkg.util"),
196            Some(vec!["pkg".into(), "util".into()])
197        );
198        assert_eq!(m.resolve_absolute("os.path"), None); // stdlib, not in batch
199        assert_eq!(m.resolve_absolute("pkg.missing"), None);
200    }
201
202    #[test]
203    fn resolve_relative_via_package_walk() {
204        let m = PyModuleMap::build(&batch());
205        let mod_ref = vec!["pkg".to_string(), "sub".into(), "mod".into()]; // regular module pkg/sub/mod.py
206        // from pkg.sub.mod (regular, is_package=false): `from .. import util` (level 2) →
207        // anchor=pkg.sub, up=1 → pkg, + "util" = pkg.util
208        assert_eq!(
209            m.resolve_relative(&mod_ref, false, 2, "util"),
210            Some(vec!["pkg".into(), "util".into()])
211        );
212        // `from . import mod` (level 1) from pkg.sub.mod → anchor=pkg.sub, up=0 → pkg.sub, + "mod"
213        assert_eq!(
214            m.resolve_relative(&mod_ref, false, 1, "mod"),
215            Some(vec!["pkg".into(), "sub".into(), "mod".into()])
216        );
217        // level exceeding depth → None
218        assert_eq!(m.resolve_relative(&["top".into()], false, 3, "x"), None);
219    }
220
221    #[test]
222    fn resolve_relative_from_package_init_anchors_at_itself() {
223        // The C1 case: referencing module is the PACKAGE __init__ (key ["pkg","sub"],
224        // is_package=true). `from . import mod` (level 1) must anchor at pkg.sub ITSELF
225        // (not pkg) → pkg.sub.mod. The off-by-one bug would give pkg.mod (None).
226        let m = PyModuleMap::build(&batch());
227        let pkg_ref = vec!["pkg".to_string(), "sub".into()]; // pkg/sub/__init__.py
228        assert_eq!(
229            m.resolve_relative(&pkg_ref, true, 1, "mod"),
230            Some(vec!["pkg".into(), "sub".into(), "mod".into()])
231        );
232        // `from .. import util` (level 2) from the pkg.sub package → anchor=pkg.sub, up=1 → pkg, +util
233        assert_eq!(
234            m.resolve_relative(&pkg_ref, true, 2, "util"),
235            Some(vec!["pkg".into(), "util".into()])
236        );
237    }
238
239    #[test]
240    fn relative_import_from_top_level_module_is_none() {
241        // top.py (no parent package): `from .util import write` is invalid Python
242        // ("no known parent package") — must NOT resolve to a root-level util. (P2 round 3)
243        let m = PyModuleMap::build(&batch());
244        assert_eq!(m.resolve_relative(&["top".into()], false, 1, "util"), None);
245    }
246}
fxrank_lang_python/module_map.rs

fxrank_lang_python/
module_map.rs
