1use crate::graph::PythonGraphSource;
6use crate::graph::hits::{
7 record_hit, record_import_hit, record_self_receiver_hit, record_unproven_hit,
8};
9use crate::graph::resolver::{
10 annotation_class_qualifier_site, annotation_reference_candidates, member_name,
11 normalized_receiver_type, receiver_annotation_matches_target,
12 resolve_callable_parameter_default_types, resolve_constructor_types, resolve_receiver_type,
13 target_owner_code_unit, top_level_identifier,
14};
15use crate::graph_support::{PythonSource, PythonUsageSource};
16use crate::imports::{PythonImportBinding, parse_python_import_bindings, resolve_fqn_candidates};
17use crate::usage_index::{
18 ModuleBindingEvent, ModuleBindingEventKind, ModuleBindingTimeline, PythonScopeFacts,
19 usage_matching_edges, usage_module_binding_timeline, usage_resolve_module_files,
20 usage_scope_facts,
21};
22use brokk_bifrost_core::analyzer::symbol_path::parse_symbol_path;
23use brokk_bifrost_core::analyzer::usages::local_inference::{
24 LocalBindingsSnapshot, LocalInferenceConfig, LocalInferenceEngine, SymbolResolution,
25};
26use brokk_bifrost_core::analyzer::usages::model::{ImportKind, UsageHit};
27use brokk_bifrost_core::analyzer::usages::{ImportEdge, ImportEdgeKind};
28use brokk_bifrost_core::analyzer::{CodeUnit, CodeUnitIndex, Language, ProjectFile, Range};
29use brokk_bifrost_core::cancellation::CancellationToken;
30use brokk_bifrost_core::hash::{HashMap, HashSet};
31use brokk_bifrost_core::text_utils::compute_line_starts;
32use rayon::prelude::*;
33use std::collections::BTreeSet;
34use std::sync::{Arc, Mutex};
35use tree_sitter::{Node, Parser, Tree};
36
37pub struct ParsedFile {
38 pub source: Arc<String>,
39 pub tree: Tree,
40}
41
42pub struct PythonProjectGraph {
43 parsed: HashMap<ProjectFile, ParsedFile>,
44}
45
46impl PythonProjectGraph {
47 pub fn scan_files(
48 &self,
49 candidate_files: &HashSet<ProjectFile>,
50 target_file: &ProjectFile,
51 ) -> HashSet<ProjectFile> {
52 candidate_files
53 .iter()
54 .cloned()
55 .chain(std::iter::once(target_file.clone()))
56 .collect()
57 }
58}
59
60pub fn build_python_graph(
61 candidate_files: &HashSet<ProjectFile>,
62 target_file: &ProjectFile,
63 cancellation: Option<&CancellationToken>,
64) -> PythonProjectGraph {
65 let parser_language = tree_sitter_python::LANGUAGE.into();
66 let files: HashSet<ProjectFile> = candidate_files
67 .iter()
68 .cloned()
69 .chain(std::iter::once(target_file.clone()))
70 .collect();
71 let mut parsed = HashMap::default();
72
73 for file in files {
74 if cancellation.is_some_and(CancellationToken::is_cancelled) {
75 break;
76 }
77 let Ok(source) = file.read_to_string() else {
78 continue;
79 };
80 if cancellation.is_some_and(CancellationToken::is_cancelled) {
81 break;
82 }
83 if source.is_empty() {
84 continue;
85 }
86 let mut parser = Parser::new();
87 if parser.set_language(&parser_language).is_err() {
88 continue;
89 }
90 let Some(tree) = parser.parse(source.as_str(), None) else {
91 continue;
92 };
93 if cancellation.is_some_and(CancellationToken::is_cancelled) {
94 break;
95 }
96 parsed.insert(
97 file,
98 ParsedFile {
99 source: Arc::new(source),
100 tree,
101 },
102 );
103 }
104
105 PythonProjectGraph { parsed }
106}
107
108pub fn scan_files_for_seeds(
109 graph: &PythonGraphSource<'_>,
110 python: &dyn PythonUsageSource,
111 project_graph: &PythonProjectGraph,
112 files: &HashSet<ProjectFile>,
113 target: &CodeUnit,
114 seeds: &BTreeSet<(ProjectFile, String)>,
115 cancellation: Option<&CancellationToken>,
116) -> ScanResult {
117 let collected: Mutex<BTreeSet<UsageHit>> = Mutex::new(BTreeSet::new());
118 let unproven_collected: Mutex<BTreeSet<UsageHit>> = Mutex::new(BTreeSet::new());
119 let target_short = top_level_identifier(graph.index, target);
120 let target_member = member_name(graph.index, target);
121 let target_owner = target_owner_code_unit(graph.index, target);
122 let member_unique_in_target_file = target_member.as_deref().is_some_and(|member| {
126 let owners: HashSet<CodeUnit> = graph
127 .index
128 .declarations(target.source())
129 .into_iter()
130 .filter(|decl| {
131 decl.identifier() == member && target_owner_code_unit(graph.index, decl).is_some()
132 })
133 .filter_map(|decl| target_owner_code_unit(graph.index, &decl))
134 .collect();
135 owners.len() == 1
136 });
137 let files_vec: Vec<&ProjectFile> = files.iter().collect();
138 let parser_language = tree_sitter_python::LANGUAGE.into();
139
140 files_vec.par_iter().for_each(|file| {
141 if cancellation.is_some_and(CancellationToken::is_cancelled) {
142 return;
143 }
144 let owned_source: Option<Arc<String>>;
145 let owned_tree: Option<Tree>;
146 let (source_str, tree_ref) = if let Some(parsed) = project_graph.parsed.get(*file) {
147 (parsed.source.as_str(), &parsed.tree)
148 } else {
149 let Ok(source) = file.read_to_string() else {
150 return;
151 };
152 if source.is_empty() {
153 return;
154 }
155 let mut parser = Parser::new();
156 if parser.set_language(&parser_language).is_err() {
157 return;
158 }
159 let Some(tree) = parser.parse(source.as_str(), None) else {
160 return;
161 };
162 owned_source = Some(Arc::new(source));
163 owned_tree = Some(tree);
164 (
165 owned_source.as_deref().unwrap().as_str(),
166 owned_tree.as_ref().unwrap(),
167 )
168 };
169 if cancellation.is_some_and(CancellationToken::is_cancelled) {
170 return;
171 }
172
173 let edges = {
174 let _scope = brokk_bifrost_core::profiling::scope("python_graph::matching_edges");
175 usage_matching_edges(python, file, seeds)
176 };
177 if !file_may_reference_target(
182 tree_ref.root_node(),
183 source_str,
184 target,
185 target_short.as_str(),
186 target_member.as_deref(),
187 &edges,
188 ) {
189 return;
190 }
191 let raw_module_bindings = {
192 let _scope =
193 brokk_bifrost_core::profiling::scope("python_graph::module_binding_timeline");
194 usage_module_binding_timeline(python, file, || {
195 collect_module_binding_timeline(tree_ref.root_node(), source_str)
196 })
197 };
198 let module_bindings = classify_module_binding_timeline(
199 python,
200 file,
201 raw_module_bindings.as_ref(),
202 seeds,
203 &edges,
204 );
205 let scoped_import_bindings = parse_python_import_bindings(source_str);
206 let target_self_file = *file == target.source();
207 let scope_facts = {
208 let _scope = brokk_bifrost_core::profiling::scope("python_graph::scope_facts");
209 usage_scope_facts(python, file, || {
210 collect_scope_facts_from_parsed_source(
211 graph,
212 python,
213 file,
214 source_str,
215 tree_ref.root_node(),
216 )
217 })
218 };
219 let scope_range_index = build_scope_range_index(graph, scope_facts.as_ref());
220
221 let mut local_hits = BTreeSet::new();
222 let mut local_unproven_hits = BTreeSet::new();
223 let line_starts = compute_line_starts(source_str);
224
225 let mut scan_ctx = ScanCtx {
226 python,
227 file,
228 source: source_str,
229 line_starts: &line_starts,
230 graph,
231 target,
232 target_short: &target_short,
233 target_member: target_member.as_deref(),
234 target_owner: target_owner.clone(),
235 target_is_module: target.is_module(),
236 target_source: target.source(),
237 seeds,
238 edges: &edges,
239 target_self_file,
240 member_best_effort_unique: target_self_file && member_unique_in_target_file,
241 raw_module_bindings: raw_module_bindings.as_ref(),
242 module_bindings: &module_bindings,
243 scoped_import_bindings: &scoped_import_bindings,
244 scope_facts: scope_facts.as_ref(),
245 scope_range_index: &scope_range_index,
246 hits: &mut local_hits,
247 unproven_hits: &mut local_unproven_hits,
248 };
249
250 {
251 let _scope = brokk_bifrost_core::profiling::scope("python_graph::scan_tree");
252 scan_node(tree_ref.root_node(), &mut scan_ctx);
253 }
254
255 if !local_hits.is_empty() {
256 let mut sink = collected
257 .lock()
258 .expect("usage hit collector mutex poisoned");
259 sink.extend(local_hits);
260 }
261 if !local_unproven_hits.is_empty() {
262 let mut sink = unproven_collected
263 .lock()
264 .expect("usage unproven hit collector mutex poisoned");
265 sink.extend(local_unproven_hits);
266 }
267 });
268
269 ScanResult {
270 hits: collected
271 .into_inner()
272 .expect("usage hit collector mutex poisoned"),
273 unproven_hits: unproven_collected
274 .into_inner()
275 .expect("usage unproven hit collector mutex poisoned"),
276 }
277}
278
279fn file_may_reference_target(
280 root: Node<'_>,
281 source: &str,
282 target: &CodeUnit,
283 target_short: &str,
284 target_member: Option<&str>,
285 edges: &[ImportEdge],
286) -> bool {
287 if target.is_module() {
288 return true;
289 }
290
291 let mut stack = vec![root];
292 while let Some(node) = stack.pop() {
293 if node.kind() == "string" {
294 return true;
298 }
299 if node.kind() == "identifier" {
300 let name = slice(node, source);
301 if name == target_short
302 || target_member.is_some_and(|member| name == member)
303 || edges.iter().any(|edge| edge.local_name == name)
304 {
305 return true;
306 }
307 }
308
309 let mut cursor = node.walk();
310 stack.extend(node.named_children(&mut cursor));
311 }
312 false
313}
314
315pub struct ScanResult {
316 pub hits: BTreeSet<UsageHit>,
317 pub unproven_hits: BTreeSet<UsageHit>,
318}
319
320pub struct ScanCtx<'a> {
321 python: &'a dyn PythonUsageSource,
322 pub file: &'a ProjectFile,
323 pub source: &'a str,
324 pub line_starts: &'a [usize],
325 pub graph: &'a PythonGraphSource<'a>,
326 target: &'a CodeUnit,
327 target_short: &'a str,
328 target_member: Option<&'a str>,
329 target_owner: Option<CodeUnit>,
330 target_is_module: bool,
331 target_source: &'a ProjectFile,
332 seeds: &'a BTreeSet<(ProjectFile, String)>,
333 edges: &'a [ImportEdge],
334 target_self_file: bool,
335 member_best_effort_unique: bool,
341 raw_module_bindings: &'a ModuleBindingTimeline,
342 module_bindings: &'a HashMap<String, Vec<ClassifiedModuleBindingEvent>>,
343 scoped_import_bindings: &'a [PythonImportBinding],
344 scope_facts: &'a HashMap<CodeUnit, LocalBindingsSnapshot<String>>,
345 scope_range_index: &'a [ScopeRangeEntry],
346 pub hits: &'a mut BTreeSet<UsageHit>,
347 pub unproven_hits: &'a mut BTreeSet<UsageHit>,
348}
349
350struct ScopeRangeEntry {
351 range: Range,
352 scope: CodeUnit,
353 prefix_max_end: usize,
354}
355
356fn build_scope_range_index(
357 graph: &PythonGraphSource<'_>,
358 scope_facts: &HashMap<CodeUnit, LocalBindingsSnapshot<String>>,
359) -> Vec<ScopeRangeEntry> {
360 let mut entries = scope_facts
361 .keys()
362 .flat_map(|scope| {
363 graph
364 .index
365 .ranges(scope)
366 .into_iter()
367 .map(|range| ScopeRangeEntry {
368 range,
369 scope: scope.clone(),
370 prefix_max_end: 0,
371 })
372 })
373 .collect::<Vec<_>>();
374 entries.sort_by(|left, right| {
375 left.range
376 .start_byte
377 .cmp(&right.range.start_byte)
378 .then_with(|| right.range.end_byte.cmp(&left.range.end_byte))
379 .then_with(|| left.scope.cmp(&right.scope))
380 });
381 let mut max_end = 0;
382 for entry in &mut entries {
383 max_end = max_end.max(entry.range.end_byte);
384 entry.prefix_max_end = max_end;
385 }
386 entries
387}
388
389fn indexed_scope_entry<'entry, 'facts>(
390 scope_range_index: &'entry [ScopeRangeEntry],
391 scope_facts: &'facts HashMap<CodeUnit, LocalBindingsSnapshot<String>>,
392 node: Node<'_>,
393 mut skip_innermost: usize,
394) -> Option<(&'entry CodeUnit, &'facts LocalBindingsSnapshot<String>)> {
395 let mut cursor =
396 scope_range_index.partition_point(|entry| entry.range.start_byte <= node.start_byte());
397 while cursor > 0 {
398 cursor -= 1;
399 let entry = &scope_range_index[cursor];
400 if entry.prefix_max_end < node.end_byte() {
401 return None;
402 }
403 if entry.range.end_byte >= node.end_byte() {
404 if skip_innermost > 0 {
405 skip_innermost -= 1;
406 continue;
407 }
408 return scope_facts
409 .get(&entry.scope)
410 .map(|facts| (&entry.scope, facts));
411 }
412 }
413 None
414}
415
416pub fn enclosing_scope_facts<'a>(
420 index: &dyn CodeUnitIndex,
421 file: &ProjectFile,
422 scope_facts: &'a HashMap<CodeUnit, LocalBindingsSnapshot<String>>,
423 node: Node<'_>,
424) -> Option<&'a LocalBindingsSnapshot<String>> {
425 let range = Range {
426 start_byte: node.start_byte(),
427 end_byte: node.end_byte(),
428 start_line: 0,
429 end_line: 0,
430 };
431 let enclosing = index.enclosing_code_unit(file, &range)?;
432 scope_facts.get(&enclosing)
433}
434
435impl ScanCtx<'_> {
436 fn scope_entry_for_node(
437 &self,
438 node: Node<'_>,
439 ) -> Option<(&CodeUnit, &LocalBindingsSnapshot<String>)> {
440 indexed_scope_entry(
441 self.scope_range_index,
442 self.scope_facts,
443 node,
444 usize::from(function_declaration_expression_is_outer_scoped(node)),
445 )
446 }
447
448 fn scope_facts_for_node(&self, node: Node<'_>) -> Option<&LocalBindingsSnapshot<String>> {
449 self.scope_entry_for_node(node).map(|(_, facts)| facts)
450 }
451
452 fn binds_target(&self, ident: &str, node: Node<'_>) -> bool {
453 let scope_entry = self.scope_entry_for_node(node);
454 if self.target_self_file
455 && ident == self.target_short
456 && scope_entry
457 .is_none_or(|(scope, facts)| scope.is_module() || !facts.is_shadowed(ident))
458 {
459 return true;
460 }
461 if scope_entry.is_some_and(|(scope, facts)| !scope.is_module() && facts.is_shadowed(ident))
462 {
463 return false;
464 }
465 self.module_binding_targets_query(ident, node)
466 }
467
468 fn receiver_binds_target(&self, expr: &str, node: Node<'_>) -> bool {
469 if self.binds_target(expr, node) {
470 return true;
471 }
472
473 if self.target_member.is_some() && self.import_edge_visible_for(expr, node) {
474 return true;
475 }
476
477 if matches!(expr, "self" | "cls") && self.self_receiver_matches_target(node) {
481 return true;
482 }
483
484 match enclosing_runtime_parameter_type(expr, node, self.source) {
485 EnclosingParameterType::Typed(raw_type) => {
486 return self.receiver_type_matches_target(&raw_type);
487 }
488 EnclosingParameterType::Untyped => return false,
489 EnclosingParameterType::NotDeclared => {}
490 }
491
492 let Some(scope_facts) = self.scope_facts_for_node(node) else {
493 return false;
494 };
495 let resolution = scope_facts.resolution_for(expr);
496 let Some(raw_type) = resolution
497 .as_precise()
498 .and_then(|targets| targets.iter().next())
499 else {
500 return false;
501 };
502 self.receiver_type_matches_target(raw_type)
503 }
504
505 fn node_directly_in_owner_class_body(&self, node: Node<'_>) -> bool {
511 let Some(target_owner) = self.target_owner.as_ref() else {
512 return false;
513 };
514 let range = Range {
515 start_byte: node.start_byte(),
516 end_byte: node.end_byte(),
517 start_line: 0,
518 end_line: 0,
519 };
520 let Some(enclosing) = self.graph.index.enclosing_code_unit(self.file, &range) else {
521 return false;
522 };
523 if &enclosing == target_owner {
524 return true;
525 }
526 if enclosing.is_function() {
527 return target_owner_code_unit(self.graph.index, &enclosing).as_ref()
528 == Some(target_owner)
529 && function_declaration_expression_is_outer_scoped(node);
530 }
531 target_owner_code_unit(self.graph.index, &enclosing).as_ref() == Some(target_owner)
532 }
533
534 fn receiver_type_is_unknown(&self, expr: &str, node: Node<'_>) -> bool {
538 match enclosing_runtime_parameter_type(expr, node, self.source) {
539 EnclosingParameterType::Typed(_) => return false,
540 EnclosingParameterType::Untyped => return true,
541 EnclosingParameterType::NotDeclared => {}
542 }
543 match self.scope_facts_for_node(node) {
544 Some(facts) => facts.resolution_for(expr).is_unknown(),
545 None => true,
546 }
547 }
548
549 fn import_edge_visible_for(&self, ident: &str, node: Node<'_>) -> bool {
550 if let Some(scope_facts) = self.scope_facts_for_node(node)
551 && scope_facts.is_shadowed(ident)
552 {
553 return false;
554 }
555 self.module_binding_targets_query(ident, node)
556 }
557
558 fn module_binding_targets_query(&self, ident: &str, node: Node<'_>) -> bool {
559 if self.target_member.is_some() {
565 return self.module_binding_targets_symbol(ident, node);
566 }
567 if let Some(matches) = self.function_import_binding_targets_query(ident, node) {
568 return matches;
569 }
570 self.module_binding_matches_query(ident, node, true, |kind| {
571 kind != ModuleBindingKind::Other
572 })
573 }
574
575 fn module_binding_targets_symbol(&self, ident: &str, node: Node<'_>) -> bool {
576 if let Some(matches) = self.function_import_binding_targets_query(ident, node) {
577 return matches;
578 }
579 let unclassified_named_import = self.edges.iter().any(|edge| {
580 edge.local_name == ident && !matches!(edge.kind, ImportEdgeKind::Namespace)
581 });
582 self.module_binding_matches_query(ident, node, unclassified_named_import, |kind| {
583 kind == ModuleBindingKind::TargetSymbolImport
584 })
585 }
586
587 fn function_import_binding_targets_query(&self, ident: &str, node: Node<'_>) -> Option<bool> {
592 let binding = self.scoped_import_bindings.iter().rev().find(|binding| {
593 binding.is_function_scoped()
594 && binding.start_byte <= node.start_byte()
595 && binding.scope_start_byte <= node.start_byte()
596 && node.end_byte() <= binding.scope_end_byte
597 && binding.local_name == ident
598 })?;
599 let candidates = resolve_fqn_candidates(self.python, &binding.qualified_name, |name| {
600 self.graph.index.definitions(name).collect()
601 });
602 let imported_target = self.target_owner.as_ref().unwrap_or(self.target);
603 Some(
604 candidates
605 .iter()
606 .any(|candidate| candidate == imported_target),
607 )
608 }
609
610 fn module_binding_matches_query(
611 &self,
612 ident: &str,
613 node: Node<'_>,
614 unclassified: bool,
615 matches: impl Fn(ModuleBindingKind) -> bool,
616 ) -> bool {
617 if !self.edges.iter().any(|edge| edge.local_name == ident) {
618 return false;
619 }
620 let Some(events) = self.module_bindings.get(ident) else {
621 return unclassified;
622 };
623 let cutoff = if reference_is_deferred_function_body(node) {
624 usize::MAX
625 } else {
626 node.start_byte()
627 };
628 let visible: Vec<_> = events
629 .iter()
630 .filter(|event| event.visible_from <= cutoff)
631 .collect();
632 let start = visible
633 .iter()
634 .rposition(|event| !event.conditional)
635 .unwrap_or(0);
636 visible[start..].iter().any(|event| matches(event.kind))
637 }
638
639 fn self_receiver_matches_target(&self, node: Node<'_>) -> bool {
643 let Some(target_owner) = self.target_owner.as_ref() else {
644 return false;
645 };
646 let range = Range {
647 start_byte: node.start_byte(),
648 end_byte: node.end_byte(),
649 start_line: 0,
650 end_line: 0,
651 };
652 let Some(enclosing) = self.graph.index.enclosing_code_unit(self.file, &range) else {
653 return false;
654 };
655 let enclosing_class = if enclosing.is_class() {
656 enclosing
657 } else {
658 match target_owner_code_unit(self.graph.index, &enclosing) {
659 Some(class) => class,
660 None => return false,
661 }
662 };
663 if &enclosing_class == target_owner {
664 return true;
665 }
666 self.graph
667 .hierarchy
668 .map(|provider| provider.get_ancestors(&enclosing_class))
669 .unwrap_or_default()
670 .into_iter()
671 .any(|ancestor| ancestor == *target_owner)
672 }
673
674 fn receiver_type_matches_target(&self, raw_type: &str) -> bool {
675 let Some(target_owner) = self.target_owner.as_ref() else {
676 return false;
677 };
678 if let Some(receiver_type) = resolve_receiver_type(
679 self.graph,
680 self.python,
681 self.file,
682 raw_type,
683 self.target_self_file,
684 ) {
685 if &receiver_type == target_owner {
686 return true;
687 }
688 return self
689 .graph
690 .hierarchy
691 .map(|provider| provider.get_ancestors(&receiver_type))
692 .unwrap_or_default()
693 .into_iter()
694 .any(|ancestor| ancestor == *target_owner);
695 }
696
697 receiver_annotation_matches_target(
702 raw_type,
703 self.edges,
704 self.target_short,
705 self.target_self_file,
706 )
707 }
708}
709
710pub(crate) fn function_declaration_expression_is_outer_scoped(node: Node<'_>) -> bool {
711 let site_start = node.start_byte();
712 let site_end = node.end_byte();
713 let mut current = node;
714 while let Some(parent) = current.parent() {
715 if parent.kind() == "function_definition" {
716 if parent
717 .child_by_field_name("body")
718 .is_some_and(|body| body.start_byte() <= site_start && site_end <= body.end_byte())
719 {
720 return false;
721 }
722 if parent
723 .child_by_field_name("name")
724 .is_some_and(|name| name.id() == node.id())
725 {
726 return false;
727 }
728 if let Some(parameters) = parent.child_by_field_name("parameters")
729 && parameters.start_byte() <= site_start
730 && site_end <= parameters.end_byte()
731 {
732 let mut parameter = node;
733 while parameter.parent() != Some(parameters) {
734 let Some(next) = parameter.parent() else {
735 return false;
736 };
737 parameter = next;
738 }
739 let binder = if parameter.kind() == "identifier" {
740 Some(parameter)
741 } else {
742 parameter.child_by_field_name("name").or_else(|| {
743 parameter
744 .named_child(0)
745 .filter(|child| child.kind() == "identifier")
746 })
747 };
748 return binder.is_none_or(|binder| binder.id() != node.id());
749 }
750 return true;
751 }
752 if parent.kind() == "decorated_definition" {
753 return current.kind() == "decorator";
754 }
755 if parent.kind() == "class_definition" {
756 break;
757 }
758 current = parent;
759 }
760 false
761}
762
763fn scan_node(node: Node<'_>, ctx: &mut ScanCtx<'_>) {
764 let mut stack = vec![node];
765 while let Some(node) = stack.pop() {
766 match node.kind() {
767 "import_statement" | "import_from_statement" => {
768 handle_import_candidate(node, ctx);
769 continue;
770 }
771 "identifier" => {
772 if handle_annotation_reference_candidate(node, ctx) {
773 continue;
774 }
775 handle_identifier_candidate(node, ctx);
776 }
777 "attribute" => {
778 if handle_annotation_reference_candidate(node, ctx) {
779 continue;
780 }
781 handle_attribute_candidate(node, ctx);
782 }
783 "string_content" => {
784 handle_annotation_reference_candidate(node, ctx);
785 }
786 "keyword_argument" => {
787 handle_keyword_argument_candidate(node, ctx);
788 if let Some(value) = node.child_by_field_name("value") {
789 stack.push(value);
790 }
791 continue;
792 }
793 _ => {}
794 }
795
796 let mut cursor = node.walk();
797 let mut children: Vec<Node<'_>> = node.named_children(&mut cursor).collect();
798 children.reverse();
799 stack.extend(children);
800 }
801}
802
803fn handle_annotation_reference_candidate(node: Node<'_>, ctx: &mut ScanCtx<'_>) -> bool {
804 let Some(candidates) = annotation_reference_candidates(
805 ctx.graph,
806 ctx.python,
807 ctx.file,
808 ctx.source,
809 node,
810 ctx.target_self_file,
811 ) else {
812 return false;
813 };
814
815 if (ctx.target.is_class() || ctx.target.is_field() || ctx.target_member.is_none())
816 && candidates.iter().all(|candidate| *candidate == *ctx.target)
817 && candidates.iter().any(|candidate| *candidate == *ctx.target)
818 {
819 let site = if node.kind() == "attribute" {
820 node.child_by_field_name("attribute").unwrap_or(node)
821 } else {
822 node
823 };
824 record_hit(site, ctx);
825 }
826
827 if let Some(site) = annotation_class_qualifier_site(
828 ctx.graph, ctx.python, ctx.file, ctx.source, node, ctx.target,
829 ) {
830 record_hit(site, ctx);
831 }
832
833 if node.kind() == "attribute" && (candidates.is_empty() || ctx.target_is_module) {
839 return false;
840 }
841
842 true
843}
844
845fn handle_keyword_argument_candidate(node: Node<'_>, ctx: &mut ScanCtx<'_>) {
846 let (Some(target_member), Some(name), Some(arguments)) = (
847 ctx.target_member,
848 node.child_by_field_name("name"),
849 node.parent(),
850 ) else {
851 return;
852 };
853 if name.kind() != "identifier"
854 || slice(name, ctx.source) != target_member
855 || arguments.kind() != "argument_list"
856 {
857 return;
858 }
859 let Some(call) = arguments.parent().filter(|parent| parent.kind() == "call") else {
860 return;
861 };
862 let Some(function) = call.child_by_field_name("function") else {
863 return;
864 };
865 if function.kind() == "identifier" && slice(function, ctx.source) == "cls" {
866 if ctx.self_receiver_matches_target(function) {
867 record_hit(name, ctx);
868 }
869 return;
870 }
871 let Some(target_owner) = ctx.target_owner.as_ref() else {
872 return;
873 };
874 let scoped_callee_matches = if function.kind() == "identifier" {
875 ctx.scope_facts_for_node(function)
876 .and_then(|facts| {
877 facts
878 .resolution_for(slice(function, ctx.source))
879 .as_precise()
880 .and_then(|targets| targets.iter().next().cloned())
881 })
882 .is_some_and(|raw_type| ctx.receiver_type_matches_target(&raw_type))
883 } else {
884 false
885 };
886 let default_callee_matches = if function.kind() == "identifier" {
887 resolve_callable_parameter_default_types(
888 ctx.graph,
889 ctx.python,
890 ctx.file,
891 ctx.source,
892 function,
893 slice(function, ctx.source),
894 )
895 .into_iter()
896 .any(|class| {
897 &class == target_owner
898 || ctx
899 .graph
900 .hierarchy
901 .map(|provider| provider.get_ancestors(&class))
902 .unwrap_or_default()
903 .into_iter()
904 .any(|ancestor| &ancestor == target_owner)
905 })
906 } else {
907 false
908 };
909 let root_shadowed = leftmost_identifier(function).is_some_and(|root| {
910 ctx.scope_facts_for_node(function)
911 .is_some_and(|facts| facts.is_shadowed(slice(root, ctx.source)))
912 });
913 if root_shadowed && !scoped_callee_matches && !default_callee_matches {
914 return;
915 }
916 let matches = scoped_callee_matches
917 || default_callee_matches
918 || (!root_shadowed
919 && resolve_constructor_types(ctx.graph, ctx.python, ctx.file, ctx.source, function)
920 .into_iter()
921 .any(|class| {
922 &class == target_owner
923 || ctx
924 .graph
925 .hierarchy
926 .map(|provider| provider.get_ancestors(&class))
927 .unwrap_or_default()
928 .into_iter()
929 .any(|ancestor| &ancestor == target_owner)
930 }));
931 if matches {
932 record_hit(name, ctx);
933 }
934}
935
936fn leftmost_identifier(mut node: Node<'_>) -> Option<Node<'_>> {
937 loop {
938 match node.kind() {
939 "identifier" => return Some(node),
940 "attribute" => node = node.child_by_field_name("object")?,
941 _ => return None,
942 }
943 }
944}
945
946fn handle_import_candidate(node: Node<'_>, ctx: &mut ScanCtx<'_>) {
952 if ctx.target_member.is_some() {
953 return;
954 }
955 if !ctx
956 .edges
957 .iter()
958 .any(|edge| edge.local_name == ctx.target_short)
959 {
960 return;
961 }
962 let mut stack = vec![node];
963 while let Some(node) = stack.pop() {
964 if node.kind() == "identifier" && slice(node, ctx.source) == ctx.target_short {
965 record_import_hit(node, ctx);
966 return;
967 }
968 let mut cursor = node.walk();
969 for child in node.named_children(&mut cursor) {
970 stack.push(child);
971 }
972 }
973}
974
975fn handle_identifier_candidate(node: Node<'_>, ctx: &mut ScanCtx<'_>) {
976 if node
977 .parent()
978 .is_some_and(|parent| parent.kind() == "attribute")
979 {
980 return;
981 }
982 let text = slice(node, ctx.source);
983 if text.is_empty() || is_declaration_identifier(node) || decorates_the_target(node, ctx) {
984 return;
985 }
986 if let Some(member) = ctx.target_member {
987 if member == "__init__" && is_call_callee(node) && ctx.binds_target(text, node) {
990 record_hit(node, ctx);
991 return;
992 }
993 if text == member && ctx.node_directly_in_owner_class_body(node) {
996 record_hit(node, ctx);
997 }
998 return;
999 }
1000 if !ctx.binds_target(text, node) {
1001 return;
1002 }
1003 if !ctx.target_is_module
1004 && ctx.edges.iter().any(|edge| edge.local_name == text)
1005 && !ctx.module_binding_targets_symbol(text, node)
1006 {
1007 return;
1008 }
1009 record_hit(node, ctx);
1010}
1011
1012fn decorates_the_target(node: Node<'_>, ctx: &ScanCtx<'_>) -> bool {
1020 if ctx.file != ctx.target_source {
1021 return false;
1022 }
1023 let mut current = node;
1024 while let Some(parent) = current.parent() {
1025 if parent.kind() == "decorated_definition" && current.kind() == "decorator" {
1026 let Some(definition) = parent.child_by_field_name("definition") else {
1027 return false;
1028 };
1029 return ctx.graph.index.ranges(ctx.target).iter().any(|range| {
1030 range.start_byte <= definition.start_byte()
1031 && definition.end_byte() <= range.end_byte
1032 });
1033 }
1034 current = parent;
1035 }
1036 false
1037}
1038
1039fn handle_attribute_candidate(node: Node<'_>, ctx: &mut ScanCtx<'_>) {
1040 let Some(object) = node.child_by_field_name("object") else {
1041 return;
1042 };
1043 let Some(attribute) = node.child_by_field_name("attribute") else {
1044 return;
1045 };
1046 let object_text = slice(object, ctx.source);
1047 let attribute_text = slice(attribute, ctx.source);
1048 if let Some(member) = ctx.target_member
1049 && attribute_text == member
1050 {
1051 let is_same_owner_receiver =
1056 matches!(object_text, "self" | "cls") && ctx.self_receiver_matches_target(node);
1057 if is_same_owner_receiver {
1058 record_self_receiver_hit(attribute, ctx);
1059 } else if ctx.receiver_binds_target(object_text, node)
1060 || (object.kind() == "call" && call_result_matches_target(object, ctx))
1061 {
1062 record_hit(attribute, ctx);
1063 } else if member_receiver_match_is_unproven(object, object_text, node, ctx) {
1064 record_unproven_hit(attribute, ctx);
1065 }
1066 }
1067
1068 let object_binds_target = if ctx.target_is_module {
1069 imported_root_targets_module(ctx, object, node)
1070 } else {
1071 ctx.binds_target(object_text, node)
1072 };
1073 if object.kind() == "identifier"
1074 && object_binds_target
1075 && (ctx.target_is_module
1076 || (ctx.target_member.is_none()
1077 && !ctx.edges.iter().any(|edge| {
1078 matches!(edge.kind, ImportEdgeKind::Namespace) && edge.local_name == object_text
1079 })))
1080 {
1081 record_hit(object, ctx);
1082 }
1083
1084 if ctx.target_is_module
1085 && let Some(module_qualifier) = module_attribute_target_hit(node, ctx)
1086 {
1087 record_hit(module_qualifier, ctx);
1088 }
1089
1090 if let Some(member) = ctx.target_member
1094 && object.kind() == "identifier"
1095 && object_text == member
1096 && ctx.node_directly_in_owner_class_body(object)
1097 {
1098 record_hit(object, ctx);
1099 }
1100
1101 if ctx.member_best_effort_unique
1107 && let Some(member) = ctx.target_member
1108 && attribute_text == member
1109 && object.kind() == "identifier"
1110 && !matches!(object_text, "self" | "cls")
1111 && !ctx.receiver_binds_target(object_text, node)
1112 && ctx.receiver_type_is_unknown(object_text, node)
1113 {
1114 record_hit(attribute, ctx);
1115 }
1116
1117 if let Some(module_binding_target) = module_binding_attribute_target_hit(node, ctx) {
1118 record_hit(module_binding_target, ctx);
1119 }
1120}
1121
1122fn module_binding_attribute_target_hit<'a>(node: Node<'a>, ctx: &ScanCtx<'_>) -> Option<Node<'a>> {
1134 if ctx.target_member.is_some() {
1135 return None;
1136 }
1137 let (root, attributes) = attribute_chain(node)?;
1138 let terminal = *attributes.last()?;
1139 let terminal_name = slice(terminal, ctx.source);
1140 if terminal_name.is_empty()
1141 || !ctx
1142 .seeds
1143 .iter()
1144 .any(|(_, seed_name)| seed_name == terminal_name)
1145 {
1146 return None;
1147 }
1148
1149 for binding in imported_module_bindings(ctx, root, node) {
1150 let mut written_module = binding.module.clone();
1151 for attribute in attributes
1152 [binding.consumed_attributes.min(attributes.len() - 1)..attributes.len() - 1]
1153 .iter()
1154 {
1155 let segment = slice(*attribute, ctx.source);
1156 if segment.is_empty() {
1157 return None;
1158 }
1159 written_module.push('.');
1160 written_module.push_str(segment);
1161 }
1162 if usage_resolve_module_files(ctx.python, ctx.file, &written_module)
1163 .iter()
1164 .any(|resolved| {
1165 ctx.seeds
1166 .contains(&(resolved.clone(), terminal_name.to_string()))
1167 })
1168 {
1169 return Some(terminal);
1170 }
1171
1172 let mut written_fqn = binding.module;
1173 for attribute in attributes.iter().skip(binding.consumed_attributes) {
1174 let segment = slice(*attribute, ctx.source);
1175 if segment.is_empty() {
1176 return None;
1177 }
1178 written_fqn.push('.');
1179 written_fqn.push_str(segment);
1180 }
1181 if resolve_fqn_candidates(ctx.python, &written_fqn, |name| {
1182 ctx.graph.index.definitions(name).collect()
1183 })
1184 .into_iter()
1185 .any(|candidate| &candidate == ctx.target)
1186 {
1187 return Some(terminal);
1188 }
1189 }
1190 None
1191}
1192
1193fn imported_root_targets_module(ctx: &ScanCtx<'_>, root: Node<'_>, reference: Node<'_>) -> bool {
1194 imported_module_bindings(ctx, root, reference)
1195 .into_iter()
1196 .any(|binding| {
1197 usage_resolve_module_files(ctx.python, ctx.file, &binding.module)
1198 .into_iter()
1199 .any(|resolved_file| &resolved_file == ctx.target_source)
1200 })
1201}
1202
1203fn module_attribute_target_hit<'a>(node: Node<'a>, ctx: &ScanCtx<'_>) -> Option<Node<'a>> {
1204 let (root, attributes) = attribute_chain(node)?;
1205 if attributes.is_empty() {
1206 return None;
1207 }
1208 for binding in imported_module_bindings(ctx, root, node) {
1209 let mut module_fqn = binding.module;
1210 for attribute in attributes.iter().skip(binding.consumed_attributes) {
1211 let segment = slice(*attribute, ctx.source);
1212 if segment.is_empty() {
1213 return None;
1214 }
1215 if module_fqn.ends_with('.') {
1216 module_fqn.push_str(segment);
1217 } else {
1218 module_fqn.push('.');
1219 module_fqn.push_str(segment);
1220 }
1221 let resolved = usage_resolve_module_files(ctx.python, ctx.file, &module_fqn);
1222 if resolved.is_empty() {
1223 break;
1224 }
1225 if resolved
1226 .iter()
1227 .any(|resolved_file| resolved_file == ctx.target_source)
1228 {
1229 return Some(*attribute);
1230 }
1231 }
1232 }
1233 None
1234}
1235
1236fn call_result_matches_target(call: Node<'_>, ctx: &ScanCtx<'_>) -> bool {
1237 let Some(target_owner) = ctx.target_owner.as_ref() else {
1238 return false;
1239 };
1240 let scope_facts = ctx.scope_facts_for_node(call);
1241 call_result_types(
1242 ctx.graph,
1243 ctx.python,
1244 ctx.file,
1245 ctx.source,
1246 call,
1247 scope_facts,
1248 )
1249 .into_iter()
1250 .any(|class| {
1251 &class == target_owner
1252 || ctx
1253 .graph
1254 .hierarchy
1255 .map(|provider| provider.get_ancestors(&class))
1256 .unwrap_or_default()
1257 .into_iter()
1258 .any(|ancestor| &ancestor == target_owner)
1259 })
1260}
1261
1262pub fn call_result_types(
1263 graph: &PythonGraphSource<'_>,
1264 python: &dyn PythonUsageSource,
1265 file: &ProjectFile,
1266 source: &str,
1267 call: Node<'_>,
1268 scope_facts: Option<&LocalBindingsSnapshot<String>>,
1269) -> Vec<CodeUnit> {
1270 let Some(function) = call.child_by_field_name("function") else {
1271 return Vec::new();
1272 };
1273 let constructed = resolve_constructor_types(graph, python, file, source, function);
1274 if !constructed.is_empty() {
1275 return constructed;
1276 }
1277 let callable_fqns = resolve_callable_fqns(graph, python, file, source, function, scope_facts);
1278 if callable_fqns.is_empty() {
1279 return Vec::new();
1280 }
1281 let callables = callable_fqns
1282 .into_iter()
1283 .flat_map(|callable_fqn| {
1284 resolve_fqn_candidates(python, &callable_fqn, |name| {
1285 graph.index.definitions(name).collect()
1286 })
1287 })
1288 .collect::<Vec<_>>();
1289 let mut classes = Vec::new();
1290 for callable in callables.into_iter().filter(CodeUnit::is_function) {
1291 let Some(raw_type) = callable_return_type_name(graph, python, &callable) else {
1292 continue;
1293 };
1294 if let Some(class) =
1295 resolve_receiver_type(graph, python, callable.source(), &raw_type, true)
1296 {
1297 classes.push(class);
1298 }
1299 }
1300 classes.sort();
1301 classes.dedup();
1302 classes
1303}
1304
1305fn resolve_callable_fqns(
1306 graph: &PythonGraphSource<'_>,
1307 python: &dyn PythonUsageSource,
1308 file: &ProjectFile,
1309 source: &str,
1310 function: Node<'_>,
1311 scope_facts: Option<&LocalBindingsSnapshot<String>>,
1312) -> Vec<String> {
1313 match function.kind() {
1314 "identifier" => {
1315 resolve_identifier_callable_fqns(graph, python, file, source, function, scope_facts)
1316 }
1317 "attribute" => {
1318 resolve_attribute_callable_fqns(graph, python, file, source, function, scope_facts)
1319 }
1320 _ => Vec::new(),
1321 }
1322}
1323
1324fn resolve_identifier_callable_fqns(
1325 graph: &PythonGraphSource<'_>,
1326 python: &dyn PythonUsageSource,
1327 file: &ProjectFile,
1328 source: &str,
1329 function: Node<'_>,
1330 scope_facts: Option<&LocalBindingsSnapshot<String>>,
1331) -> Vec<String> {
1332 let local = slice(function, source);
1333 if local.is_empty() || scope_facts.is_some_and(|facts| facts.is_shadowed(local)) {
1334 return Vec::new();
1335 }
1336 let binder = python.import_binder_of(file);
1337 match binder.bindings.get(local) {
1338 Some(binding) if binding.kind == ImportKind::Named => binding
1339 .imported_name
1340 .as_ref()
1341 .map(|imported| vec![format!("{}.{}", binding.module_specifier, imported)])
1342 .unwrap_or_default(),
1343 _ => graph
1344 .index
1345 .declarations(file)
1346 .into_iter()
1347 .find(|unit| unit.is_function() && unit.identifier() == local)
1348 .map(|unit| vec![unit.fq_name()])
1349 .unwrap_or_default(),
1350 }
1351}
1352
1353fn resolve_attribute_callable_fqns(
1354 graph: &PythonGraphSource<'_>,
1355 python: &dyn PythonUsageSource,
1356 file: &ProjectFile,
1357 source: &str,
1358 function: Node<'_>,
1359 scope_facts: Option<&LocalBindingsSnapshot<String>>,
1360) -> Vec<String> {
1361 let Some(receiver) = function.child_by_field_name("object") else {
1362 return Vec::new();
1363 };
1364 let Some(method) = function.child_by_field_name("attribute") else {
1365 return Vec::new();
1366 };
1367 let method = slice(method, source);
1368 if method.is_empty() {
1369 return Vec::new();
1370 }
1371 let mut fqns = attribute_receiver_classes(graph, python, file, source, receiver, scope_facts)
1372 .into_iter()
1373 .map(|class| format!("{}.{}", class.fq_name(), method))
1374 .collect::<Vec<_>>();
1375 fqns.sort();
1376 fqns.dedup();
1377 fqns
1378}
1379
1380fn attribute_receiver_classes(
1381 graph: &PythonGraphSource<'_>,
1382 python: &dyn PythonUsageSource,
1383 file: &ProjectFile,
1384 source: &str,
1385 receiver: Node<'_>,
1386 scope_facts: Option<&LocalBindingsSnapshot<String>>,
1387) -> Vec<CodeUnit> {
1388 let mut classes = match receiver.kind() {
1389 "identifier" => {
1390 identifier_receiver_classes(graph, python, file, source, receiver, scope_facts)
1391 }
1392 "attribute" => {
1393 if let Some(root) = leftmost_identifier(receiver)
1394 && scope_facts.is_some_and(|facts| facts.is_shadowed(slice(root, source)))
1395 {
1396 Vec::new()
1397 } else {
1398 resolve_constructor_types(graph, python, file, source, receiver)
1399 }
1400 }
1401 _ => Vec::new(),
1402 };
1403 classes.sort();
1404 classes.dedup();
1405 classes
1406}
1407
1408fn identifier_receiver_classes(
1409 graph: &PythonGraphSource<'_>,
1410 python: &dyn PythonUsageSource,
1411 file: &ProjectFile,
1412 source: &str,
1413 receiver: Node<'_>,
1414 scope_facts: Option<&LocalBindingsSnapshot<String>>,
1415) -> Vec<CodeUnit> {
1416 let ident = slice(receiver, source);
1417 if ident.is_empty() {
1418 return Vec::new();
1419 }
1420 if matches!(ident, "self" | "cls")
1421 && let Some(class) = enclosing_class_for_node(graph, file, receiver)
1422 {
1423 return vec![class];
1424 }
1425 if let Some(facts) = scope_facts {
1426 if let Some(raw_type) = facts
1427 .resolution_for(ident)
1428 .as_precise()
1429 .and_then(|targets| targets.iter().next())
1430 && let Some(class) = resolve_receiver_type(graph, python, file, raw_type, false)
1431 {
1432 return vec![class];
1433 }
1434 if facts.is_shadowed(ident) {
1435 return Vec::new();
1436 }
1437 }
1438 resolve_receiver_type(graph, python, file, ident, false)
1439 .into_iter()
1440 .collect()
1441}
1442
1443fn enclosing_class_for_node(
1444 graph: &PythonGraphSource<'_>,
1445 file: &ProjectFile,
1446 node: Node<'_>,
1447) -> Option<CodeUnit> {
1448 let range = Range {
1449 start_byte: node.start_byte(),
1450 end_byte: node.end_byte(),
1451 start_line: 0,
1452 end_line: 0,
1453 };
1454 let enclosing = graph.index.enclosing_code_unit(file, &range)?;
1455 brokk_bifrost_core::analyzer::usages::common::enclosing_owner_chain(enclosing, |unit| {
1461 graph.index.parent_of(unit)
1462 })
1463 .take(2)
1464 .find(|unit| unit.is_class() && unit.source() == file)
1465}
1466
1467fn attribute_chain<'a>(node: Node<'a>) -> Option<(Node<'a>, Vec<Node<'a>>)> {
1468 let mut attributes = Vec::new();
1469 let mut current = node;
1470 loop {
1471 if current.kind() != "attribute" {
1472 return None;
1473 }
1474 attributes.push(current.child_by_field_name("attribute")?);
1475 current = current.child_by_field_name("object")?;
1476 if current.kind() == "identifier" {
1477 attributes.reverse();
1478 return Some((current, attributes));
1479 }
1480 }
1481}
1482
1483struct ImportedModuleBinding {
1484 module: String,
1485 consumed_attributes: usize,
1486}
1487
1488fn imported_module_bindings(
1489 ctx: &ScanCtx<'_>,
1490 root: Node<'_>,
1491 reference: Node<'_>,
1492) -> Vec<ImportedModuleBinding> {
1493 let root_text = slice(root, ctx.source);
1494 if root_text.is_empty() || import_root_shadowed(ctx, root_text, root, reference) {
1495 return Vec::new();
1496 }
1497
1498 if let Some(binding) = ctx.scoped_import_bindings.iter().rev().find(|binding| {
1499 binding.is_function_scoped()
1500 && binding.start_byte <= reference.start_byte()
1501 && binding.scope_start_byte <= reference.start_byte()
1502 && reference.end_byte() <= binding.scope_end_byte
1503 && binding.local_name == root_text
1504 }) {
1505 return if usage_resolve_module_files(ctx.python, ctx.file, &binding.qualified_name)
1506 .is_empty()
1507 {
1508 Vec::new()
1509 } else {
1510 vec![ImportedModuleBinding {
1511 module: binding.qualified_name.clone(),
1512 consumed_attributes: binding.consumed_attributes,
1513 }]
1514 };
1515 }
1516
1517 let Some(events) = ctx.raw_module_bindings.get(root_text) else {
1518 return Vec::new();
1519 };
1520 let cutoff = if reference_is_deferred_function_body(reference) {
1521 usize::MAX
1522 } else {
1523 reference.start_byte()
1524 };
1525 let visible: Vec<_> = events
1526 .iter()
1527 .filter(|event| event.visible_from <= cutoff)
1528 .collect();
1529 let start = visible
1530 .iter()
1531 .rposition(|event| !event.conditional)
1532 .unwrap_or(0);
1533 let mut modules = visible[start..]
1534 .iter()
1535 .filter_map(|event| match &event.kind {
1536 ModuleBindingEventKind::ImportModule {
1537 module,
1538 consumed_attributes,
1539 } => Some(ImportedModuleBinding {
1540 module: module.clone(),
1541 consumed_attributes: *consumed_attributes,
1542 }),
1543 ModuleBindingEventKind::FromImport {
1544 module,
1545 imported_name,
1546 } => {
1547 let submodule = if module.ends_with('.') {
1548 format!("{module}{imported_name}")
1549 } else {
1550 format!("{module}.{imported_name}")
1551 };
1552 (!usage_resolve_module_files(ctx.python, ctx.file, &submodule).is_empty())
1553 .then_some(ImportedModuleBinding {
1554 module: submodule,
1555 consumed_attributes: 0,
1556 })
1557 }
1558 ModuleBindingEventKind::Other => None,
1559 })
1560 .collect::<Vec<_>>();
1561 modules.sort_by(|left, right| {
1562 left.module
1563 .cmp(&right.module)
1564 .then_with(|| left.consumed_attributes.cmp(&right.consumed_attributes))
1565 });
1566 modules.dedup_by(|left, right| {
1567 left.module == right.module && left.consumed_attributes == right.consumed_attributes
1568 });
1569 modules
1570}
1571
1572fn import_root_shadowed(
1573 ctx: &ScanCtx<'_>,
1574 root_text: &str,
1575 root: Node<'_>,
1576 reference: Node<'_>,
1577) -> bool {
1578 ctx.scope_entry_for_node(root)
1579 .or_else(|| ctx.scope_entry_for_node(reference))
1580 .is_some_and(|(scope, facts)| !scope.is_module() && facts.is_shadowed(root_text))
1581 || enclosing_parameters_shadow(root_text, reference, ctx.source)
1582}
1583
1584fn enclosing_parameters_shadow(root_text: &str, reference: Node<'_>, source: &str) -> bool {
1585 let mut current = reference;
1586 while let Some(parent) = current.parent() {
1587 if matches!(parent.kind(), "function_definition" | "lambda") {
1588 let Some(parameters) = parent.child_by_field_name("parameters") else {
1589 return false;
1590 };
1591 let mut cursor = parameters.walk();
1592 return parameters.named_children(&mut cursor).any(|parameter| {
1593 parameter_symbol(parameter, source).as_deref() == Some(root_text)
1594 });
1595 }
1596 current = parent;
1597 }
1598 false
1599}
1600
1601enum EnclosingParameterType {
1602 NotDeclared,
1603 Untyped,
1604 Typed(String),
1605}
1606
1607fn enclosing_runtime_parameter_type(
1608 name: &str,
1609 reference: Node<'_>,
1610 source: &str,
1611) -> EnclosingParameterType {
1612 let site_start = reference.start_byte();
1613 let site_end = reference.end_byte();
1614 let mut current = reference;
1615 while let Some(parent) = current.parent() {
1616 if matches!(parent.kind(), "function_definition" | "lambda")
1617 && parent
1618 .child_by_field_name("body")
1619 .is_some_and(|body| body.start_byte() <= site_start && site_end <= body.end_byte())
1620 && let Some(parameters) = parent.child_by_field_name("parameters")
1621 {
1622 let mut cursor = parameters.walk();
1623 for parameter in parameters.named_children(&mut cursor) {
1624 if parameter_symbol(parameter, source).as_deref() != Some(name) {
1625 continue;
1626 }
1627 return parameter
1628 .child_by_field_name("type")
1629 .and_then(|annotation| normalized_receiver_type(slice(annotation, source)))
1630 .map_or(EnclosingParameterType::Untyped, |raw_type| {
1631 EnclosingParameterType::Typed(raw_type)
1632 });
1633 }
1634 }
1635 current = parent;
1636 }
1637 EnclosingParameterType::NotDeclared
1638}
1639
1640fn member_receiver_match_is_unproven(
1641 object: Node<'_>,
1642 object_text: &str,
1643 node: Node<'_>,
1644 ctx: &ScanCtx<'_>,
1645) -> bool {
1646 if matches!(object_text, "self" | "cls") {
1647 return false;
1648 }
1649 match object.kind() {
1650 "identifier" => {
1651 ctx.receiver_type_is_unknown(object_text, node) && !ctx.member_best_effort_unique
1652 }
1653 "attribute" => true,
1654 _ => false,
1655 }
1656}
1657
1658pub fn slice<'a>(node: Node<'_>, source: &'a str) -> &'a str {
1659 brokk_bifrost_core::analyzer::common::node_source_text(node, source)
1660}
1661
1662fn is_call_callee(node: Node<'_>) -> bool {
1664 node.parent().is_some_and(|parent| {
1665 parent.kind() == "call"
1666 && parent
1667 .child_by_field_name("function")
1668 .is_some_and(|function| function.id() == node.id())
1669 })
1670}
1671
1672pub fn is_declaration_identifier(node: Node<'_>) -> bool {
1673 let Some(parent) = node.parent() else {
1674 return false;
1675 };
1676 let contains = |container: Node<'_>| {
1677 container.start_byte() <= node.start_byte() && node.end_byte() <= container.end_byte()
1678 };
1679 match parent.kind() {
1680 "class_definition" | "function_definition" => parent
1681 .child_by_field_name("name")
1682 .is_some_and(|name| name.id() == node.id()),
1683 "parameters" | "lambda_parameters" | "list_splat_pattern" | "dictionary_splat_pattern" => {
1684 true
1685 }
1686 "default_parameter" | "typed_parameter" | "typed_default_parameter" => {
1687 parent.child_by_field_name("name").is_some_and(contains)
1688 }
1689 "assignment" | "augmented_assignment" | "for_statement" | "for_in_clause" => {
1690 parent.child_by_field_name("left").is_some_and(contains)
1691 }
1692 "named_expression" => parent.child_by_field_name("name").is_some_and(contains),
1693 "aliased_import" | "import_from_statement" | "import_statement" => true,
1694 _ => false,
1695 }
1696}
1697
1698#[derive(Clone, Copy, Debug, PartialEq, Eq)]
1699enum ModuleBindingKind {
1700 TargetSymbolImport,
1701 TargetModuleImport,
1702 Other,
1703}
1704
1705#[derive(Clone, Copy, Debug)]
1706struct ClassifiedModuleBindingEvent {
1707 visible_from: usize,
1708 conditional: bool,
1709 kind: ModuleBindingKind,
1710}
1711
1712pub fn collect_module_binding_timeline(root: Node<'_>, source: &str) -> ModuleBindingTimeline {
1713 let mut timeline = ModuleBindingTimeline::default();
1714 let mut stack = vec![root];
1715 while let Some(node) = stack.pop() {
1716 match node.kind() {
1717 "function_definition" | "class_definition" => {
1718 if let Some(name) = node.child_by_field_name("name") {
1719 record_module_binding(
1720 &mut timeline,
1721 slice(name, source),
1722 node.end_byte(),
1723 binding_is_conditional(node),
1724 ModuleBindingEventKind::Other,
1725 );
1726 }
1727 continue;
1728 }
1729 "import_statement" | "import_from_statement" => {
1730 collect_import_binding_events(node, source, &mut timeline);
1731 continue;
1732 }
1733 "assignment" | "augmented_assignment" | "named_expression" => {
1734 if let Some(left) = node.child_by_field_name("left") {
1735 record_local_binding_targets(
1736 left,
1737 source,
1738 node.end_byte(),
1739 binding_is_conditional(node),
1740 &mut timeline,
1741 );
1742 }
1743 continue;
1744 }
1745 "for_statement" => {
1746 if let Some(left) = node.child_by_field_name("left") {
1747 record_local_binding_targets(
1748 left,
1749 source,
1750 left.end_byte(),
1751 true,
1752 &mut timeline,
1753 );
1754 }
1755 }
1756 _ => {}
1757 }
1758
1759 let mut cursor = node.walk();
1760 let mut children: Vec<Node<'_>> = node.named_children(&mut cursor).collect();
1761 children.reverse();
1762 stack.extend(children);
1763 }
1764 for events in timeline.values_mut() {
1765 events.sort_by_key(|event| event.visible_from);
1766 }
1767 timeline
1768}
1769
1770fn collect_import_binding_events(
1771 node: Node<'_>,
1772 source: &str,
1773 timeline: &mut ModuleBindingTimeline,
1774) {
1775 if node.kind() == "import_statement" {
1776 let mut cursor = node.walk();
1777 for imported in node.children_by_field_name("name", &mut cursor) {
1778 let name = imported.child_by_field_name("name").unwrap_or(imported);
1779 let Some(local) = imported
1780 .child_by_field_name("alias")
1781 .or_else(|| first_identifier(name))
1782 else {
1783 continue;
1784 };
1785 let module = slice(name, source).trim();
1786 let consumed_attributes = if imported.child_by_field_name("alias").is_some() {
1787 0
1788 } else {
1789 parse_symbol_path(Language::Python, module)
1790 .len()
1791 .saturating_sub(1)
1792 };
1793 record_module_binding(
1794 timeline,
1795 slice(local, source),
1796 node.end_byte(),
1797 binding_is_conditional(node),
1798 ModuleBindingEventKind::ImportModule {
1799 module: module.to_string(),
1800 consumed_attributes,
1801 },
1802 );
1803 }
1804 return;
1805 }
1806
1807 let Some(module_node) = node.child_by_field_name("module_name") else {
1808 return;
1809 };
1810 let module = slice(module_node, source).trim();
1811 let mut cursor = node.walk();
1812 for imported in node.children_by_field_name("name", &mut cursor) {
1813 if imported.kind() == "wildcard_import" {
1814 continue;
1815 }
1816 let name = imported.child_by_field_name("name").unwrap_or(imported);
1817 let Some(imported_identifier) = last_identifier(name) else {
1818 continue;
1819 };
1820 let imported_name = slice(imported_identifier, source).trim();
1821 let Some(local) = imported
1822 .child_by_field_name("alias")
1823 .or_else(|| last_identifier(name))
1824 else {
1825 continue;
1826 };
1827 record_module_binding(
1828 timeline,
1829 slice(local, source),
1830 node.end_byte(),
1831 binding_is_conditional(node),
1832 ModuleBindingEventKind::FromImport {
1833 module: module.to_string(),
1834 imported_name: imported_name.to_string(),
1835 },
1836 );
1837 }
1838}
1839
1840fn classify_module_binding_timeline(
1841 python: &dyn PythonUsageSource,
1842 file: &ProjectFile,
1843 timeline: &ModuleBindingTimeline,
1844 seeds: &BTreeSet<(ProjectFile, String)>,
1845 edges: &[ImportEdge],
1846) -> HashMap<String, Vec<ClassifiedModuleBindingEvent>> {
1847 let mut classified = HashMap::default();
1848 let mut module_targets: HashMap<String, bool> = HashMap::default();
1849 let relevant_locals: HashSet<&str> =
1850 edges.iter().map(|edge| edge.local_name.as_str()).collect();
1851 for (local, events) in timeline {
1852 if !relevant_locals.contains(local.as_str()) {
1853 continue;
1854 }
1855 let classified_events = events
1856 .iter()
1857 .map(|event| {
1858 let kind = match &event.kind {
1859 ModuleBindingEventKind::ImportModule { module, .. } => {
1860 if *module_targets
1861 .entry(module.clone())
1862 .or_insert_with(|| module_contains_seed(python, file, module, seeds))
1863 {
1864 ModuleBindingKind::TargetModuleImport
1865 } else {
1866 ModuleBindingKind::Other
1867 }
1868 }
1869 ModuleBindingEventKind::FromImport {
1870 module,
1871 imported_name,
1872 } => {
1873 let direct = usage_resolve_module_files(python, file, module).iter().any(
1874 |resolved| seeds.contains(&(resolved.clone(), imported_name.clone())),
1875 );
1876 let submodule = if module.ends_with('.') {
1877 format!("{module}{imported_name}")
1878 } else {
1879 format!("{module}.{imported_name}")
1880 };
1881 let imports_target_module =
1882 *module_targets.entry(submodule.clone()).or_insert_with(|| {
1883 module_contains_seed(python, file, &submodule, seeds)
1884 });
1885 if direct {
1886 ModuleBindingKind::TargetSymbolImport
1887 } else if imports_target_module {
1888 ModuleBindingKind::TargetModuleImport
1889 } else {
1890 ModuleBindingKind::Other
1891 }
1892 }
1893 ModuleBindingEventKind::Other => ModuleBindingKind::Other,
1894 };
1895 ClassifiedModuleBindingEvent {
1896 visible_from: event.visible_from,
1897 conditional: event.conditional,
1898 kind,
1899 }
1900 })
1901 .collect();
1902 classified.insert(local.clone(), classified_events);
1903 }
1904 classified
1905}
1906
1907fn module_contains_seed(
1908 python: &dyn PythonUsageSource,
1909 file: &ProjectFile,
1910 module: &str,
1911 seeds: &BTreeSet<(ProjectFile, String)>,
1912) -> bool {
1913 usage_resolve_module_files(python, file, module)
1914 .iter()
1915 .any(|resolved| seeds.iter().any(|(seed_file, _)| seed_file == resolved))
1916}
1917
1918fn record_module_binding(
1919 timeline: &mut ModuleBindingTimeline,
1920 name: &str,
1921 visible_from: usize,
1922 conditional: bool,
1923 kind: ModuleBindingEventKind,
1924) {
1925 let name = name.trim();
1926 if name.is_empty() {
1927 return;
1928 }
1929 timeline
1930 .entry(name.to_string())
1931 .or_default()
1932 .push(ModuleBindingEvent {
1933 visible_from,
1934 conditional,
1935 kind,
1936 });
1937}
1938
1939fn record_local_binding_targets(
1940 target: Node<'_>,
1941 source: &str,
1942 visible_from: usize,
1943 conditional: bool,
1944 timeline: &mut ModuleBindingTimeline,
1945) {
1946 let mut stack = vec![target];
1947 while let Some(node) = stack.pop() {
1948 if node.kind() == "identifier" {
1949 record_module_binding(
1950 timeline,
1951 slice(node, source),
1952 visible_from,
1953 conditional,
1954 ModuleBindingEventKind::Other,
1955 );
1956 continue;
1957 }
1958 if matches!(node.kind(), "attribute" | "subscript") {
1959 continue;
1960 }
1961 let mut cursor = node.walk();
1962 let mut children: Vec<Node<'_>> = node.named_children(&mut cursor).collect();
1963 children.reverse();
1964 stack.extend(children);
1965 }
1966}
1967
1968fn binding_is_conditional(mut node: Node<'_>) -> bool {
1969 while let Some(parent) = node.parent() {
1970 if matches!(
1971 parent.kind(),
1972 "if_statement"
1973 | "try_statement"
1974 | "except_clause"
1975 | "match_statement"
1976 | "case_clause"
1977 | "for_statement"
1978 | "while_statement"
1979 ) {
1980 return true;
1981 }
1982 if matches!(
1983 parent.kind(),
1984 "module" | "function_definition" | "class_definition"
1985 ) {
1986 return false;
1987 }
1988 node = parent;
1989 }
1990 false
1991}
1992
1993fn first_identifier(node: Node<'_>) -> Option<Node<'_>> {
1994 identifier_extreme(node, false)
1995}
1996
1997fn last_identifier(node: Node<'_>) -> Option<Node<'_>> {
1998 identifier_extreme(node, true)
1999}
2000
2001fn identifier_extreme(node: Node<'_>, last: bool) -> Option<Node<'_>> {
2002 let mut best = None;
2003 let mut stack = vec![node];
2004 while let Some(node) = stack.pop() {
2005 if node.kind() == "identifier" {
2006 if best.is_none_or(|current: Node<'_>| {
2007 if last {
2008 node.start_byte() > current.start_byte()
2009 } else {
2010 node.start_byte() < current.start_byte()
2011 }
2012 }) {
2013 best = Some(node);
2014 }
2015 continue;
2016 }
2017 let mut cursor = node.walk();
2018 stack.extend(node.named_children(&mut cursor));
2019 }
2020 best
2021}
2022
2023fn reference_is_deferred_function_body(node: Node<'_>) -> bool {
2024 let site_start = node.start_byte();
2025 let site_end = node.end_byte();
2026 let mut current = node;
2027 while let Some(parent) = current.parent() {
2028 if matches!(parent.kind(), "function_definition" | "lambda")
2029 && parent
2030 .child_by_field_name("body")
2031 .is_some_and(|body| body.start_byte() <= site_start && site_end <= body.end_byte())
2032 {
2033 return true;
2034 }
2035 current = parent;
2036 }
2037 false
2038}
2039
2040pub fn collect_assigned_identifiers(node: Node<'_>, source: &str, out: &mut HashSet<String>) {
2041 let mut stack = vec![node];
2042 while let Some(node) = stack.pop() {
2043 if node.kind() == "identifier" {
2044 let text = slice(node, source).trim();
2045 if !text.is_empty() {
2046 out.insert(text.to_string());
2047 }
2048 continue;
2049 }
2050
2051 let mut cursor = node.walk();
2052 let mut children: Vec<Node<'_>> = node.named_children(&mut cursor).collect();
2053 children.reverse();
2054 stack.extend(children);
2055 }
2056}
2057
2058pub fn collect_scope_facts_from_parsed_source(
2059 graph: &PythonGraphSource<'_>,
2060 python: &dyn PythonUsageSource,
2061 file: &ProjectFile,
2062 source: &str,
2063 root: Node<'_>,
2064) -> PythonScopeFacts {
2065 let mut factory_return_types = collect_factory_return_types_from_root(root, source);
2066 collect_imported_factory_return_types(graph, python, file, &mut factory_return_types);
2067 collect_scope_facts_with_factory_returns(graph, file, source, &factory_return_types)
2068}
2069
2070fn collect_imported_factory_return_types(
2071 graph: &PythonGraphSource<'_>,
2072 python: &dyn PythonUsageSource,
2073 file: &ProjectFile,
2074 factory_return_types: &mut HashMap<String, String>,
2075) {
2076 let binder = python.import_binder_of(file);
2077 for (local, binding) in &binder.bindings {
2078 if !matches!(binding.kind, ImportKind::Named) {
2079 continue;
2080 }
2081 let Some(imported) = binding.imported_name.as_deref() else {
2082 continue;
2083 };
2084 let fqn = format!("{}.{}", binding.module_specifier, imported);
2085 let units =
2086 resolve_fqn_candidates(python, &fqn, |name| graph.index.definitions(name).collect());
2087 for unit in units {
2088 if unit.is_function() {
2089 if let Some(return_type) = callable_return_type_name(graph, python, &unit) {
2090 factory_return_types
2091 .entry(local.clone())
2092 .or_insert(return_type);
2093 }
2094 continue;
2095 }
2096 if !unit.is_class() {
2097 continue;
2098 }
2099 factory_return_types
2100 .entry(local.clone())
2101 .or_insert_with(|| unit.identifier().to_string());
2102 collect_imported_class_method_return_types(
2103 graph,
2104 python,
2105 local,
2106 &unit,
2107 factory_return_types,
2108 );
2109 }
2110 }
2111}
2112
2113fn collect_imported_class_method_return_types(
2114 graph: &PythonGraphSource<'_>,
2115 python: &dyn PythonSource,
2116 local_class_name: &str,
2117 class_unit: &CodeUnit,
2118 factory_return_types: &mut HashMap<String, String>,
2119) {
2120 for member in graph.index.direct_children(class_unit) {
2121 if !member.is_function() {
2122 continue;
2123 }
2124 let Some(return_type) = callable_return_type_name(graph, python, &member) else {
2125 continue;
2126 };
2127 factory_return_types
2128 .entry(format!("{}.{}", local_class_name, member.identifier()))
2129 .or_insert(return_type);
2130 }
2131}
2132
2133fn callable_return_type_name(
2134 graph: &PythonGraphSource<'_>,
2135 python: &dyn PythonSource,
2136 callable: &CodeUnit,
2137) -> Option<String> {
2138 if let Some(prepared) = python.prepared_syntax(callable.source()) {
2143 #[cfg(any(test, feature = "test-support"))]
2144 note_callable_return_type_lookup_for_test(true);
2145 return callable_return_type_name_in_tree(
2146 graph,
2147 callable,
2148 prepared.source(),
2149 prepared.tree().root_node(),
2150 );
2151 }
2152 #[cfg(any(test, feature = "test-support"))]
2153 note_callable_return_type_lookup_for_test(false);
2154 let source = graph.index.indexed_source(callable.source())?;
2155 declaration_source_slices(graph, callable, &source)
2156 .into_iter()
2157 .find_map(|declaration_source| {
2158 let mut parser = Parser::new();
2159 parser
2160 .set_language(&tree_sitter_python::LANGUAGE.into())
2161 .ok()?;
2162 let tree = parser.parse(declaration_source, None)?;
2163 let function = first_function_definition(tree.root_node())?;
2164 factory_return_type(function, declaration_source)
2165 })
2166}
2167
2168#[cfg(any(test, feature = "test-support"))]
2174#[derive(Debug, Default, Clone, Copy, PartialEq, Eq)]
2175pub struct CallableReturnTypeLookupCounts {
2176 pub prepared: usize,
2178 pub reparsed: usize,
2180}
2181
2182#[cfg(any(test, feature = "test-support"))]
2183thread_local! {
2184 static CALLABLE_RETURN_TYPE_LOOKUPS_FOR_TEST: std::cell::Cell<CallableReturnTypeLookupCounts> =
2185 const { std::cell::Cell::new(CallableReturnTypeLookupCounts { prepared: 0, reparsed: 0 }) };
2186}
2187
2188#[cfg(any(test, feature = "test-support"))]
2189fn note_callable_return_type_lookup_for_test(from_prepared_syntax: bool) {
2190 CALLABLE_RETURN_TYPE_LOOKUPS_FOR_TEST.with(|counts| {
2191 let mut observed = counts.get();
2192 if from_prepared_syntax {
2193 observed.prepared += 1;
2194 } else {
2195 observed.reparsed += 1;
2196 }
2197 counts.set(observed);
2198 });
2199}
2200
2201#[cfg(any(test, feature = "test-support"))]
2204pub fn with_callable_return_type_lookup_counter_for_test<T>(
2205 body: impl FnOnce() -> T,
2206) -> (T, CallableReturnTypeLookupCounts) {
2207 CALLABLE_RETURN_TYPE_LOOKUPS_FOR_TEST.with(|counts| {
2208 counts.set(CallableReturnTypeLookupCounts::default());
2209 let result = body();
2210 let observed = counts.get();
2211 counts.set(CallableReturnTypeLookupCounts::default());
2212 (result, observed)
2213 })
2214}
2215
2216fn callable_return_type_name_in_tree(
2222 graph: &PythonGraphSource<'_>,
2223 callable: &CodeUnit,
2224 source: &str,
2225 root: Node<'_>,
2226) -> Option<String> {
2227 let mut ranges = graph.index.ranges(callable);
2228 ranges.sort_by_key(|range| range.start_byte);
2229 ranges.into_iter().find_map(|range| {
2230 let declaration = root.descendant_for_byte_range(range.start_byte, range.end_byte)?;
2231 let function = first_function_definition(declaration)?;
2232 factory_return_type(function, source)
2233 })
2234}
2235
2236fn first_function_definition(root: Node<'_>) -> Option<Node<'_>> {
2237 let mut stack = vec![root];
2238 while let Some(node) = stack.pop() {
2239 if node.kind() == "function_definition" {
2240 return Some(node);
2241 }
2242 let mut cursor = node.walk();
2243 let mut children: Vec<Node<'_>> = node.named_children(&mut cursor).collect();
2244 children.reverse();
2245 stack.extend(children);
2246 }
2247 None
2248}
2249
2250fn collect_scope_facts_with_factory_returns(
2251 graph: &PythonGraphSource<'_>,
2252 file: &ProjectFile,
2253 source: &str,
2254 factory_return_types: &HashMap<String, String>,
2255) -> PythonScopeFacts {
2256 let declarations = graph.index.declarations(file);
2257 let mut class_facts_by_name: HashMap<String, LocalBindingsSnapshot<String>> =
2258 HashMap::default();
2259 for declaration in declarations
2260 .iter()
2261 .filter(|declaration| declaration.is_class())
2262 {
2263 let Some(declaration_source) = declaration_source(graph, declaration, source) else {
2264 continue;
2265 };
2266 let facts = collect_scope_facts_from_source(
2267 &declaration_source,
2268 ScopeFactTraversal::Class,
2269 true,
2270 Some(declaration.short_name()),
2271 factory_return_types,
2272 );
2273 class_facts_by_name.insert(
2274 declaration.short_name().to_string(),
2275 facts.filtered_visible_bindings(|symbol, _| symbol.starts_with("self.")),
2276 );
2277 }
2278
2279 let mut scope_facts = HashMap::default();
2280 for declaration in declarations
2281 .iter()
2282 .filter(|declaration| declaration.is_function())
2283 {
2284 let Some(declaration_source) = declaration_source(graph, declaration, source) else {
2285 continue;
2286 };
2287 let owner = declaration
2292 .short_name()
2293 .rsplit_once('.')
2294 .map(|(owner, _)| owner);
2295 let mut facts = collect_scope_facts_from_source(
2296 &declaration_source,
2297 ScopeFactTraversal::Function,
2298 false,
2299 owner,
2300 factory_return_types,
2301 );
2302 if let Some(owner) = owner
2303 && let Some(class_facts) = class_facts_by_name.get(owner)
2304 {
2305 facts = facts.merged_with_visible(class_facts);
2306 }
2307 scope_facts.insert(declaration.clone(), facts);
2308 }
2309
2310 for declaration in declarations.iter().filter(|d| d.is_module()) {
2315 let Some(declaration_source) = declaration_source(graph, declaration, source) else {
2316 continue;
2317 };
2318 let facts = collect_scope_facts_from_source(
2319 &declaration_source,
2320 ScopeFactTraversal::Module,
2321 false,
2322 None,
2323 factory_return_types,
2324 );
2325 scope_facts.insert(declaration.clone(), facts);
2326 }
2327 scope_facts
2328}
2329
2330fn declaration_source(
2331 graph: &PythonGraphSource<'_>,
2332 declaration: &CodeUnit,
2333 file_source: &str,
2334) -> Option<String> {
2335 let slices = declaration_source_slices(graph, declaration, file_source);
2336 (!slices.is_empty()).then(|| slices.join("\n\n"))
2337}
2338
2339fn declaration_source_slices<'a>(
2340 graph: &PythonGraphSource<'_>,
2341 declaration: &CodeUnit,
2342 file_source: &'a str,
2343) -> Vec<&'a str> {
2344 let mut ranges = graph.index.ranges(declaration);
2345 ranges.sort_by_key(|range| range.start_byte);
2346 ranges
2347 .into_iter()
2348 .filter_map(|range| file_source.get(range.start_byte..range.end_byte))
2349 .collect()
2350}
2351
2352fn collect_scope_facts_from_source(
2353 source: &str,
2354 traversal: ScopeFactTraversal,
2355 allow_self_receivers: bool,
2356 current_class: Option<&str>,
2357 factory_return_types: &HashMap<String, String>,
2358) -> LocalBindingsSnapshot<String> {
2359 let events = collect_scope_fact_events(source, traversal);
2360 collect_scope_facts_from_events(
2361 &events,
2362 allow_self_receivers,
2363 current_class,
2364 factory_return_types,
2365 )
2366}
2367
2368pub fn collect_function_scope_facts_from_node(
2369 function: Node<'_>,
2370 source: &str,
2371) -> LocalBindingsSnapshot<String> {
2372 let mut events = Vec::new();
2373 if function.kind() == "lambda" {
2374 if let Some(parameters) = function.child_by_field_name("parameters") {
2375 collect_parameter_events(parameters, source, &mut events);
2376 }
2377 } else {
2378 collect_scope_fact_events_from_node(
2379 function,
2380 source,
2381 ScopeFactTraversal::Function,
2382 &mut events,
2383 );
2384 }
2385 collect_scope_facts_from_events(&events, false, None, &HashMap::default())
2386}
2387
2388fn collect_scope_facts_from_events(
2389 events: &[ScopeFactEvent],
2390 allow_self_receivers: bool,
2391 current_class: Option<&str>,
2392 factory_return_types: &HashMap<String, String>,
2393) -> LocalBindingsSnapshot<String> {
2394 let mut engine = LocalInferenceEngine::new(LocalInferenceConfig::default());
2395 let globals: HashSet<&str> = events
2396 .iter()
2397 .filter_map(|event| match event {
2398 ScopeFactEvent::Global { symbol } => Some(symbol.as_str()),
2399 _ => None,
2400 })
2401 .collect();
2402 let nonlocals: HashSet<&str> = events
2403 .iter()
2404 .filter_map(|event| match event {
2405 ScopeFactEvent::Nonlocal { symbol } => Some(symbol.as_str()),
2406 _ => None,
2407 })
2408 .collect();
2409 for symbol in &nonlocals {
2410 engine.declare_shadow((*symbol).to_string());
2411 }
2412 for event in events {
2413 if let ScopeFactEvent::Parameter { symbol, .. } = event
2414 && !globals.contains(symbol.as_str())
2415 && !nonlocals.contains(symbol.as_str())
2416 && !engine.is_shadowed(symbol)
2417 {
2418 engine.declare_shadow(symbol.clone());
2419 }
2420 }
2421
2422 let mut changed = true;
2423 while changed {
2424 changed = false;
2425 let mut aliases = Vec::new();
2426 for event in events {
2427 match event {
2428 ScopeFactEvent::Parameter {
2429 symbol,
2430 annotation: Some(annotation),
2431 }
2432 | ScopeFactEvent::Annotation { symbol, annotation } => {
2433 if globals.contains(symbol.as_str()) || nonlocals.contains(symbol.as_str()) {
2434 continue;
2435 }
2436 apply_annotation_event(
2437 symbol,
2438 annotation,
2439 allow_self_receivers,
2440 &mut engine,
2441 &mut changed,
2442 );
2443 }
2444 ScopeFactEvent::Parameter {
2445 annotation: None, ..
2446 } => {}
2447 ScopeFactEvent::Assignment { lhs, rhs } => {
2448 if globals.contains(lhs.as_str()) {
2449 continue;
2450 }
2451 if !engine.is_shadowed(lhs) {
2452 engine.declare_shadow(lhs.clone());
2453 }
2454 if lhs.starts_with("self.") && !allow_self_receivers {
2455 continue;
2456 }
2457
2458 match rhs {
2459 AssignmentRhs::Call(callee) => {
2460 if !engine.is_shadowed(callee) {
2461 if let Some(receiver_type) = factory_return_type_for_callee(
2462 callee,
2463 current_class,
2464 factory_return_types,
2465 ) && engine.resolve_symbol(lhs).is_unknown()
2466 {
2467 engine.seed_symbol(lhs.clone(), receiver_type.clone());
2468 changed = true;
2469 continue;
2470 }
2471
2472 if let Some(receiver_type) = normalized_receiver_type(callee)
2473 && engine.resolve_symbol(lhs).is_unknown()
2474 {
2475 engine.seed_symbol(lhs.clone(), receiver_type);
2476 changed = true;
2477 continue;
2478 }
2479 }
2480 }
2481 AssignmentRhs::Symbol(rhs_symbol) => {
2482 if !engine.is_shadowed(rhs_symbol)
2483 && let Some(receiver_type) = normalized_receiver_type(rhs_symbol)
2484 && engine.resolve_symbol(lhs).is_unknown()
2485 {
2486 engine.seed_symbol(lhs.clone(), receiver_type);
2487 changed = true;
2488 continue;
2489 }
2490
2491 if let SymbolResolution::Precise(targets) =
2492 engine.resolve_symbol(rhs_symbol)
2493 && !targets.is_empty()
2494 {
2495 aliases.push((lhs.clone(), rhs_symbol.clone()));
2496 }
2497 }
2498 AssignmentRhs::Unknown => {}
2499 }
2500 }
2501 ScopeFactEvent::Global { .. } | ScopeFactEvent::Nonlocal { .. } => {}
2502 }
2503 }
2504 let before = engine.snapshot();
2505 engine.apply_aliases_until_stable(aliases);
2506 if engine.snapshot() != before {
2507 changed = true;
2508 }
2509 }
2510
2511 engine.snapshot()
2512}
2513
2514fn factory_return_type_for_callee<'a>(
2515 callee: &str,
2516 current_class: Option<&str>,
2517 factory_return_types: &'a HashMap<String, String>,
2518) -> Option<&'a String> {
2519 if let Some(receiver_type) = factory_return_types.get(callee) {
2520 return Some(receiver_type);
2521 }
2522 let class_name = current_class?;
2523 let method = callee
2524 .strip_prefix("self.")
2525 .or_else(|| callee.strip_prefix("cls."))?;
2526 factory_return_types.get(&format!("{class_name}.{method}"))
2527}
2528
2529fn apply_annotation_event(
2530 symbol: &str,
2531 annotation: &str,
2532 allow_self_receivers: bool,
2533 engine: &mut LocalInferenceEngine<String>,
2534 changed: &mut bool,
2535) {
2536 if symbol.starts_with("self.") && !allow_self_receivers {
2537 return;
2538 }
2539 if let Some(receiver_type) = normalized_receiver_type(annotation)
2540 && engine.resolve_symbol(symbol).is_unknown()
2541 {
2542 engine.seed_symbol(symbol.to_string(), receiver_type);
2543 *changed = true;
2544 }
2545}
2546
2547enum ScopeFactEvent {
2548 Global {
2549 symbol: String,
2550 },
2551 Nonlocal {
2552 symbol: String,
2553 },
2554 Parameter {
2555 symbol: String,
2556 annotation: Option<String>,
2557 },
2558 Annotation {
2559 symbol: String,
2560 annotation: String,
2561 },
2562 Assignment {
2563 lhs: String,
2564 rhs: AssignmentRhs,
2565 },
2566}
2567
2568enum AssignmentRhs {
2569 Symbol(String),
2570 Call(String),
2571 Unknown,
2572}
2573
2574#[derive(Clone, Copy)]
2575enum ScopeFactTraversal {
2576 Module,
2577 Function,
2578 Class,
2579}
2580
2581fn collect_scope_fact_events(source: &str, traversal: ScopeFactTraversal) -> Vec<ScopeFactEvent> {
2582 if source.trim().is_empty() {
2583 return Vec::new();
2584 }
2585
2586 let mut parser = Parser::new();
2587 if parser
2588 .set_language(&tree_sitter_python::LANGUAGE.into())
2589 .is_err()
2590 {
2591 return Vec::new();
2592 }
2593 let Some(tree) = parser.parse(source, None) else {
2594 return Vec::new();
2595 };
2596
2597 let mut events = Vec::new();
2598 collect_scope_fact_events_from_node(tree.root_node(), source, traversal, &mut events);
2599 events
2600}
2601
2602fn collect_scope_fact_events_from_node(
2603 root: Node<'_>,
2604 source: &str,
2605 traversal: ScopeFactTraversal,
2606 events: &mut Vec<ScopeFactEvent>,
2607) {
2608 let mut stack = vec![(root, false)];
2609 while let Some((node, inside_function)) = stack.pop() {
2610 let next_inside_function = match traversal {
2611 ScopeFactTraversal::Module => {
2612 if matches!(
2613 node.kind(),
2614 "function_definition" | "class_definition" | "lambda"
2615 ) {
2616 continue;
2617 }
2618 false
2619 }
2620 ScopeFactTraversal::Function => match node.kind() {
2621 "function_definition" if inside_function => continue,
2622 "function_definition" => true,
2623 "class_definition" | "lambda" => continue,
2624 _ => inside_function,
2625 },
2626 ScopeFactTraversal::Class => inside_function,
2627 };
2628 if matches!(traversal, ScopeFactTraversal::Function) && !next_inside_function {
2629 let mut cursor = node.walk();
2630 let mut children: Vec<Node<'_>> = node.named_children(&mut cursor).collect();
2631 children.reverse();
2632 stack.extend(children.into_iter().map(|child| (child, false)));
2633 continue;
2634 }
2635 match node.kind() {
2636 "global_statement" => collect_scope_directive_events(node, source, events, |symbol| {
2637 ScopeFactEvent::Global { symbol }
2638 }),
2639 "nonlocal_statement" => {
2640 collect_scope_directive_events(node, source, events, |symbol| {
2641 ScopeFactEvent::Nonlocal { symbol }
2642 })
2643 }
2644 "parameters" | "lambda_parameters" => collect_parameter_events(node, source, events),
2645 "assignment" => collect_assignment_events(node, source, events),
2646 _ => {}
2647 }
2648
2649 let mut cursor = node.walk();
2650 let mut children: Vec<Node<'_>> = node.named_children(&mut cursor).collect();
2651 children.reverse();
2652 stack.extend(
2653 children
2654 .into_iter()
2655 .map(|child| (child, next_inside_function)),
2656 );
2657 }
2658}
2659
2660fn collect_scope_directive_events(
2661 node: Node<'_>,
2662 source: &str,
2663 events: &mut Vec<ScopeFactEvent>,
2664 make_event: impl Fn(String) -> ScopeFactEvent,
2665) {
2666 let mut cursor = node.walk();
2667 for identifier in node
2668 .named_children(&mut cursor)
2669 .filter(|child| child.kind() == "identifier")
2670 {
2671 let Some(symbol) = non_empty_node_text(identifier, source) else {
2672 continue;
2673 };
2674 events.push(make_event(symbol));
2675 }
2676}
2677
2678fn collect_parameter_events(node: Node<'_>, source: &str, events: &mut Vec<ScopeFactEvent>) {
2679 let mut cursor = node.walk();
2680 for child in node.named_children(&mut cursor) {
2681 if child.kind() == "type_parameter" {
2682 continue;
2683 }
2684 let Some(symbol) = parameter_symbol(child, source) else {
2685 continue;
2686 };
2687 if matches!(symbol.as_str(), "self" | "cls" | "/") {
2688 continue;
2689 }
2690 let annotation = child
2691 .child_by_field_name("type")
2692 .map(|annotation| slice(annotation, source).trim().to_string())
2693 .filter(|annotation| !annotation.is_empty());
2694 events.push(ScopeFactEvent::Parameter { symbol, annotation });
2695 }
2696}
2697
2698fn parameter_symbol(node: Node<'_>, source: &str) -> Option<String> {
2699 if node.kind() == "identifier" {
2700 return non_empty_node_text(node, source);
2701 }
2702 if let Some(name) = node.child_by_field_name("name") {
2703 return non_empty_node_text(name, source);
2704 }
2705 let mut cursor = node.walk();
2706 node.named_children(&mut cursor)
2707 .find(|child| child.kind() == "identifier")
2708 .and_then(|identifier| non_empty_node_text(identifier, source))
2709}
2710
2711fn collect_assignment_events(node: Node<'_>, source: &str, events: &mut Vec<ScopeFactEvent>) {
2712 let Some(left) = node.child_by_field_name("left") else {
2713 return;
2714 };
2715 let Some(lhs) = receiver_symbol(left, source) else {
2716 return;
2717 };
2718
2719 if let Some(annotation) = node
2720 .child_by_field_name("type")
2721 .map(|annotation| slice(annotation, source).trim().to_string())
2722 .filter(|annotation| !annotation.is_empty())
2723 {
2724 events.push(ScopeFactEvent::Annotation {
2725 symbol: lhs,
2726 annotation,
2727 });
2728 return;
2729 }
2730
2731 let rhs = node
2732 .child_by_field_name("right")
2733 .and_then(|right| rhs_symbol(right, source))
2734 .unwrap_or(AssignmentRhs::Unknown);
2735 events.push(ScopeFactEvent::Assignment { lhs, rhs });
2736}
2737
2738fn receiver_symbol(node: Node<'_>, source: &str) -> Option<String> {
2739 match node.kind() {
2740 "identifier" | "attribute" => non_empty_node_text(node, source),
2741 _ => None,
2742 }
2743}
2744
2745fn rhs_symbol(node: Node<'_>, source: &str) -> Option<AssignmentRhs> {
2746 match node.kind() {
2747 "identifier" | "attribute" => non_empty_node_text(node, source).map(AssignmentRhs::Symbol),
2748 "call" => node
2749 .child_by_field_name("function")
2750 .or_else(|| node.named_child(0))
2751 .and_then(|callee| receiver_symbol(callee, source))
2752 .map(AssignmentRhs::Call),
2753 _ => None,
2754 }
2755}
2756
2757fn non_empty_node_text(node: Node<'_>, source: &str) -> Option<String> {
2758 let text = slice(node, source).trim();
2759 (!text.is_empty()).then(|| text.to_string())
2760}
2761
2762fn collect_factory_return_types_from_root(root: Node<'_>, source: &str) -> HashMap<String, String> {
2763 let mut returns = HashMap::default();
2764 let mut stack = vec![(root, None::<String>)];
2765 while let Some((node, class_name)) = stack.pop() {
2766 match node.kind() {
2767 "class_definition" => {
2768 let next_class = node
2769 .child_by_field_name("name")
2770 .and_then(|name| non_empty_node_text(name, source))
2771 .or(class_name);
2772 push_factory_index_children(node, next_class, &mut stack);
2773 }
2774 "function_definition" => {
2775 if let Some(name) = node
2776 .child_by_field_name("name")
2777 .and_then(|name| non_empty_node_text(name, source))
2778 && let Some(return_type) = factory_return_type(node, source)
2779 {
2780 let key = class_name
2781 .as_ref()
2782 .map(|class| format!("{class}.{name}"))
2783 .unwrap_or(name);
2784 returns.insert(key, return_type);
2785 }
2786 }
2787 _ => push_factory_index_children(node, class_name, &mut stack),
2788 }
2789 }
2790 returns
2791}
2792
2793fn push_factory_index_children<'tree>(
2794 node: Node<'tree>,
2795 class_name: Option<String>,
2796 stack: &mut Vec<(Node<'tree>, Option<String>)>,
2797) {
2798 let mut cursor = node.walk();
2799 let mut children: Vec<Node<'tree>> = node.named_children(&mut cursor).collect();
2800 children.reverse();
2801 stack.extend(
2802 children
2803 .into_iter()
2804 .map(|child| (child, class_name.clone())),
2805 );
2806}
2807
2808fn factory_return_type(function: Node<'_>, source: &str) -> Option<String> {
2809 if let Some(return_type) = function.child_by_field_name("return_type") {
2810 return receiver_type_from_annotation_node(return_type, source);
2811 }
2812
2813 let body = function.child_by_field_name("body")?;
2814 let mut candidates = HashSet::default();
2815 let mut saw_return = false;
2816 let mut saw_unknown_return = false;
2817 let mut stack = vec![body];
2818 while let Some(node) = stack.pop() {
2819 if node != body && matches!(node.kind(), "function_definition" | "class_definition") {
2820 continue;
2821 }
2822 if node.kind() == "return_statement" {
2823 saw_return = true;
2824 match node
2825 .named_child(0)
2826 .and_then(|value| returned_receiver_type(value, source))
2827 {
2828 Some(returned_type) => {
2829 candidates.insert(returned_type);
2830 }
2831 None => saw_unknown_return = true,
2832 }
2833 }
2834 let mut cursor = node.walk();
2835 let mut children: Vec<Node<'_>> = node.named_children(&mut cursor).collect();
2836 children.reverse();
2837 stack.extend(children);
2838 }
2839 if !saw_return || saw_unknown_return {
2840 return None;
2841 }
2842 (candidates.len() == 1)
2843 .then(|| candidates.into_iter().next())
2844 .flatten()
2845}
2846
2847fn receiver_type_from_annotation_node(annotation: Node<'_>, source: &str) -> Option<String> {
2853 match annotation.kind() {
2854 "type" => receiver_type_from_annotation_node(annotation.named_child(0)?, source),
2855 "identifier" | "attribute" | "member_type" | "string" => {
2856 normalized_receiver_type(slice(annotation, source).trim())
2857 }
2858 "generic_type" => {
2859 let base = annotation.named_child(0)?;
2860 if optional_annotation_wrapper(base, source) {
2861 let parameter = annotation.named_child(1)?;
2862 return receiver_type_from_annotation_node(parameter.named_child(0)?, source);
2863 }
2864 receiver_type_from_annotation_node(base, source)
2865 }
2866 "subscript" => {
2867 let value = annotation.child_by_field_name("value")?;
2868 if optional_annotation_wrapper(value, source) {
2869 let inner = annotation.child_by_field_name("subscript")?;
2870 return receiver_type_from_annotation_node(inner, source);
2871 }
2872 normalized_receiver_type(slice(value, source).trim())
2873 }
2874 _ => None,
2875 }
2876}
2877
2878fn optional_annotation_wrapper(node: Node<'_>, source: &str) -> bool {
2879 match node.kind() {
2880 "identifier" => slice(node, source) == "Optional",
2881 "attribute" => {
2882 let (Some(object), Some(attribute)) = (
2883 node.child_by_field_name("object"),
2884 node.child_by_field_name("attribute"),
2885 ) else {
2886 return false;
2887 };
2888 object.kind() == "identifier"
2889 && attribute.kind() == "identifier"
2890 && slice(object, source) == "typing"
2891 && slice(attribute, source) == "Optional"
2892 }
2893 _ => false,
2894 }
2895}
2896
2897fn returned_receiver_type(node: Node<'_>, source: &str) -> Option<String> {
2898 let raw = match node.kind() {
2899 "identifier" => non_empty_node_text(node, source),
2900 "call" => node
2901 .child_by_field_name("function")
2902 .or_else(|| node.named_child(0))
2903 .filter(|callee| callee.kind() == "identifier")
2904 .and_then(|callee| non_empty_node_text(callee, source)),
2905 _ => None,
2906 }?;
2907 normalized_receiver_type(&raw)
2908}
2909
2910#[cfg(test)]
2911mod tests {
2912 use super::*;
2913 use std::path::PathBuf;
2914
2915 #[test]
2916 fn lambda_scope_facts_preserve_untyped_parameter_shadowing() {
2917 let source = "shadowed = lambda method: method.signature\n";
2918 let mut parser = Parser::new();
2919 parser
2920 .set_language(&tree_sitter_python::LANGUAGE.into())
2921 .unwrap();
2922 let tree = parser.parse(source, None).unwrap();
2923 let mut nodes = vec![tree.root_node()];
2924 let lambda = loop {
2925 let node = nodes.pop().unwrap();
2926 if node.kind() == "lambda" {
2927 break node;
2928 }
2929 let mut cursor = node.walk();
2930 nodes.extend(node.named_children(&mut cursor));
2931 };
2932
2933 let facts = collect_function_scope_facts_from_node(lambda, source);
2934
2935 assert!(facts.is_shadowed("method"));
2936 assert!(facts.resolution_for("method").is_unknown());
2937 }
2938
2939 #[test]
2940 fn pre_cancelled_graph_build_skips_python_file_parsing() {
2941 let temp = tempfile::tempdir().unwrap();
2942 let root = temp.path().canonicalize().unwrap();
2943 std::fs::write(root.join("target.py"), "def target():\n pass\n").unwrap();
2944 let file = ProjectFile::new(root.clone(), PathBuf::from("target.py"));
2945 let files = [file.clone()].into_iter().collect();
2946 let cancellation = CancellationToken::default();
2947 cancellation.cancel();
2948
2949 let graph = build_python_graph(&files, &file, Some(&cancellation));
2950
2951 assert!(graph.parsed.is_empty());
2952 }
2953
2954 #[test]
2955 fn graph_build_parses_only_candidates_and_target_not_transitive_imports() {
2956 let temp = tempfile::tempdir().unwrap();
2957 let root = temp.path().canonicalize().unwrap();
2958 std::fs::write(root.join("target.py"), "from dependency import value\n").unwrap();
2959 std::fs::write(
2960 root.join("candidate.py"),
2961 "from transitively_imported import value\n",
2962 )
2963 .unwrap();
2964 std::fs::write(root.join("dependency.py"), "value = 1\n").unwrap();
2965 std::fs::write(root.join("transitively_imported.py"), "value = 2\n").unwrap();
2966 let target = ProjectFile::new(root.clone(), PathBuf::from("target.py"));
2967 let candidate = ProjectFile::new(root.clone(), PathBuf::from("candidate.py"));
2968 let dependency = ProjectFile::new(root.clone(), PathBuf::from("dependency.py"));
2969 let transitive = ProjectFile::new(root.clone(), PathBuf::from("transitively_imported.py"));
2970 let candidates = [candidate.clone()].into_iter().collect();
2971
2972 let graph = build_python_graph(&candidates, &target, None);
2973
2974 assert_eq!(graph.parsed.len(), 2);
2975 assert!(graph.parsed.contains_key(&target));
2976 assert!(graph.parsed.contains_key(&candidate));
2977 assert!(!graph.parsed.contains_key(&dependency));
2978 assert!(!graph.parsed.contains_key(&transitive));
2979 }
2980}