_diffctx/edges/semantic/
python.rs

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

_diffctx/edges/semantic/
python.rs
