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