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