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