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