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