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 let Some(matches) = self.function_import_binding_targets_query(ident, node) {
560 return matches;
561 }
562 self.module_binding_matches_query(ident, node, true, |kind| {
563 kind != ModuleBindingKind::Other
564 })
565 }
566
567 fn module_binding_targets_symbol(&self, ident: &str, node: Node<'_>) -> bool {
568 if let Some(matches) = self.function_import_binding_targets_query(ident, node) {
569 return matches;
570 }
571 let unclassified_named_import = self.edges.iter().any(|edge| {
572 edge.local_name == ident && !matches!(edge.kind, ImportEdgeKind::Namespace)
573 });
574 self.module_binding_matches_query(ident, node, unclassified_named_import, |kind| {
575 kind == ModuleBindingKind::TargetSymbolImport
576 })
577 }
578
579 fn function_import_binding_targets_query(&self, ident: &str, node: Node<'_>) -> Option<bool> {
584 let binding = self.scoped_import_bindings.iter().rev().find(|binding| {
585 binding.is_function_scoped()
586 && binding.start_byte <= node.start_byte()
587 && binding.scope_start_byte <= node.start_byte()
588 && node.end_byte() <= binding.scope_end_byte
589 && binding.local_name == ident
590 })?;
591 let candidates = resolve_fqn_candidates(self.python, &binding.qualified_name, |name| {
592 self.graph.index.definitions(name).collect()
593 });
594 Some(candidates.iter().any(|candidate| candidate == self.target))
595 }
596
597 fn module_binding_matches_query(
598 &self,
599 ident: &str,
600 node: Node<'_>,
601 unclassified: bool,
602 matches: impl Fn(ModuleBindingKind) -> bool,
603 ) -> bool {
604 if !self.edges.iter().any(|edge| edge.local_name == ident) {
605 return false;
606 }
607 let Some(events) = self.module_bindings.get(ident) else {
608 return unclassified;
609 };
610 let cutoff = if reference_is_deferred_function_body(node) {
611 usize::MAX
612 } else {
613 node.start_byte()
614 };
615 let visible: Vec<_> = events
616 .iter()
617 .filter(|event| event.visible_from <= cutoff)
618 .collect();
619 let start = visible
620 .iter()
621 .rposition(|event| !event.conditional)
622 .unwrap_or(0);
623 visible[start..].iter().any(|event| matches(event.kind))
624 }
625
626 fn self_receiver_matches_target(&self, node: Node<'_>) -> bool {
630 let Some(target_owner) = self.target_owner.as_ref() else {
631 return false;
632 };
633 let range = Range {
634 start_byte: node.start_byte(),
635 end_byte: node.end_byte(),
636 start_line: 0,
637 end_line: 0,
638 };
639 let Some(enclosing) = self.graph.index.enclosing_code_unit(self.file, &range) else {
640 return false;
641 };
642 let enclosing_class = if enclosing.is_class() {
643 enclosing
644 } else {
645 match target_owner_code_unit(self.graph.index, &enclosing) {
646 Some(class) => class,
647 None => return false,
648 }
649 };
650 if &enclosing_class == target_owner {
651 return true;
652 }
653 self.graph
654 .hierarchy
655 .map(|provider| provider.get_ancestors(&enclosing_class))
656 .unwrap_or_default()
657 .into_iter()
658 .any(|ancestor| ancestor == *target_owner)
659 }
660
661 fn receiver_type_matches_target(&self, raw_type: &str) -> bool {
662 if receiver_annotation_matches_target(
663 raw_type,
664 self.edges,
665 self.target_short,
666 self.target_self_file,
667 ) {
668 return true;
669 }
670
671 let Some(target_owner) = self.target_owner.as_ref() else {
672 return false;
673 };
674 let Some(receiver_type) = resolve_receiver_type(
675 self.graph,
676 self.python,
677 self.file,
678 raw_type,
679 self.target_self_file,
680 ) else {
681 return false;
682 };
683 if &receiver_type == target_owner {
684 return true;
685 }
686 self.graph
687 .hierarchy
688 .map(|provider| provider.get_ancestors(&receiver_type))
689 .unwrap_or_default()
690 .into_iter()
691 .any(|ancestor| ancestor == *target_owner)
692 }
693}
694
695pub(crate) fn function_declaration_expression_is_outer_scoped(node: Node<'_>) -> bool {
696 let site_start = node.start_byte();
697 let site_end = node.end_byte();
698 let mut current = node;
699 while let Some(parent) = current.parent() {
700 if parent.kind() == "function_definition" {
701 if parent
702 .child_by_field_name("body")
703 .is_some_and(|body| body.start_byte() <= site_start && site_end <= body.end_byte())
704 {
705 return false;
706 }
707 if parent
708 .child_by_field_name("name")
709 .is_some_and(|name| name.id() == node.id())
710 {
711 return false;
712 }
713 if let Some(parameters) = parent.child_by_field_name("parameters")
714 && parameters.start_byte() <= site_start
715 && site_end <= parameters.end_byte()
716 {
717 let mut parameter = node;
718 while parameter.parent() != Some(parameters) {
719 let Some(next) = parameter.parent() else {
720 return false;
721 };
722 parameter = next;
723 }
724 let binder = if parameter.kind() == "identifier" {
725 Some(parameter)
726 } else {
727 parameter.child_by_field_name("name").or_else(|| {
728 parameter
729 .named_child(0)
730 .filter(|child| child.kind() == "identifier")
731 })
732 };
733 return binder.is_none_or(|binder| binder.id() != node.id());
734 }
735 return true;
736 }
737 if parent.kind() == "decorated_definition" {
738 return current.kind() == "decorator";
739 }
740 if parent.kind() == "class_definition" {
741 break;
742 }
743 current = parent;
744 }
745 false
746}
747
748fn scan_node(node: Node<'_>, ctx: &mut ScanCtx<'_>) {
749 let mut stack = vec![node];
750 while let Some(node) = stack.pop() {
751 match node.kind() {
752 "import_statement" | "import_from_statement" => {
753 handle_import_candidate(node, ctx);
754 continue;
755 }
756 "identifier" => {
757 if handle_annotation_reference_candidate(node, ctx) {
758 continue;
759 }
760 handle_identifier_candidate(node, ctx);
761 }
762 "attribute" => {
763 if handle_annotation_reference_candidate(node, ctx) {
764 continue;
765 }
766 handle_attribute_candidate(node, ctx);
767 }
768 "string_content" => {
769 handle_annotation_reference_candidate(node, ctx);
770 }
771 "keyword_argument" => {
772 handle_keyword_argument_candidate(node, ctx);
773 if let Some(value) = node.child_by_field_name("value") {
774 stack.push(value);
775 }
776 continue;
777 }
778 _ => {}
779 }
780
781 let mut cursor = node.walk();
782 let mut children: Vec<Node<'_>> = node.named_children(&mut cursor).collect();
783 children.reverse();
784 stack.extend(children);
785 }
786}
787
788fn handle_annotation_reference_candidate(node: Node<'_>, ctx: &mut ScanCtx<'_>) -> bool {
789 let Some(candidates) = annotation_reference_candidates(
790 ctx.graph,
791 ctx.python,
792 ctx.file,
793 ctx.source,
794 node,
795 ctx.target_self_file,
796 ) else {
797 return false;
798 };
799
800 if (ctx.target.is_class() || ctx.target.is_field() || ctx.target_member.is_none())
801 && candidates.iter().all(|candidate| *candidate == *ctx.target)
802 && candidates.iter().any(|candidate| *candidate == *ctx.target)
803 {
804 let site = if node.kind() == "attribute" {
805 node.child_by_field_name("attribute").unwrap_or(node)
806 } else {
807 node
808 };
809 record_hit(site, ctx);
810 }
811
812 if let Some(site) = annotation_class_qualifier_site(
813 ctx.graph, ctx.python, ctx.file, ctx.source, node, ctx.target,
814 ) {
815 record_hit(site, ctx);
816 }
817
818 if node.kind() == "attribute" && (candidates.is_empty() || ctx.target_is_module) {
824 return false;
825 }
826
827 true
828}
829
830fn handle_keyword_argument_candidate(node: Node<'_>, ctx: &mut ScanCtx<'_>) {
831 let (Some(target_member), Some(name), Some(arguments)) = (
832 ctx.target_member,
833 node.child_by_field_name("name"),
834 node.parent(),
835 ) else {
836 return;
837 };
838 if name.kind() != "identifier"
839 || slice(name, ctx.source) != target_member
840 || arguments.kind() != "argument_list"
841 {
842 return;
843 }
844 let Some(call) = arguments.parent().filter(|parent| parent.kind() == "call") else {
845 return;
846 };
847 let Some(function) = call.child_by_field_name("function") else {
848 return;
849 };
850 if function.kind() == "identifier" && slice(function, ctx.source) == "cls" {
851 if ctx.self_receiver_matches_target(function) {
852 record_hit(name, ctx);
853 }
854 return;
855 }
856 let Some(target_owner) = ctx.target_owner.as_ref() else {
857 return;
858 };
859 let scoped_callee_matches = if function.kind() == "identifier" {
860 ctx.scope_facts_for_node(function)
861 .and_then(|facts| {
862 facts
863 .resolution_for(slice(function, ctx.source))
864 .as_precise()
865 .and_then(|targets| targets.iter().next().cloned())
866 })
867 .is_some_and(|raw_type| ctx.receiver_type_matches_target(&raw_type))
868 } else {
869 false
870 };
871 let default_callee_matches = if function.kind() == "identifier" {
872 resolve_callable_parameter_default_types(
873 ctx.graph,
874 ctx.python,
875 ctx.file,
876 ctx.source,
877 function,
878 slice(function, ctx.source),
879 )
880 .into_iter()
881 .any(|class| {
882 &class == target_owner
883 || ctx
884 .graph
885 .hierarchy
886 .map(|provider| provider.get_ancestors(&class))
887 .unwrap_or_default()
888 .into_iter()
889 .any(|ancestor| &ancestor == target_owner)
890 })
891 } else {
892 false
893 };
894 let root_shadowed = leftmost_identifier(function).is_some_and(|root| {
895 ctx.scope_facts_for_node(function)
896 .is_some_and(|facts| facts.is_shadowed(slice(root, ctx.source)))
897 });
898 if root_shadowed && !scoped_callee_matches && !default_callee_matches {
899 return;
900 }
901 let matches = scoped_callee_matches
902 || default_callee_matches
903 || (!root_shadowed
904 && resolve_constructor_types(ctx.graph, ctx.python, ctx.file, ctx.source, function)
905 .into_iter()
906 .any(|class| {
907 &class == target_owner
908 || ctx
909 .graph
910 .hierarchy
911 .map(|provider| provider.get_ancestors(&class))
912 .unwrap_or_default()
913 .into_iter()
914 .any(|ancestor| &ancestor == target_owner)
915 }));
916 if matches {
917 record_hit(name, ctx);
918 }
919}
920
921fn leftmost_identifier(mut node: Node<'_>) -> Option<Node<'_>> {
922 loop {
923 match node.kind() {
924 "identifier" => return Some(node),
925 "attribute" => node = node.child_by_field_name("object")?,
926 _ => return None,
927 }
928 }
929}
930
931fn handle_import_candidate(node: Node<'_>, ctx: &mut ScanCtx<'_>) {
937 if ctx.target_member.is_some() {
938 return;
939 }
940 if !ctx
941 .edges
942 .iter()
943 .any(|edge| edge.local_name == ctx.target_short)
944 {
945 return;
946 }
947 let mut stack = vec![node];
948 while let Some(node) = stack.pop() {
949 if node.kind() == "identifier" && slice(node, ctx.source) == ctx.target_short {
950 record_import_hit(node, ctx);
951 return;
952 }
953 let mut cursor = node.walk();
954 for child in node.named_children(&mut cursor) {
955 stack.push(child);
956 }
957 }
958}
959
960fn handle_identifier_candidate(node: Node<'_>, ctx: &mut ScanCtx<'_>) {
961 if node
962 .parent()
963 .is_some_and(|parent| parent.kind() == "attribute")
964 {
965 return;
966 }
967 let text = slice(node, ctx.source);
968 if text.is_empty() || is_declaration_identifier(node) || decorates_the_target(node, ctx) {
969 return;
970 }
971 if let Some(member) = ctx.target_member {
972 if member == "__init__" && is_call_callee(node) && ctx.binds_target(text, node) {
975 record_hit(node, ctx);
976 return;
977 }
978 if text == member && ctx.node_directly_in_owner_class_body(node) {
981 record_hit(node, ctx);
982 }
983 return;
984 }
985 if !ctx.binds_target(text, node) {
986 return;
987 }
988 if !ctx.target_is_module
989 && ctx.edges.iter().any(|edge| edge.local_name == text)
990 && !ctx.module_binding_targets_symbol(text, node)
991 {
992 return;
993 }
994 record_hit(node, ctx);
995}
996
997fn decorates_the_target(node: Node<'_>, ctx: &ScanCtx<'_>) -> bool {
1005 if ctx.file != ctx.target_source {
1006 return false;
1007 }
1008 let mut current = node;
1009 while let Some(parent) = current.parent() {
1010 if parent.kind() == "decorated_definition" && current.kind() == "decorator" {
1011 let Some(definition) = parent.child_by_field_name("definition") else {
1012 return false;
1013 };
1014 return ctx.graph.index.ranges(ctx.target).iter().any(|range| {
1015 range.start_byte <= definition.start_byte()
1016 && definition.end_byte() <= range.end_byte
1017 });
1018 }
1019 current = parent;
1020 }
1021 false
1022}
1023
1024fn handle_attribute_candidate(node: Node<'_>, ctx: &mut ScanCtx<'_>) {
1025 let Some(object) = node.child_by_field_name("object") else {
1026 return;
1027 };
1028 let Some(attribute) = node.child_by_field_name("attribute") else {
1029 return;
1030 };
1031 let object_text = slice(object, ctx.source);
1032 let attribute_text = slice(attribute, ctx.source);
1033 if let Some(member) = ctx.target_member
1034 && attribute_text == member
1035 {
1036 let is_same_owner_receiver =
1041 matches!(object_text, "self" | "cls") && ctx.self_receiver_matches_target(node);
1042 if is_same_owner_receiver {
1043 record_self_receiver_hit(attribute, ctx);
1044 } else if ctx.receiver_binds_target(object_text, node)
1045 || (object.kind() == "call" && call_result_matches_target(object, ctx))
1046 {
1047 record_hit(attribute, ctx);
1048 } else if member_receiver_match_is_unproven(object, object_text, node, ctx) {
1049 record_unproven_hit(attribute, ctx);
1050 }
1051 }
1052
1053 let object_binds_target = if ctx.target_is_module {
1054 imported_root_targets_module(ctx, object, node)
1055 } else {
1056 ctx.binds_target(object_text, node)
1057 };
1058 if object.kind() == "identifier"
1059 && object_binds_target
1060 && (ctx.target_is_module
1061 || (ctx.target_member.is_none()
1062 && !ctx.edges.iter().any(|edge| {
1063 matches!(edge.kind, ImportEdgeKind::Namespace) && edge.local_name == object_text
1064 })))
1065 {
1066 record_hit(object, ctx);
1067 }
1068
1069 if ctx.target_is_module
1070 && let Some(module_qualifier) = module_attribute_target_hit(node, ctx)
1071 {
1072 record_hit(module_qualifier, ctx);
1073 }
1074
1075 if let Some(member) = ctx.target_member
1079 && object.kind() == "identifier"
1080 && object_text == member
1081 && ctx.node_directly_in_owner_class_body(object)
1082 {
1083 record_hit(object, ctx);
1084 }
1085
1086 if ctx.member_best_effort_unique
1092 && let Some(member) = ctx.target_member
1093 && attribute_text == member
1094 && object.kind() == "identifier"
1095 && !matches!(object_text, "self" | "cls")
1096 && !ctx.receiver_binds_target(object_text, node)
1097 && ctx.receiver_type_is_unknown(object_text, node)
1098 {
1099 record_hit(attribute, ctx);
1100 }
1101
1102 if let Some(module_binding_target) = module_binding_attribute_target_hit(node, ctx) {
1103 record_hit(module_binding_target, ctx);
1104 }
1105}
1106
1107fn module_binding_attribute_target_hit<'a>(node: Node<'a>, ctx: &ScanCtx<'_>) -> Option<Node<'a>> {
1119 if ctx.target_member.is_some() {
1120 return None;
1121 }
1122 let (root, attributes) = attribute_chain(node)?;
1123 let terminal = *attributes.last()?;
1124 let terminal_name = slice(terminal, ctx.source);
1125 if terminal_name.is_empty()
1126 || !ctx
1127 .seeds
1128 .iter()
1129 .any(|(_, seed_name)| seed_name == terminal_name)
1130 {
1131 return None;
1132 }
1133
1134 for binding in imported_module_bindings(ctx, root, node) {
1135 let mut written_module = binding.module.clone();
1136 for attribute in attributes
1137 [binding.consumed_attributes.min(attributes.len() - 1)..attributes.len() - 1]
1138 .iter()
1139 {
1140 let segment = slice(*attribute, ctx.source);
1141 if segment.is_empty() {
1142 return None;
1143 }
1144 written_module.push('.');
1145 written_module.push_str(segment);
1146 }
1147 if usage_resolve_module_files(ctx.python, ctx.file, &written_module)
1148 .iter()
1149 .any(|resolved| {
1150 ctx.seeds
1151 .contains(&(resolved.clone(), terminal_name.to_string()))
1152 })
1153 {
1154 return Some(terminal);
1155 }
1156
1157 let mut written_fqn = binding.module;
1158 for attribute in attributes.iter().skip(binding.consumed_attributes) {
1159 let segment = slice(*attribute, ctx.source);
1160 if segment.is_empty() {
1161 return None;
1162 }
1163 written_fqn.push('.');
1164 written_fqn.push_str(segment);
1165 }
1166 if resolve_fqn_candidates(ctx.python, &written_fqn, |name| {
1167 ctx.graph.index.definitions(name).collect()
1168 })
1169 .into_iter()
1170 .any(|candidate| &candidate == ctx.target)
1171 {
1172 return Some(terminal);
1173 }
1174 }
1175 None
1176}
1177
1178fn imported_root_targets_module(ctx: &ScanCtx<'_>, root: Node<'_>, reference: Node<'_>) -> bool {
1179 imported_module_bindings(ctx, root, reference)
1180 .into_iter()
1181 .any(|binding| {
1182 usage_resolve_module_files(ctx.python, ctx.file, &binding.module)
1183 .into_iter()
1184 .any(|resolved_file| &resolved_file == ctx.target_source)
1185 })
1186}
1187
1188fn module_attribute_target_hit<'a>(node: Node<'a>, ctx: &ScanCtx<'_>) -> Option<Node<'a>> {
1189 let (root, attributes) = attribute_chain(node)?;
1190 if attributes.is_empty() {
1191 return None;
1192 }
1193 for binding in imported_module_bindings(ctx, root, node) {
1194 let mut module_fqn = binding.module;
1195 for attribute in attributes.iter().skip(binding.consumed_attributes) {
1196 let segment = slice(*attribute, ctx.source);
1197 if segment.is_empty() {
1198 return None;
1199 }
1200 if module_fqn.ends_with('.') {
1201 module_fqn.push_str(segment);
1202 } else {
1203 module_fqn.push('.');
1204 module_fqn.push_str(segment);
1205 }
1206 let resolved = usage_resolve_module_files(ctx.python, ctx.file, &module_fqn);
1207 if resolved.is_empty() {
1208 break;
1209 }
1210 if resolved
1211 .iter()
1212 .any(|resolved_file| resolved_file == ctx.target_source)
1213 {
1214 return Some(*attribute);
1215 }
1216 }
1217 }
1218 None
1219}
1220
1221fn call_result_matches_target(call: Node<'_>, ctx: &ScanCtx<'_>) -> bool {
1222 let Some(target_owner) = ctx.target_owner.as_ref() else {
1223 return false;
1224 };
1225 let scope_facts = ctx.scope_facts_for_node(call);
1226 call_result_types(
1227 ctx.graph,
1228 ctx.python,
1229 ctx.file,
1230 ctx.source,
1231 call,
1232 scope_facts,
1233 )
1234 .into_iter()
1235 .any(|class| {
1236 &class == target_owner
1237 || ctx
1238 .graph
1239 .hierarchy
1240 .map(|provider| provider.get_ancestors(&class))
1241 .unwrap_or_default()
1242 .into_iter()
1243 .any(|ancestor| &ancestor == target_owner)
1244 })
1245}
1246
1247pub fn call_result_types(
1248 graph: &PythonGraphSource<'_>,
1249 python: &dyn PythonUsageSource,
1250 file: &ProjectFile,
1251 source: &str,
1252 call: Node<'_>,
1253 scope_facts: Option<&LocalBindingsSnapshot<String>>,
1254) -> Vec<CodeUnit> {
1255 let Some(function) = call.child_by_field_name("function") else {
1256 return Vec::new();
1257 };
1258 let constructed = resolve_constructor_types(graph, python, file, source, function);
1259 if !constructed.is_empty() {
1260 return constructed;
1261 }
1262 let callable_fqns = resolve_callable_fqns(graph, python, file, source, function, scope_facts);
1263 if callable_fqns.is_empty() {
1264 return Vec::new();
1265 }
1266 let callables = callable_fqns
1267 .into_iter()
1268 .flat_map(|callable_fqn| {
1269 resolve_fqn_candidates(python, &callable_fqn, |name| {
1270 graph.index.definitions(name).collect()
1271 })
1272 })
1273 .collect::<Vec<_>>();
1274 let mut classes = Vec::new();
1275 for callable in callables.into_iter().filter(CodeUnit::is_function) {
1276 let Some(raw_type) = callable_return_type_name(graph, python, &callable) else {
1277 continue;
1278 };
1279 if let Some(class) =
1280 resolve_receiver_type(graph, python, callable.source(), &raw_type, true)
1281 {
1282 classes.push(class);
1283 }
1284 }
1285 classes.sort();
1286 classes.dedup();
1287 classes
1288}
1289
1290fn resolve_callable_fqns(
1291 graph: &PythonGraphSource<'_>,
1292 python: &dyn PythonUsageSource,
1293 file: &ProjectFile,
1294 source: &str,
1295 function: Node<'_>,
1296 scope_facts: Option<&LocalBindingsSnapshot<String>>,
1297) -> Vec<String> {
1298 match function.kind() {
1299 "identifier" => {
1300 resolve_identifier_callable_fqns(graph, python, file, source, function, scope_facts)
1301 }
1302 "attribute" => {
1303 resolve_attribute_callable_fqns(graph, python, file, source, function, scope_facts)
1304 }
1305 _ => Vec::new(),
1306 }
1307}
1308
1309fn resolve_identifier_callable_fqns(
1310 graph: &PythonGraphSource<'_>,
1311 python: &dyn PythonUsageSource,
1312 file: &ProjectFile,
1313 source: &str,
1314 function: Node<'_>,
1315 scope_facts: Option<&LocalBindingsSnapshot<String>>,
1316) -> Vec<String> {
1317 let local = slice(function, source);
1318 if local.is_empty() || scope_facts.is_some_and(|facts| facts.is_shadowed(local)) {
1319 return Vec::new();
1320 }
1321 let binder = python.import_binder_of(file);
1322 match binder.bindings.get(local) {
1323 Some(binding) if binding.kind == ImportKind::Named => binding
1324 .imported_name
1325 .as_ref()
1326 .map(|imported| vec![format!("{}.{}", binding.module_specifier, imported)])
1327 .unwrap_or_default(),
1328 _ => graph
1329 .index
1330 .declarations(file)
1331 .into_iter()
1332 .find(|unit| unit.is_function() && unit.identifier() == local)
1333 .map(|unit| vec![unit.fq_name()])
1334 .unwrap_or_default(),
1335 }
1336}
1337
1338fn resolve_attribute_callable_fqns(
1339 graph: &PythonGraphSource<'_>,
1340 python: &dyn PythonUsageSource,
1341 file: &ProjectFile,
1342 source: &str,
1343 function: Node<'_>,
1344 scope_facts: Option<&LocalBindingsSnapshot<String>>,
1345) -> Vec<String> {
1346 let Some(receiver) = function.child_by_field_name("object") else {
1347 return Vec::new();
1348 };
1349 let Some(method) = function.child_by_field_name("attribute") else {
1350 return Vec::new();
1351 };
1352 let method = slice(method, source);
1353 if method.is_empty() {
1354 return Vec::new();
1355 }
1356 let mut fqns = attribute_receiver_classes(graph, python, file, source, receiver, scope_facts)
1357 .into_iter()
1358 .map(|class| format!("{}.{}", class.fq_name(), method))
1359 .collect::<Vec<_>>();
1360 fqns.sort();
1361 fqns.dedup();
1362 fqns
1363}
1364
1365fn attribute_receiver_classes(
1366 graph: &PythonGraphSource<'_>,
1367 python: &dyn PythonUsageSource,
1368 file: &ProjectFile,
1369 source: &str,
1370 receiver: Node<'_>,
1371 scope_facts: Option<&LocalBindingsSnapshot<String>>,
1372) -> Vec<CodeUnit> {
1373 let mut classes = match receiver.kind() {
1374 "identifier" => {
1375 identifier_receiver_classes(graph, python, file, source, receiver, scope_facts)
1376 }
1377 "attribute" => {
1378 if let Some(root) = leftmost_identifier(receiver)
1379 && scope_facts.is_some_and(|facts| facts.is_shadowed(slice(root, source)))
1380 {
1381 Vec::new()
1382 } else {
1383 resolve_constructor_types(graph, python, file, source, receiver)
1384 }
1385 }
1386 _ => Vec::new(),
1387 };
1388 classes.sort();
1389 classes.dedup();
1390 classes
1391}
1392
1393fn identifier_receiver_classes(
1394 graph: &PythonGraphSource<'_>,
1395 python: &dyn PythonUsageSource,
1396 file: &ProjectFile,
1397 source: &str,
1398 receiver: Node<'_>,
1399 scope_facts: Option<&LocalBindingsSnapshot<String>>,
1400) -> Vec<CodeUnit> {
1401 let ident = slice(receiver, source);
1402 if ident.is_empty() {
1403 return Vec::new();
1404 }
1405 if matches!(ident, "self" | "cls")
1406 && let Some(class) = enclosing_class_for_node(graph, file, receiver)
1407 {
1408 return vec![class];
1409 }
1410 if let Some(facts) = scope_facts {
1411 if let Some(raw_type) = facts
1412 .resolution_for(ident)
1413 .as_precise()
1414 .and_then(|targets| targets.iter().next())
1415 && let Some(class) = resolve_receiver_type(graph, python, file, raw_type, false)
1416 {
1417 return vec![class];
1418 }
1419 if facts.is_shadowed(ident) {
1420 return Vec::new();
1421 }
1422 }
1423 resolve_receiver_type(graph, python, file, ident, false)
1424 .into_iter()
1425 .collect()
1426}
1427
1428fn enclosing_class_for_node(
1429 graph: &PythonGraphSource<'_>,
1430 file: &ProjectFile,
1431 node: Node<'_>,
1432) -> Option<CodeUnit> {
1433 let range = Range {
1434 start_byte: node.start_byte(),
1435 end_byte: node.end_byte(),
1436 start_line: 0,
1437 end_line: 0,
1438 };
1439 let enclosing = graph.index.enclosing_code_unit(file, &range)?;
1440 brokk_bifrost_core::analyzer::usages::common::enclosing_owner_chain(enclosing, |unit| {
1446 graph.index.parent_of(unit)
1447 })
1448 .take(2)
1449 .find(|unit| unit.is_class() && unit.source() == file)
1450}
1451
1452fn attribute_chain<'a>(node: Node<'a>) -> Option<(Node<'a>, Vec<Node<'a>>)> {
1453 let mut attributes = Vec::new();
1454 let mut current = node;
1455 loop {
1456 if current.kind() != "attribute" {
1457 return None;
1458 }
1459 attributes.push(current.child_by_field_name("attribute")?);
1460 current = current.child_by_field_name("object")?;
1461 if current.kind() == "identifier" {
1462 attributes.reverse();
1463 return Some((current, attributes));
1464 }
1465 }
1466}
1467
1468struct ImportedModuleBinding {
1469 module: String,
1470 consumed_attributes: usize,
1471}
1472
1473fn imported_module_bindings(
1474 ctx: &ScanCtx<'_>,
1475 root: Node<'_>,
1476 reference: Node<'_>,
1477) -> Vec<ImportedModuleBinding> {
1478 let root_text = slice(root, ctx.source);
1479 if root_text.is_empty() || import_root_shadowed(ctx, root_text, root, reference) {
1480 return Vec::new();
1481 }
1482
1483 if let Some(binding) = ctx.scoped_import_bindings.iter().rev().find(|binding| {
1484 binding.is_function_scoped()
1485 && binding.start_byte <= reference.start_byte()
1486 && binding.scope_start_byte <= reference.start_byte()
1487 && reference.end_byte() <= binding.scope_end_byte
1488 && binding.local_name == root_text
1489 }) {
1490 return if usage_resolve_module_files(ctx.python, ctx.file, &binding.qualified_name)
1491 .is_empty()
1492 {
1493 Vec::new()
1494 } else {
1495 vec![ImportedModuleBinding {
1496 module: binding.qualified_name.clone(),
1497 consumed_attributes: binding.consumed_attributes,
1498 }]
1499 };
1500 }
1501
1502 let Some(events) = ctx.raw_module_bindings.get(root_text) else {
1503 return Vec::new();
1504 };
1505 let cutoff = if reference_is_deferred_function_body(reference) {
1506 usize::MAX
1507 } else {
1508 reference.start_byte()
1509 };
1510 let visible: Vec<_> = events
1511 .iter()
1512 .filter(|event| event.visible_from <= cutoff)
1513 .collect();
1514 let start = visible
1515 .iter()
1516 .rposition(|event| !event.conditional)
1517 .unwrap_or(0);
1518 let mut modules = visible[start..]
1519 .iter()
1520 .filter_map(|event| match &event.kind {
1521 ModuleBindingEventKind::ImportModule {
1522 module,
1523 consumed_attributes,
1524 } => Some(ImportedModuleBinding {
1525 module: module.clone(),
1526 consumed_attributes: *consumed_attributes,
1527 }),
1528 ModuleBindingEventKind::FromImport {
1529 module,
1530 imported_name,
1531 } => {
1532 let submodule = if module.ends_with('.') {
1533 format!("{module}{imported_name}")
1534 } else {
1535 format!("{module}.{imported_name}")
1536 };
1537 (!usage_resolve_module_files(ctx.python, ctx.file, &submodule).is_empty())
1538 .then_some(ImportedModuleBinding {
1539 module: submodule,
1540 consumed_attributes: 0,
1541 })
1542 }
1543 ModuleBindingEventKind::Other => None,
1544 })
1545 .collect::<Vec<_>>();
1546 modules.sort_by(|left, right| {
1547 left.module
1548 .cmp(&right.module)
1549 .then_with(|| left.consumed_attributes.cmp(&right.consumed_attributes))
1550 });
1551 modules.dedup_by(|left, right| {
1552 left.module == right.module && left.consumed_attributes == right.consumed_attributes
1553 });
1554 modules
1555}
1556
1557fn import_root_shadowed(
1558 ctx: &ScanCtx<'_>,
1559 root_text: &str,
1560 root: Node<'_>,
1561 reference: Node<'_>,
1562) -> bool {
1563 ctx.scope_entry_for_node(root)
1564 .or_else(|| ctx.scope_entry_for_node(reference))
1565 .is_some_and(|(scope, facts)| !scope.is_module() && facts.is_shadowed(root_text))
1566 || enclosing_parameters_shadow(root_text, reference, ctx.source)
1567}
1568
1569fn enclosing_parameters_shadow(root_text: &str, reference: Node<'_>, source: &str) -> bool {
1570 let mut current = reference;
1571 while let Some(parent) = current.parent() {
1572 if matches!(parent.kind(), "function_definition" | "lambda") {
1573 let Some(parameters) = parent.child_by_field_name("parameters") else {
1574 return false;
1575 };
1576 let mut cursor = parameters.walk();
1577 return parameters.named_children(&mut cursor).any(|parameter| {
1578 parameter_symbol(parameter, source).as_deref() == Some(root_text)
1579 });
1580 }
1581 current = parent;
1582 }
1583 false
1584}
1585
1586enum EnclosingParameterType {
1587 NotDeclared,
1588 Untyped,
1589 Typed(String),
1590}
1591
1592fn enclosing_runtime_parameter_type(
1593 name: &str,
1594 reference: Node<'_>,
1595 source: &str,
1596) -> EnclosingParameterType {
1597 let site_start = reference.start_byte();
1598 let site_end = reference.end_byte();
1599 let mut current = reference;
1600 while let Some(parent) = current.parent() {
1601 if matches!(parent.kind(), "function_definition" | "lambda")
1602 && parent
1603 .child_by_field_name("body")
1604 .is_some_and(|body| body.start_byte() <= site_start && site_end <= body.end_byte())
1605 && let Some(parameters) = parent.child_by_field_name("parameters")
1606 {
1607 let mut cursor = parameters.walk();
1608 for parameter in parameters.named_children(&mut cursor) {
1609 if parameter_symbol(parameter, source).as_deref() != Some(name) {
1610 continue;
1611 }
1612 return parameter
1613 .child_by_field_name("type")
1614 .and_then(|annotation| normalized_receiver_type(slice(annotation, source)))
1615 .map_or(EnclosingParameterType::Untyped, |raw_type| {
1616 EnclosingParameterType::Typed(raw_type)
1617 });
1618 }
1619 }
1620 current = parent;
1621 }
1622 EnclosingParameterType::NotDeclared
1623}
1624
1625fn member_receiver_match_is_unproven(
1626 object: Node<'_>,
1627 object_text: &str,
1628 node: Node<'_>,
1629 ctx: &ScanCtx<'_>,
1630) -> bool {
1631 if matches!(object_text, "self" | "cls") {
1632 return false;
1633 }
1634 match object.kind() {
1635 "identifier" => {
1636 ctx.receiver_type_is_unknown(object_text, node) && !ctx.member_best_effort_unique
1637 }
1638 "attribute" => true,
1639 _ => false,
1640 }
1641}
1642
1643pub fn slice<'a>(node: Node<'_>, source: &'a str) -> &'a str {
1644 brokk_bifrost_core::analyzer::common::node_source_text(node, source)
1645}
1646
1647fn is_call_callee(node: Node<'_>) -> bool {
1649 node.parent().is_some_and(|parent| {
1650 parent.kind() == "call"
1651 && parent
1652 .child_by_field_name("function")
1653 .is_some_and(|function| function.id() == node.id())
1654 })
1655}
1656
1657pub fn is_declaration_identifier(node: Node<'_>) -> bool {
1658 let Some(parent) = node.parent() else {
1659 return false;
1660 };
1661 let contains = |container: Node<'_>| {
1662 container.start_byte() <= node.start_byte() && node.end_byte() <= container.end_byte()
1663 };
1664 match parent.kind() {
1665 "class_definition" | "function_definition" => parent
1666 .child_by_field_name("name")
1667 .is_some_and(|name| name.id() == node.id()),
1668 "parameters" | "lambda_parameters" | "list_splat_pattern" | "dictionary_splat_pattern" => {
1669 true
1670 }
1671 "default_parameter" | "typed_parameter" | "typed_default_parameter" => {
1672 parent.child_by_field_name("name").is_some_and(contains)
1673 }
1674 "assignment" | "augmented_assignment" | "for_statement" | "for_in_clause" => {
1675 parent.child_by_field_name("left").is_some_and(contains)
1676 }
1677 "named_expression" => parent.child_by_field_name("name").is_some_and(contains),
1678 "aliased_import" | "import_from_statement" | "import_statement" => true,
1679 _ => false,
1680 }
1681}
1682
1683#[derive(Clone, Copy, Debug, PartialEq, Eq)]
1684enum ModuleBindingKind {
1685 TargetSymbolImport,
1686 TargetModuleImport,
1687 Other,
1688}
1689
1690#[derive(Clone, Copy, Debug)]
1691struct ClassifiedModuleBindingEvent {
1692 visible_from: usize,
1693 conditional: bool,
1694 kind: ModuleBindingKind,
1695}
1696
1697pub fn collect_module_binding_timeline(root: Node<'_>, source: &str) -> ModuleBindingTimeline {
1698 let mut timeline = ModuleBindingTimeline::default();
1699 let mut stack = vec![root];
1700 while let Some(node) = stack.pop() {
1701 match node.kind() {
1702 "function_definition" | "class_definition" => {
1703 if let Some(name) = node.child_by_field_name("name") {
1704 record_module_binding(
1705 &mut timeline,
1706 slice(name, source),
1707 node.end_byte(),
1708 binding_is_conditional(node),
1709 ModuleBindingEventKind::Other,
1710 );
1711 }
1712 continue;
1713 }
1714 "import_statement" | "import_from_statement" => {
1715 collect_import_binding_events(node, source, &mut timeline);
1716 continue;
1717 }
1718 "assignment" | "augmented_assignment" | "named_expression" => {
1719 if let Some(left) = node.child_by_field_name("left") {
1720 record_local_binding_targets(
1721 left,
1722 source,
1723 node.end_byte(),
1724 binding_is_conditional(node),
1725 &mut timeline,
1726 );
1727 }
1728 continue;
1729 }
1730 "for_statement" => {
1731 if let Some(left) = node.child_by_field_name("left") {
1732 record_local_binding_targets(
1733 left,
1734 source,
1735 left.end_byte(),
1736 true,
1737 &mut timeline,
1738 );
1739 }
1740 }
1741 _ => {}
1742 }
1743
1744 let mut cursor = node.walk();
1745 let mut children: Vec<Node<'_>> = node.named_children(&mut cursor).collect();
1746 children.reverse();
1747 stack.extend(children);
1748 }
1749 for events in timeline.values_mut() {
1750 events.sort_by_key(|event| event.visible_from);
1751 }
1752 timeline
1753}
1754
1755fn collect_import_binding_events(
1756 node: Node<'_>,
1757 source: &str,
1758 timeline: &mut ModuleBindingTimeline,
1759) {
1760 if node.kind() == "import_statement" {
1761 let mut cursor = node.walk();
1762 for imported in node.children_by_field_name("name", &mut cursor) {
1763 let name = imported.child_by_field_name("name").unwrap_or(imported);
1764 let Some(local) = imported
1765 .child_by_field_name("alias")
1766 .or_else(|| first_identifier(name))
1767 else {
1768 continue;
1769 };
1770 let module = slice(name, source).trim();
1771 let consumed_attributes = if imported.child_by_field_name("alias").is_some() {
1772 0
1773 } else {
1774 parse_symbol_path(Language::Python, module)
1775 .len()
1776 .saturating_sub(1)
1777 };
1778 record_module_binding(
1779 timeline,
1780 slice(local, source),
1781 node.end_byte(),
1782 binding_is_conditional(node),
1783 ModuleBindingEventKind::ImportModule {
1784 module: module.to_string(),
1785 consumed_attributes,
1786 },
1787 );
1788 }
1789 return;
1790 }
1791
1792 let Some(module_node) = node.child_by_field_name("module_name") else {
1793 return;
1794 };
1795 let module = slice(module_node, source).trim();
1796 let mut cursor = node.walk();
1797 for imported in node.children_by_field_name("name", &mut cursor) {
1798 if imported.kind() == "wildcard_import" {
1799 continue;
1800 }
1801 let name = imported.child_by_field_name("name").unwrap_or(imported);
1802 let Some(imported_identifier) = last_identifier(name) else {
1803 continue;
1804 };
1805 let imported_name = slice(imported_identifier, source).trim();
1806 let Some(local) = imported
1807 .child_by_field_name("alias")
1808 .or_else(|| last_identifier(name))
1809 else {
1810 continue;
1811 };
1812 record_module_binding(
1813 timeline,
1814 slice(local, source),
1815 node.end_byte(),
1816 binding_is_conditional(node),
1817 ModuleBindingEventKind::FromImport {
1818 module: module.to_string(),
1819 imported_name: imported_name.to_string(),
1820 },
1821 );
1822 }
1823}
1824
1825fn classify_module_binding_timeline(
1826 python: &dyn PythonUsageSource,
1827 file: &ProjectFile,
1828 timeline: &ModuleBindingTimeline,
1829 seeds: &BTreeSet<(ProjectFile, String)>,
1830 edges: &[ImportEdge],
1831) -> HashMap<String, Vec<ClassifiedModuleBindingEvent>> {
1832 let mut classified = HashMap::default();
1833 let mut module_targets: HashMap<String, bool> = HashMap::default();
1834 let relevant_locals: HashSet<&str> =
1835 edges.iter().map(|edge| edge.local_name.as_str()).collect();
1836 for (local, events) in timeline {
1837 if !relevant_locals.contains(local.as_str()) {
1838 continue;
1839 }
1840 let classified_events = events
1841 .iter()
1842 .map(|event| {
1843 let kind = match &event.kind {
1844 ModuleBindingEventKind::ImportModule { module, .. } => {
1845 if *module_targets
1846 .entry(module.clone())
1847 .or_insert_with(|| module_contains_seed(python, file, module, seeds))
1848 {
1849 ModuleBindingKind::TargetModuleImport
1850 } else {
1851 ModuleBindingKind::Other
1852 }
1853 }
1854 ModuleBindingEventKind::FromImport {
1855 module,
1856 imported_name,
1857 } => {
1858 let direct = usage_resolve_module_files(python, file, module).iter().any(
1859 |resolved| seeds.contains(&(resolved.clone(), imported_name.clone())),
1860 );
1861 let submodule = if module.ends_with('.') {
1862 format!("{module}{imported_name}")
1863 } else {
1864 format!("{module}.{imported_name}")
1865 };
1866 let imports_target_module =
1867 *module_targets.entry(submodule.clone()).or_insert_with(|| {
1868 module_contains_seed(python, file, &submodule, seeds)
1869 });
1870 if direct {
1871 ModuleBindingKind::TargetSymbolImport
1872 } else if imports_target_module {
1873 ModuleBindingKind::TargetModuleImport
1874 } else {
1875 ModuleBindingKind::Other
1876 }
1877 }
1878 ModuleBindingEventKind::Other => ModuleBindingKind::Other,
1879 };
1880 ClassifiedModuleBindingEvent {
1881 visible_from: event.visible_from,
1882 conditional: event.conditional,
1883 kind,
1884 }
1885 })
1886 .collect();
1887 classified.insert(local.clone(), classified_events);
1888 }
1889 classified
1890}
1891
1892fn module_contains_seed(
1893 python: &dyn PythonUsageSource,
1894 file: &ProjectFile,
1895 module: &str,
1896 seeds: &BTreeSet<(ProjectFile, String)>,
1897) -> bool {
1898 usage_resolve_module_files(python, file, module)
1899 .iter()
1900 .any(|resolved| seeds.iter().any(|(seed_file, _)| seed_file == resolved))
1901}
1902
1903fn record_module_binding(
1904 timeline: &mut ModuleBindingTimeline,
1905 name: &str,
1906 visible_from: usize,
1907 conditional: bool,
1908 kind: ModuleBindingEventKind,
1909) {
1910 let name = name.trim();
1911 if name.is_empty() {
1912 return;
1913 }
1914 timeline
1915 .entry(name.to_string())
1916 .or_default()
1917 .push(ModuleBindingEvent {
1918 visible_from,
1919 conditional,
1920 kind,
1921 });
1922}
1923
1924fn record_local_binding_targets(
1925 target: Node<'_>,
1926 source: &str,
1927 visible_from: usize,
1928 conditional: bool,
1929 timeline: &mut ModuleBindingTimeline,
1930) {
1931 let mut stack = vec![target];
1932 while let Some(node) = stack.pop() {
1933 if node.kind() == "identifier" {
1934 record_module_binding(
1935 timeline,
1936 slice(node, source),
1937 visible_from,
1938 conditional,
1939 ModuleBindingEventKind::Other,
1940 );
1941 continue;
1942 }
1943 if matches!(node.kind(), "attribute" | "subscript") {
1944 continue;
1945 }
1946 let mut cursor = node.walk();
1947 let mut children: Vec<Node<'_>> = node.named_children(&mut cursor).collect();
1948 children.reverse();
1949 stack.extend(children);
1950 }
1951}
1952
1953fn binding_is_conditional(mut node: Node<'_>) -> bool {
1954 while let Some(parent) = node.parent() {
1955 if matches!(
1956 parent.kind(),
1957 "if_statement"
1958 | "try_statement"
1959 | "except_clause"
1960 | "match_statement"
1961 | "case_clause"
1962 | "for_statement"
1963 | "while_statement"
1964 ) {
1965 return true;
1966 }
1967 if matches!(
1968 parent.kind(),
1969 "module" | "function_definition" | "class_definition"
1970 ) {
1971 return false;
1972 }
1973 node = parent;
1974 }
1975 false
1976}
1977
1978fn first_identifier(node: Node<'_>) -> Option<Node<'_>> {
1979 identifier_extreme(node, false)
1980}
1981
1982fn last_identifier(node: Node<'_>) -> Option<Node<'_>> {
1983 identifier_extreme(node, true)
1984}
1985
1986fn identifier_extreme(node: Node<'_>, last: bool) -> Option<Node<'_>> {
1987 let mut best = None;
1988 let mut stack = vec![node];
1989 while let Some(node) = stack.pop() {
1990 if node.kind() == "identifier" {
1991 if best.is_none_or(|current: Node<'_>| {
1992 if last {
1993 node.start_byte() > current.start_byte()
1994 } else {
1995 node.start_byte() < current.start_byte()
1996 }
1997 }) {
1998 best = Some(node);
1999 }
2000 continue;
2001 }
2002 let mut cursor = node.walk();
2003 stack.extend(node.named_children(&mut cursor));
2004 }
2005 best
2006}
2007
2008fn reference_is_deferred_function_body(node: Node<'_>) -> bool {
2009 let site_start = node.start_byte();
2010 let site_end = node.end_byte();
2011 let mut current = node;
2012 while let Some(parent) = current.parent() {
2013 if matches!(parent.kind(), "function_definition" | "lambda")
2014 && parent
2015 .child_by_field_name("body")
2016 .is_some_and(|body| body.start_byte() <= site_start && site_end <= body.end_byte())
2017 {
2018 return true;
2019 }
2020 current = parent;
2021 }
2022 false
2023}
2024
2025pub fn collect_assigned_identifiers(node: Node<'_>, source: &str, out: &mut HashSet<String>) {
2026 let mut stack = vec![node];
2027 while let Some(node) = stack.pop() {
2028 if node.kind() == "identifier" {
2029 let text = slice(node, source).trim();
2030 if !text.is_empty() {
2031 out.insert(text.to_string());
2032 }
2033 continue;
2034 }
2035
2036 let mut cursor = node.walk();
2037 let mut children: Vec<Node<'_>> = node.named_children(&mut cursor).collect();
2038 children.reverse();
2039 stack.extend(children);
2040 }
2041}
2042
2043pub fn collect_scope_facts_from_parsed_source(
2044 graph: &PythonGraphSource<'_>,
2045 python: &dyn PythonUsageSource,
2046 file: &ProjectFile,
2047 source: &str,
2048 root: Node<'_>,
2049) -> PythonScopeFacts {
2050 let mut factory_return_types = collect_factory_return_types_from_root(root, source);
2051 collect_imported_factory_return_types(graph, python, file, &mut factory_return_types);
2052 collect_scope_facts_with_factory_returns(graph, file, source, &factory_return_types)
2053}
2054
2055fn collect_imported_factory_return_types(
2056 graph: &PythonGraphSource<'_>,
2057 python: &dyn PythonUsageSource,
2058 file: &ProjectFile,
2059 factory_return_types: &mut HashMap<String, String>,
2060) {
2061 let binder = python.import_binder_of(file);
2062 for (local, binding) in &binder.bindings {
2063 if !matches!(binding.kind, ImportKind::Named) {
2064 continue;
2065 }
2066 let Some(imported) = binding.imported_name.as_deref() else {
2067 continue;
2068 };
2069 let fqn = format!("{}.{}", binding.module_specifier, imported);
2070 let units =
2071 resolve_fqn_candidates(python, &fqn, |name| graph.index.definitions(name).collect());
2072 for unit in units {
2073 if unit.is_function() {
2074 if let Some(return_type) = callable_return_type_name(graph, python, &unit) {
2075 factory_return_types
2076 .entry(local.clone())
2077 .or_insert(return_type);
2078 }
2079 continue;
2080 }
2081 if !unit.is_class() {
2082 continue;
2083 }
2084 factory_return_types
2085 .entry(local.clone())
2086 .or_insert_with(|| unit.identifier().to_string());
2087 collect_imported_class_method_return_types(
2088 graph,
2089 python,
2090 local,
2091 &unit,
2092 factory_return_types,
2093 );
2094 }
2095 }
2096}
2097
2098fn collect_imported_class_method_return_types(
2099 graph: &PythonGraphSource<'_>,
2100 python: &dyn PythonSource,
2101 local_class_name: &str,
2102 class_unit: &CodeUnit,
2103 factory_return_types: &mut HashMap<String, String>,
2104) {
2105 for member in graph.index.direct_children(class_unit) {
2106 if !member.is_function() {
2107 continue;
2108 }
2109 let Some(return_type) = callable_return_type_name(graph, python, &member) else {
2110 continue;
2111 };
2112 factory_return_types
2113 .entry(format!("{}.{}", local_class_name, member.identifier()))
2114 .or_insert(return_type);
2115 }
2116}
2117
2118fn callable_return_type_name(
2119 graph: &PythonGraphSource<'_>,
2120 python: &dyn PythonSource,
2121 callable: &CodeUnit,
2122) -> Option<String> {
2123 if let Some(prepared) = python.prepared_syntax(callable.source()) {
2128 #[cfg(any(test, feature = "test-support"))]
2129 note_callable_return_type_lookup_for_test(true);
2130 return callable_return_type_name_in_tree(
2131 graph,
2132 callable,
2133 prepared.source(),
2134 prepared.tree().root_node(),
2135 );
2136 }
2137 #[cfg(any(test, feature = "test-support"))]
2138 note_callable_return_type_lookup_for_test(false);
2139 let source = graph.index.indexed_source(callable.source())?;
2140 declaration_source_slices(graph, callable, &source)
2141 .into_iter()
2142 .find_map(|declaration_source| {
2143 let mut parser = Parser::new();
2144 parser
2145 .set_language(&tree_sitter_python::LANGUAGE.into())
2146 .ok()?;
2147 let tree = parser.parse(declaration_source, None)?;
2148 let function = first_function_definition(tree.root_node())?;
2149 factory_return_type(function, declaration_source)
2150 })
2151}
2152
2153#[cfg(any(test, feature = "test-support"))]
2159#[derive(Debug, Default, Clone, Copy, PartialEq, Eq)]
2160pub struct CallableReturnTypeLookupCounts {
2161 pub prepared: usize,
2163 pub reparsed: usize,
2165}
2166
2167#[cfg(any(test, feature = "test-support"))]
2168thread_local! {
2169 static CALLABLE_RETURN_TYPE_LOOKUPS_FOR_TEST: std::cell::Cell<CallableReturnTypeLookupCounts> =
2170 const { std::cell::Cell::new(CallableReturnTypeLookupCounts { prepared: 0, reparsed: 0 }) };
2171}
2172
2173#[cfg(any(test, feature = "test-support"))]
2174fn note_callable_return_type_lookup_for_test(from_prepared_syntax: bool) {
2175 CALLABLE_RETURN_TYPE_LOOKUPS_FOR_TEST.with(|counts| {
2176 let mut observed = counts.get();
2177 if from_prepared_syntax {
2178 observed.prepared += 1;
2179 } else {
2180 observed.reparsed += 1;
2181 }
2182 counts.set(observed);
2183 });
2184}
2185
2186#[cfg(any(test, feature = "test-support"))]
2189pub fn with_callable_return_type_lookup_counter_for_test<T>(
2190 body: impl FnOnce() -> T,
2191) -> (T, CallableReturnTypeLookupCounts) {
2192 CALLABLE_RETURN_TYPE_LOOKUPS_FOR_TEST.with(|counts| {
2193 counts.set(CallableReturnTypeLookupCounts::default());
2194 let result = body();
2195 let observed = counts.get();
2196 counts.set(CallableReturnTypeLookupCounts::default());
2197 (result, observed)
2198 })
2199}
2200
2201fn callable_return_type_name_in_tree(
2207 graph: &PythonGraphSource<'_>,
2208 callable: &CodeUnit,
2209 source: &str,
2210 root: Node<'_>,
2211) -> Option<String> {
2212 let mut ranges = graph.index.ranges(callable);
2213 ranges.sort_by_key(|range| range.start_byte);
2214 ranges.into_iter().find_map(|range| {
2215 let declaration = root.descendant_for_byte_range(range.start_byte, range.end_byte)?;
2216 let function = first_function_definition(declaration)?;
2217 factory_return_type(function, source)
2218 })
2219}
2220
2221fn first_function_definition(root: Node<'_>) -> Option<Node<'_>> {
2222 let mut stack = vec![root];
2223 while let Some(node) = stack.pop() {
2224 if node.kind() == "function_definition" {
2225 return Some(node);
2226 }
2227 let mut cursor = node.walk();
2228 let mut children: Vec<Node<'_>> = node.named_children(&mut cursor).collect();
2229 children.reverse();
2230 stack.extend(children);
2231 }
2232 None
2233}
2234
2235fn collect_scope_facts_with_factory_returns(
2236 graph: &PythonGraphSource<'_>,
2237 file: &ProjectFile,
2238 source: &str,
2239 factory_return_types: &HashMap<String, String>,
2240) -> PythonScopeFacts {
2241 let declarations = graph.index.declarations(file);
2242 let mut class_facts_by_name: HashMap<String, LocalBindingsSnapshot<String>> =
2243 HashMap::default();
2244 for declaration in declarations
2245 .iter()
2246 .filter(|declaration| declaration.is_class())
2247 {
2248 let Some(declaration_source) = declaration_source(graph, declaration, source) else {
2249 continue;
2250 };
2251 let facts = collect_scope_facts_from_source(
2252 &declaration_source,
2253 ScopeFactTraversal::Class,
2254 true,
2255 Some(declaration.short_name()),
2256 factory_return_types,
2257 );
2258 class_facts_by_name.insert(
2259 declaration.short_name().to_string(),
2260 facts.filtered_visible_bindings(|symbol, _| symbol.starts_with("self.")),
2261 );
2262 }
2263
2264 let mut scope_facts = HashMap::default();
2265 for declaration in declarations
2266 .iter()
2267 .filter(|declaration| declaration.is_function())
2268 {
2269 let Some(declaration_source) = declaration_source(graph, declaration, source) else {
2270 continue;
2271 };
2272 let owner = declaration
2277 .short_name()
2278 .rsplit_once('.')
2279 .map(|(owner, _)| owner);
2280 let mut facts = collect_scope_facts_from_source(
2281 &declaration_source,
2282 ScopeFactTraversal::Function,
2283 false,
2284 owner,
2285 factory_return_types,
2286 );
2287 if let Some(owner) = owner
2288 && let Some(class_facts) = class_facts_by_name.get(owner)
2289 {
2290 facts = facts.merged_with_visible(class_facts);
2291 }
2292 scope_facts.insert(declaration.clone(), facts);
2293 }
2294
2295 for declaration in declarations.iter().filter(|d| d.is_module()) {
2300 let Some(declaration_source) = declaration_source(graph, declaration, source) else {
2301 continue;
2302 };
2303 let facts = collect_scope_facts_from_source(
2304 &declaration_source,
2305 ScopeFactTraversal::Module,
2306 false,
2307 None,
2308 factory_return_types,
2309 );
2310 scope_facts.insert(declaration.clone(), facts);
2311 }
2312 scope_facts
2313}
2314
2315fn declaration_source(
2316 graph: &PythonGraphSource<'_>,
2317 declaration: &CodeUnit,
2318 file_source: &str,
2319) -> Option<String> {
2320 let slices = declaration_source_slices(graph, declaration, file_source);
2321 (!slices.is_empty()).then(|| slices.join("\n\n"))
2322}
2323
2324fn declaration_source_slices<'a>(
2325 graph: &PythonGraphSource<'_>,
2326 declaration: &CodeUnit,
2327 file_source: &'a str,
2328) -> Vec<&'a str> {
2329 let mut ranges = graph.index.ranges(declaration);
2330 ranges.sort_by_key(|range| range.start_byte);
2331 ranges
2332 .into_iter()
2333 .filter_map(|range| file_source.get(range.start_byte..range.end_byte))
2334 .collect()
2335}
2336
2337fn collect_scope_facts_from_source(
2338 source: &str,
2339 traversal: ScopeFactTraversal,
2340 allow_self_receivers: bool,
2341 current_class: Option<&str>,
2342 factory_return_types: &HashMap<String, String>,
2343) -> LocalBindingsSnapshot<String> {
2344 let events = collect_scope_fact_events(source, traversal);
2345 collect_scope_facts_from_events(
2346 &events,
2347 allow_self_receivers,
2348 current_class,
2349 factory_return_types,
2350 )
2351}
2352
2353pub fn collect_function_scope_facts_from_node(
2354 function: Node<'_>,
2355 source: &str,
2356) -> LocalBindingsSnapshot<String> {
2357 let mut events = Vec::new();
2358 if function.kind() == "lambda" {
2359 if let Some(parameters) = function.child_by_field_name("parameters") {
2360 collect_parameter_events(parameters, source, &mut events);
2361 }
2362 } else {
2363 collect_scope_fact_events_from_node(
2364 function,
2365 source,
2366 ScopeFactTraversal::Function,
2367 &mut events,
2368 );
2369 }
2370 collect_scope_facts_from_events(&events, false, None, &HashMap::default())
2371}
2372
2373fn collect_scope_facts_from_events(
2374 events: &[ScopeFactEvent],
2375 allow_self_receivers: bool,
2376 current_class: Option<&str>,
2377 factory_return_types: &HashMap<String, String>,
2378) -> LocalBindingsSnapshot<String> {
2379 let mut engine = LocalInferenceEngine::new(LocalInferenceConfig::default());
2380 let globals: HashSet<&str> = events
2381 .iter()
2382 .filter_map(|event| match event {
2383 ScopeFactEvent::Global { symbol } => Some(symbol.as_str()),
2384 _ => None,
2385 })
2386 .collect();
2387 let nonlocals: HashSet<&str> = events
2388 .iter()
2389 .filter_map(|event| match event {
2390 ScopeFactEvent::Nonlocal { symbol } => Some(symbol.as_str()),
2391 _ => None,
2392 })
2393 .collect();
2394 for symbol in &nonlocals {
2395 engine.declare_shadow((*symbol).to_string());
2396 }
2397 for event in events {
2398 if let ScopeFactEvent::Parameter { symbol, .. } = event
2399 && !globals.contains(symbol.as_str())
2400 && !nonlocals.contains(symbol.as_str())
2401 && !engine.is_shadowed(symbol)
2402 {
2403 engine.declare_shadow(symbol.clone());
2404 }
2405 }
2406
2407 let mut changed = true;
2408 while changed {
2409 changed = false;
2410 let mut aliases = Vec::new();
2411 for event in events {
2412 match event {
2413 ScopeFactEvent::Parameter {
2414 symbol,
2415 annotation: Some(annotation),
2416 }
2417 | ScopeFactEvent::Annotation { symbol, annotation } => {
2418 if globals.contains(symbol.as_str()) || nonlocals.contains(symbol.as_str()) {
2419 continue;
2420 }
2421 apply_annotation_event(
2422 symbol,
2423 annotation,
2424 allow_self_receivers,
2425 &mut engine,
2426 &mut changed,
2427 );
2428 }
2429 ScopeFactEvent::Parameter {
2430 annotation: None, ..
2431 } => {}
2432 ScopeFactEvent::Assignment { lhs, rhs } => {
2433 if globals.contains(lhs.as_str()) {
2434 continue;
2435 }
2436 if !engine.is_shadowed(lhs) {
2437 engine.declare_shadow(lhs.clone());
2438 }
2439 if lhs.starts_with("self.") && !allow_self_receivers {
2440 continue;
2441 }
2442
2443 match rhs {
2444 AssignmentRhs::Call(callee) => {
2445 if !engine.is_shadowed(callee) {
2446 if let Some(receiver_type) = factory_return_type_for_callee(
2447 callee,
2448 current_class,
2449 factory_return_types,
2450 ) && engine.resolve_symbol(lhs).is_unknown()
2451 {
2452 engine.seed_symbol(lhs.clone(), receiver_type.clone());
2453 changed = true;
2454 continue;
2455 }
2456
2457 if let Some(receiver_type) = normalized_receiver_type(callee)
2458 && engine.resolve_symbol(lhs).is_unknown()
2459 {
2460 engine.seed_symbol(lhs.clone(), receiver_type);
2461 changed = true;
2462 continue;
2463 }
2464 }
2465 }
2466 AssignmentRhs::Symbol(rhs_symbol) => {
2467 if !engine.is_shadowed(rhs_symbol)
2468 && let Some(receiver_type) = normalized_receiver_type(rhs_symbol)
2469 && engine.resolve_symbol(lhs).is_unknown()
2470 {
2471 engine.seed_symbol(lhs.clone(), receiver_type);
2472 changed = true;
2473 continue;
2474 }
2475
2476 if let SymbolResolution::Precise(targets) =
2477 engine.resolve_symbol(rhs_symbol)
2478 && !targets.is_empty()
2479 {
2480 aliases.push((lhs.clone(), rhs_symbol.clone()));
2481 }
2482 }
2483 AssignmentRhs::Unknown => {}
2484 }
2485 }
2486 ScopeFactEvent::Global { .. } | ScopeFactEvent::Nonlocal { .. } => {}
2487 }
2488 }
2489 let before = engine.snapshot();
2490 engine.apply_aliases_until_stable(aliases);
2491 if engine.snapshot() != before {
2492 changed = true;
2493 }
2494 }
2495
2496 engine.snapshot()
2497}
2498
2499fn factory_return_type_for_callee<'a>(
2500 callee: &str,
2501 current_class: Option<&str>,
2502 factory_return_types: &'a HashMap<String, String>,
2503) -> Option<&'a String> {
2504 if let Some(receiver_type) = factory_return_types.get(callee) {
2505 return Some(receiver_type);
2506 }
2507 let class_name = current_class?;
2508 let method = callee
2509 .strip_prefix("self.")
2510 .or_else(|| callee.strip_prefix("cls."))?;
2511 factory_return_types.get(&format!("{class_name}.{method}"))
2512}
2513
2514fn apply_annotation_event(
2515 symbol: &str,
2516 annotation: &str,
2517 allow_self_receivers: bool,
2518 engine: &mut LocalInferenceEngine<String>,
2519 changed: &mut bool,
2520) {
2521 if symbol.starts_with("self.") && !allow_self_receivers {
2522 return;
2523 }
2524 if let Some(receiver_type) = normalized_receiver_type(annotation)
2525 && engine.resolve_symbol(symbol).is_unknown()
2526 {
2527 engine.seed_symbol(symbol.to_string(), receiver_type);
2528 *changed = true;
2529 }
2530}
2531
2532enum ScopeFactEvent {
2533 Global {
2534 symbol: String,
2535 },
2536 Nonlocal {
2537 symbol: String,
2538 },
2539 Parameter {
2540 symbol: String,
2541 annotation: Option<String>,
2542 },
2543 Annotation {
2544 symbol: String,
2545 annotation: String,
2546 },
2547 Assignment {
2548 lhs: String,
2549 rhs: AssignmentRhs,
2550 },
2551}
2552
2553enum AssignmentRhs {
2554 Symbol(String),
2555 Call(String),
2556 Unknown,
2557}
2558
2559#[derive(Clone, Copy)]
2560enum ScopeFactTraversal {
2561 Module,
2562 Function,
2563 Class,
2564}
2565
2566fn collect_scope_fact_events(source: &str, traversal: ScopeFactTraversal) -> Vec<ScopeFactEvent> {
2567 if source.trim().is_empty() {
2568 return Vec::new();
2569 }
2570
2571 let mut parser = Parser::new();
2572 if parser
2573 .set_language(&tree_sitter_python::LANGUAGE.into())
2574 .is_err()
2575 {
2576 return Vec::new();
2577 }
2578 let Some(tree) = parser.parse(source, None) else {
2579 return Vec::new();
2580 };
2581
2582 let mut events = Vec::new();
2583 collect_scope_fact_events_from_node(tree.root_node(), source, traversal, &mut events);
2584 events
2585}
2586
2587fn collect_scope_fact_events_from_node(
2588 root: Node<'_>,
2589 source: &str,
2590 traversal: ScopeFactTraversal,
2591 events: &mut Vec<ScopeFactEvent>,
2592) {
2593 let mut stack = vec![(root, false)];
2594 while let Some((node, inside_function)) = stack.pop() {
2595 let next_inside_function = match traversal {
2596 ScopeFactTraversal::Module => {
2597 if matches!(
2598 node.kind(),
2599 "function_definition" | "class_definition" | "lambda"
2600 ) {
2601 continue;
2602 }
2603 false
2604 }
2605 ScopeFactTraversal::Function => match node.kind() {
2606 "function_definition" if inside_function => continue,
2607 "function_definition" => true,
2608 "class_definition" | "lambda" => continue,
2609 _ => inside_function,
2610 },
2611 ScopeFactTraversal::Class => inside_function,
2612 };
2613 if matches!(traversal, ScopeFactTraversal::Function) && !next_inside_function {
2614 let mut cursor = node.walk();
2615 let mut children: Vec<Node<'_>> = node.named_children(&mut cursor).collect();
2616 children.reverse();
2617 stack.extend(children.into_iter().map(|child| (child, false)));
2618 continue;
2619 }
2620 match node.kind() {
2621 "global_statement" => collect_scope_directive_events(node, source, events, |symbol| {
2622 ScopeFactEvent::Global { symbol }
2623 }),
2624 "nonlocal_statement" => {
2625 collect_scope_directive_events(node, source, events, |symbol| {
2626 ScopeFactEvent::Nonlocal { symbol }
2627 })
2628 }
2629 "parameters" | "lambda_parameters" => collect_parameter_events(node, source, events),
2630 "assignment" => collect_assignment_events(node, source, events),
2631 _ => {}
2632 }
2633
2634 let mut cursor = node.walk();
2635 let mut children: Vec<Node<'_>> = node.named_children(&mut cursor).collect();
2636 children.reverse();
2637 stack.extend(
2638 children
2639 .into_iter()
2640 .map(|child| (child, next_inside_function)),
2641 );
2642 }
2643}
2644
2645fn collect_scope_directive_events(
2646 node: Node<'_>,
2647 source: &str,
2648 events: &mut Vec<ScopeFactEvent>,
2649 make_event: impl Fn(String) -> ScopeFactEvent,
2650) {
2651 let mut cursor = node.walk();
2652 for identifier in node
2653 .named_children(&mut cursor)
2654 .filter(|child| child.kind() == "identifier")
2655 {
2656 let Some(symbol) = non_empty_node_text(identifier, source) else {
2657 continue;
2658 };
2659 events.push(make_event(symbol));
2660 }
2661}
2662
2663fn collect_parameter_events(node: Node<'_>, source: &str, events: &mut Vec<ScopeFactEvent>) {
2664 let mut cursor = node.walk();
2665 for child in node.named_children(&mut cursor) {
2666 if child.kind() == "type_parameter" {
2667 continue;
2668 }
2669 let Some(symbol) = parameter_symbol(child, source) else {
2670 continue;
2671 };
2672 if matches!(symbol.as_str(), "self" | "cls" | "/") {
2673 continue;
2674 }
2675 let annotation = child
2676 .child_by_field_name("type")
2677 .map(|annotation| slice(annotation, source).trim().to_string())
2678 .filter(|annotation| !annotation.is_empty());
2679 events.push(ScopeFactEvent::Parameter { symbol, annotation });
2680 }
2681}
2682
2683fn parameter_symbol(node: Node<'_>, source: &str) -> Option<String> {
2684 if node.kind() == "identifier" {
2685 return non_empty_node_text(node, source);
2686 }
2687 if let Some(name) = node.child_by_field_name("name") {
2688 return non_empty_node_text(name, source);
2689 }
2690 let mut cursor = node.walk();
2691 node.named_children(&mut cursor)
2692 .find(|child| child.kind() == "identifier")
2693 .and_then(|identifier| non_empty_node_text(identifier, source))
2694}
2695
2696fn collect_assignment_events(node: Node<'_>, source: &str, events: &mut Vec<ScopeFactEvent>) {
2697 let Some(left) = node.child_by_field_name("left") else {
2698 return;
2699 };
2700 let Some(lhs) = receiver_symbol(left, source) else {
2701 return;
2702 };
2703
2704 if let Some(annotation) = node
2705 .child_by_field_name("type")
2706 .map(|annotation| slice(annotation, source).trim().to_string())
2707 .filter(|annotation| !annotation.is_empty())
2708 {
2709 events.push(ScopeFactEvent::Annotation {
2710 symbol: lhs,
2711 annotation,
2712 });
2713 return;
2714 }
2715
2716 let rhs = node
2717 .child_by_field_name("right")
2718 .and_then(|right| rhs_symbol(right, source))
2719 .unwrap_or(AssignmentRhs::Unknown);
2720 events.push(ScopeFactEvent::Assignment { lhs, rhs });
2721}
2722
2723fn receiver_symbol(node: Node<'_>, source: &str) -> Option<String> {
2724 match node.kind() {
2725 "identifier" | "attribute" => non_empty_node_text(node, source),
2726 _ => None,
2727 }
2728}
2729
2730fn rhs_symbol(node: Node<'_>, source: &str) -> Option<AssignmentRhs> {
2731 match node.kind() {
2732 "identifier" | "attribute" => non_empty_node_text(node, source).map(AssignmentRhs::Symbol),
2733 "call" => node
2734 .child_by_field_name("function")
2735 .or_else(|| node.named_child(0))
2736 .and_then(|callee| receiver_symbol(callee, source))
2737 .map(AssignmentRhs::Call),
2738 _ => None,
2739 }
2740}
2741
2742fn non_empty_node_text(node: Node<'_>, source: &str) -> Option<String> {
2743 let text = slice(node, source).trim();
2744 (!text.is_empty()).then(|| text.to_string())
2745}
2746
2747fn collect_factory_return_types_from_root(root: Node<'_>, source: &str) -> HashMap<String, String> {
2748 let mut returns = HashMap::default();
2749 let mut stack = vec![(root, None::<String>)];
2750 while let Some((node, class_name)) = stack.pop() {
2751 match node.kind() {
2752 "class_definition" => {
2753 let next_class = node
2754 .child_by_field_name("name")
2755 .and_then(|name| non_empty_node_text(name, source))
2756 .or(class_name);
2757 push_factory_index_children(node, next_class, &mut stack);
2758 }
2759 "function_definition" => {
2760 if let Some(name) = node
2761 .child_by_field_name("name")
2762 .and_then(|name| non_empty_node_text(name, source))
2763 && let Some(return_type) = factory_return_type(node, source)
2764 {
2765 let key = class_name
2766 .as_ref()
2767 .map(|class| format!("{class}.{name}"))
2768 .unwrap_or(name);
2769 returns.insert(key, return_type);
2770 }
2771 }
2772 _ => push_factory_index_children(node, class_name, &mut stack),
2773 }
2774 }
2775 returns
2776}
2777
2778fn push_factory_index_children<'tree>(
2779 node: Node<'tree>,
2780 class_name: Option<String>,
2781 stack: &mut Vec<(Node<'tree>, Option<String>)>,
2782) {
2783 let mut cursor = node.walk();
2784 let mut children: Vec<Node<'tree>> = node.named_children(&mut cursor).collect();
2785 children.reverse();
2786 stack.extend(
2787 children
2788 .into_iter()
2789 .map(|child| (child, class_name.clone())),
2790 );
2791}
2792
2793fn factory_return_type(function: Node<'_>, source: &str) -> Option<String> {
2794 if let Some(return_type) = function.child_by_field_name("return_type") {
2795 return receiver_type_from_annotation_node(return_type, source);
2796 }
2797
2798 let body = function.child_by_field_name("body")?;
2799 let mut candidates = HashSet::default();
2800 let mut saw_return = false;
2801 let mut saw_unknown_return = false;
2802 let mut stack = vec![body];
2803 while let Some(node) = stack.pop() {
2804 if node != body && matches!(node.kind(), "function_definition" | "class_definition") {
2805 continue;
2806 }
2807 if node.kind() == "return_statement" {
2808 saw_return = true;
2809 match node
2810 .named_child(0)
2811 .and_then(|value| returned_receiver_type(value, source))
2812 {
2813 Some(returned_type) => {
2814 candidates.insert(returned_type);
2815 }
2816 None => saw_unknown_return = true,
2817 }
2818 }
2819 let mut cursor = node.walk();
2820 let mut children: Vec<Node<'_>> = node.named_children(&mut cursor).collect();
2821 children.reverse();
2822 stack.extend(children);
2823 }
2824 if !saw_return || saw_unknown_return {
2825 return None;
2826 }
2827 (candidates.len() == 1)
2828 .then(|| candidates.into_iter().next())
2829 .flatten()
2830}
2831
2832fn receiver_type_from_annotation_node(annotation: Node<'_>, source: &str) -> Option<String> {
2838 match annotation.kind() {
2839 "type" => receiver_type_from_annotation_node(annotation.named_child(0)?, source),
2840 "identifier" | "attribute" | "member_type" | "string" => {
2841 normalized_receiver_type(slice(annotation, source).trim())
2842 }
2843 "generic_type" => {
2844 let base = annotation.named_child(0)?;
2845 if optional_annotation_wrapper(base, source) {
2846 let parameter = annotation.named_child(1)?;
2847 return receiver_type_from_annotation_node(parameter.named_child(0)?, source);
2848 }
2849 receiver_type_from_annotation_node(base, source)
2850 }
2851 "subscript" => {
2852 let value = annotation.child_by_field_name("value")?;
2853 if optional_annotation_wrapper(value, source) {
2854 let inner = annotation.child_by_field_name("subscript")?;
2855 return receiver_type_from_annotation_node(inner, source);
2856 }
2857 normalized_receiver_type(slice(value, source).trim())
2858 }
2859 _ => None,
2860 }
2861}
2862
2863fn optional_annotation_wrapper(node: Node<'_>, source: &str) -> bool {
2864 match node.kind() {
2865 "identifier" => slice(node, source) == "Optional",
2866 "attribute" => {
2867 let (Some(object), Some(attribute)) = (
2868 node.child_by_field_name("object"),
2869 node.child_by_field_name("attribute"),
2870 ) else {
2871 return false;
2872 };
2873 object.kind() == "identifier"
2874 && attribute.kind() == "identifier"
2875 && slice(object, source) == "typing"
2876 && slice(attribute, source) == "Optional"
2877 }
2878 _ => false,
2879 }
2880}
2881
2882fn returned_receiver_type(node: Node<'_>, source: &str) -> Option<String> {
2883 let raw = match node.kind() {
2884 "identifier" => non_empty_node_text(node, source),
2885 "call" => node
2886 .child_by_field_name("function")
2887 .or_else(|| node.named_child(0))
2888 .filter(|callee| callee.kind() == "identifier")
2889 .and_then(|callee| non_empty_node_text(callee, source)),
2890 _ => None,
2891 }?;
2892 normalized_receiver_type(&raw)
2893}
2894
2895#[cfg(test)]
2896mod tests {
2897 use super::*;
2898 use std::path::PathBuf;
2899
2900 #[test]
2901 fn lambda_scope_facts_preserve_untyped_parameter_shadowing() {
2902 let source = "shadowed = lambda method: method.signature\n";
2903 let mut parser = Parser::new();
2904 parser
2905 .set_language(&tree_sitter_python::LANGUAGE.into())
2906 .unwrap();
2907 let tree = parser.parse(source, None).unwrap();
2908 let mut nodes = vec![tree.root_node()];
2909 let lambda = loop {
2910 let node = nodes.pop().unwrap();
2911 if node.kind() == "lambda" {
2912 break node;
2913 }
2914 let mut cursor = node.walk();
2915 nodes.extend(node.named_children(&mut cursor));
2916 };
2917
2918 let facts = collect_function_scope_facts_from_node(lambda, source);
2919
2920 assert!(facts.is_shadowed("method"));
2921 assert!(facts.resolution_for("method").is_unknown());
2922 }
2923
2924 #[test]
2925 fn pre_cancelled_graph_build_skips_python_file_parsing() {
2926 let temp = tempfile::tempdir().unwrap();
2927 let root = temp.path().canonicalize().unwrap();
2928 std::fs::write(root.join("target.py"), "def target():\n pass\n").unwrap();
2929 let file = ProjectFile::new(root.clone(), PathBuf::from("target.py"));
2930 let files = [file.clone()].into_iter().collect();
2931 let cancellation = CancellationToken::default();
2932 cancellation.cancel();
2933
2934 let graph = build_python_graph(&files, &file, Some(&cancellation));
2935
2936 assert!(graph.parsed.is_empty());
2937 }
2938
2939 #[test]
2940 fn graph_build_parses_only_candidates_and_target_not_transitive_imports() {
2941 let temp = tempfile::tempdir().unwrap();
2942 let root = temp.path().canonicalize().unwrap();
2943 std::fs::write(root.join("target.py"), "from dependency import value\n").unwrap();
2944 std::fs::write(
2945 root.join("candidate.py"),
2946 "from transitively_imported import value\n",
2947 )
2948 .unwrap();
2949 std::fs::write(root.join("dependency.py"), "value = 1\n").unwrap();
2950 std::fs::write(root.join("transitively_imported.py"), "value = 2\n").unwrap();
2951 let target = ProjectFile::new(root.clone(), PathBuf::from("target.py"));
2952 let candidate = ProjectFile::new(root.clone(), PathBuf::from("candidate.py"));
2953 let dependency = ProjectFile::new(root.clone(), PathBuf::from("dependency.py"));
2954 let transitive = ProjectFile::new(root.clone(), PathBuf::from("transitively_imported.py"));
2955 let candidates = [candidate.clone()].into_iter().collect();
2956
2957 let graph = build_python_graph(&candidates, &target, None);
2958
2959 assert_eq!(graph.parsed.len(), 2);
2960 assert!(graph.parsed.contains_key(&target));
2961 assert!(graph.parsed.contains_key(&candidate));
2962 assert!(!graph.parsed.contains_key(&dependency));
2963 assert!(!graph.parsed.contains_key(&transitive));
2964 }
2965}