1use std::path::{Path, PathBuf};
2
3use once_cell::sync::Lazy;
4use regex::Regex;
5use rustc_hash::{FxHashMap, FxHashSet};
6
7use crate::config::edge_weights::PYTHON_SEMANTIC;
8use crate::config::extensions::PYTHON_EXTENSIONS;
9use crate::config::weights::LANG_WEIGHTS;
10use crate::types::{Fragment, FragmentId};
11
12use super::super::EdgeDict;
13use super::super::base::{self, EdgeBuilder, path_to_module};
14
15fn is_python_file(path: &Path) -> bool {
16 let ext = base::file_ext(path);
17 PYTHON_EXTENSIONS.contains(ext.as_str())
18}
19
20static IMPORT_RE: Lazy<Regex> = Lazy::new(|| Regex::new(r"(?m)^\s*import\s+([\w.]+)").unwrap());
21static FROM_IMPORT_RE: Lazy<Regex> =
22 Lazy::new(|| Regex::new(r"(?m)^\s*from\s+([\w.]+)\s+import\s+(.+)").unwrap());
23static CALL_RE: Lazy<Regex> = Lazy::new(|| Regex::new(r"\b([A-Za-z_]\w*)\s*\(").unwrap());
24static TYPE_REF_RE: Lazy<Regex> = Lazy::new(|| Regex::new(r"\b([A-Z]\w*)\b").unwrap());
25static DEF_RE: Lazy<Regex> =
26 Lazy::new(|| Regex::new(r"(?m)^\s*(?:def|class|async\s+def)\s+([A-Za-z_]\w*)").unwrap());
27
28fn extract_imports(content: &str, path: &Path, repo_root: Option<&Path>) -> FxHashSet<String> {
29 let mut imports = FxHashSet::default();
30 for cap in IMPORT_RE.captures_iter(content) {
31 imports.insert(cap[1].to_string());
32 }
33 for cap in FROM_IMPORT_RE.captures_iter(content) {
34 let module = &cap[1];
35 if module.starts_with('.') {
36 if let Some(parent) = path.parent() {
37 let module_path = path_to_module(parent, repo_root);
38 if !module_path.is_empty() {
39 imports.insert(module_path);
40 }
41 }
42 } else {
43 imports.insert(module.to_string());
44 let parts: Vec<&str> = module.split('.').collect();
45 for i in 1..parts.len() {
46 imports.insert(parts[..i].join("."));
47 }
48 }
49 }
50 imports
51}
52
53fn extract_defines(content: &str) -> FxHashSet<String> {
54 DEF_RE
55 .captures_iter(content)
56 .map(|c| c[1].to_string())
57 .collect()
58}
59
60fn extract_calls(content: &str) -> FxHashSet<String> {
61 CALL_RE
62 .captures_iter(content)
63 .map(|c| c[1].to_string())
64 .filter(|n| !PY_KEYWORDS.contains(n.as_str()))
65 .collect()
66}
67
68fn extract_type_refs(content: &str) -> FxHashSet<String> {
69 TYPE_REF_RE
70 .captures_iter(content)
71 .map(|c| c[1].to_string())
72 .collect()
73}
74
75static PY_KEYWORDS: Lazy<FxHashSet<&str>> = Lazy::new(|| {
76 [
77 "if",
78 "for",
79 "while",
80 "return",
81 "def",
82 "class",
83 "import",
84 "from",
85 "as",
86 "with",
87 "try",
88 "except",
89 "finally",
90 "raise",
91 "pass",
92 "break",
93 "continue",
94 "yield",
95 "lambda",
96 "assert",
97 "del",
98 "elif",
99 "else",
100 "global",
101 "nonlocal",
102 "and",
103 "or",
104 "not",
105 "is",
106 "in",
107 "async",
108 "await",
109 "True",
110 "False",
111 "None",
112 "print",
113 "len",
114 "range",
115 "type",
116 "list",
117 "dict",
118 "set",
119 "tuple",
120 "str",
121 "int",
122 "float",
123 "bool",
124 "super",
125 "isinstance",
126 "hasattr",
127 "getattr",
128 "setattr",
129 "property",
130 "staticmethod",
131 "classmethod",
132 ]
133 .iter()
134 .copied()
135 .collect()
136});
137
138pub struct PythonEdgeBuilder;
139
140impl EdgeBuilder for PythonEdgeBuilder {
141 fn build(&self, fragments: &[Fragment], repo_root: Option<&Path>) -> EdgeDict {
142 let py_frags: Vec<&Fragment> = fragments
143 .iter()
144 .filter(|f| is_python_file(Path::new(f.path())))
145 .collect();
146 if py_frags.is_empty() {
147 return FxHashMap::default();
148 }
149
150 let weights = LANG_WEIGHTS.get("python").expect("python weights");
151 let call_weight = weights.call;
152 let symbol_ref_weight = weights.symbol_ref;
153 let type_ref_weight = weights.type_ref;
154
155 let mut name_to_defs: FxHashMap<String, Vec<FragmentId>> = FxHashMap::default();
156 let mut frag_defines: FxHashMap<FragmentId, FxHashSet<String>> = FxHashMap::default();
157 let mut module_to_frags: FxHashMap<String, Vec<FragmentId>> = FxHashMap::default();
158
159 for f in &py_frags {
160 let defines = extract_defines(&f.content);
161 for name in &defines {
162 name_to_defs
163 .entry(name.clone())
164 .or_default()
165 .push(f.id.clone());
166 }
167 frag_defines.insert(f.id.clone(), defines);
168
169 let module = path_to_module(Path::new(f.path()), repo_root);
170 if !module.is_empty() {
171 module_to_frags
172 .entry(module)
173 .or_default()
174 .push(f.id.clone());
175 }
176 }
177
178 let frag_imports: FxHashMap<FragmentId, FxHashSet<String>> = py_frags
179 .iter()
180 .map(|f| {
181 let imports = extract_imports(&f.content, Path::new(f.path()), repo_root);
182 (f.id.clone(), imports)
183 })
184 .collect();
185
186 let frag_to_module: FxHashMap<FragmentId, String> = py_frags
187 .iter()
188 .filter_map(|f| {
189 let m = path_to_module(Path::new(f.path()), repo_root);
190 if m.is_empty() {
191 None
192 } else {
193 Some((f.id.clone(), m))
194 }
195 })
196 .collect();
197
198 let mut edges: EdgeDict = FxHashMap::default();
199
200 for f in &py_frags {
201 let self_defs = frag_defines.get(&f.id).cloned().unwrap_or_default();
202 let src_imports = frag_imports.get(&f.id).cloned().unwrap_or_default();
203
204 let calls = extract_calls(&f.content);
205 let type_refs = extract_type_refs(&f.content);
206 let refs: FxHashSet<String> = f
207 .identifiers
208 .iter()
209 .filter(|id| !self_defs.contains(*id))
210 .cloned()
211 .collect();
212
213 for (ref_set, base_weight) in [
214 (&calls, call_weight),
215 (&refs, symbol_ref_weight),
216 (&type_refs, type_ref_weight),
217 ] {
218 for name in ref_set {
219 if self_defs.contains(name) {
220 continue;
221 }
222 if let Some(dst_ids) = name_to_defs.get(name) {
223 for dst_id in dst_ids {
224 if dst_id == &f.id {
225 continue;
226 }
227 let dst_module =
228 frag_to_module.get(dst_id).map(|s| s.as_str()).unwrap_or("");
229 let confirmed =
230 !dst_module.is_empty() && src_imports.contains(dst_module);
231 let factor = if confirmed {
232 PYTHON_SEMANTIC.import_confirmed_boost
233 } else {
234 PYTHON_SEMANTIC.import_unconfirmed_penalty
235 };
236 let w = base_weight * factor;
237 let key_fwd = (f.id.clone(), dst_id.clone());
238 let existing = edges.get(&key_fwd).copied().unwrap_or(0.0);
239 if w > existing {
240 edges.insert(key_fwd, w);
241 }
242 let rev_w = w * PYTHON_SEMANTIC.reverse_factor;
243 let key_rev = (dst_id.clone(), f.id.clone());
244 let existing_rev = edges.get(&key_rev).copied().unwrap_or(0.0);
245 if rev_w > existing_rev {
246 edges.insert(key_rev, rev_w);
247 }
248 }
249 }
250 }
251 }
252
253 for imp in &src_imports {
254 if let Some(targets) = module_to_frags.get(imp) {
255 for tgt in targets {
256 if tgt == &f.id {
257 continue;
258 }
259 base::add_edge(
260 &mut edges,
261 &f.id,
262 tgt,
263 PYTHON_SEMANTIC.import_weight,
264 PYTHON_SEMANTIC.reverse_factor,
265 );
266 }
267 }
268 }
269 }
270
271 edges
272 }
273
274 fn discover_related_files(
275 &self,
276 changed: &[PathBuf],
277 candidates: &[PathBuf],
278 repo_root: Option<&Path>,
279 file_cache: Option<&FxHashMap<PathBuf, String>>,
280 ) -> Vec<PathBuf> {
281 let py_changed: Vec<&PathBuf> = changed.iter().filter(|f| is_python_file(f)).collect();
282 if py_changed.is_empty() {
283 return vec![];
284 }
285
286 let mut file_to_module: FxHashMap<PathBuf, String> = FxHashMap::default();
287 let mut module_to_files: FxHashMap<String, Vec<PathBuf>> = FxHashMap::default();
288 let mut file_to_imports: FxHashMap<PathBuf, FxHashSet<String>> = FxHashMap::default();
289
290 for f in candidates {
291 if !is_python_file(f) {
292 continue;
293 }
294 let module = path_to_module(f, repo_root);
295 if !module.is_empty() {
296 file_to_module.insert(f.clone(), module.clone());
297 module_to_files
298 .entry(module.clone())
299 .or_default()
300 .push(f.clone());
301 let parts: Vec<&str> = module.split('.').collect();
302 for i in 1..parts.len() {
303 module_to_files
304 .entry(parts[..i].join("."))
305 .or_default()
306 .push(f.clone());
307 }
308 }
309 let content = base::read_file_cached(f, file_cache);
310 if let Some(c) = content {
311 file_to_imports.insert(f.clone(), extract_imports(&c, f, repo_root));
312 }
313 }
314
315 let changed_set: FxHashSet<PathBuf> = changed.iter().cloned().collect();
316 let mut discovered: FxHashSet<PathBuf> = FxHashSet::default();
317 let mut frontier: FxHashSet<PathBuf> = py_changed.iter().map(|f| (*f).clone()).collect();
318
319 for _ in 0..2 {
320 let mut next_frontier: FxHashSet<PathBuf> = FxHashSet::default();
321 for f in &frontier {
322 let f_imports = file_to_imports.get(f).cloned().unwrap_or_default();
323 for imp in &f_imports {
324 if let Some(targets) = module_to_files.get(imp) {
325 for target in targets {
326 if !changed_set.contains(target) && !discovered.contains(target) {
327 discovered.insert(target.clone());
328 next_frontier.insert(target.clone());
329 }
330 }
331 }
332 }
333 let f_module = file_to_module
334 .get(f)
335 .cloned()
336 .unwrap_or_else(|| path_to_module(f, repo_root));
337 if !f_module.is_empty() {
338 for (candidate, cand_imports) in &file_to_imports {
339 if !changed_set.contains(candidate)
340 && !discovered.contains(candidate)
341 && cand_imports.contains(&f_module)
342 {
343 discovered.insert(candidate.clone());
344 next_frontier.insert(candidate.clone());
345 }
346 }
347 }
348 }
349 if next_frontier.is_empty() {
350 break;
351 }
352 frontier = next_frontier;
353 }
354
355 let mut result: Vec<PathBuf> = discovered.into_iter().collect();
356 result.sort();
357 result
358 }
359}