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