1use crate::graph::PythonGraphSource;
9use crate::graph_support::PythonUsageSource;
10use crate::imports::resolve_fqn_candidates;
11use crate::syntax::{
12 attribute_chain, python_deferred_annotation_identifier_ranges, python_deferred_annotation_tree,
13 python_node_is_in_annotation,
14};
15use crate::usage_index::{
16 ModuleBindingEventKind, ModuleBindingTimeline, usage_resolve_module_files, usage_seeds,
17 visible_module_binding_suffix,
18};
19use brokk_bifrost_core::analyzer::common::node_source_text as node_text;
20use brokk_bifrost_core::analyzer::fq_name::{FqName, SegmentKind, segment_interner};
21use brokk_bifrost_core::analyzer::symbol_path::parse_symbol_path;
22use brokk_bifrost_core::analyzer::usages::model::{ExportEntry, ImportBinder, ImportKind};
23use brokk_bifrost_core::analyzer::usages::{ImportEdge, ImportEdgeKind};
24use brokk_bifrost_core::analyzer::{
25 CodeUnit, CodeUnitIndex, DefinitionLanguageScope, Language, ProjectFile, Range,
26 RelationalDefinitionQuery, RelationalDefinitionQuestion, RelationalDefinitionValue,
27 RelationalName,
28};
29use brokk_bifrost_core::hash::HashSet;
30use std::collections::BTreeSet;
31use tree_sitter::Node;
32
33const MAX_MEMBER_OWNER_FRONTIER: usize = 512;
34
35pub(crate) fn resolved_member_declarations(
39 graph: &PythonGraphSource<'_>,
40 receiver: &CodeUnit,
41 member: &str,
42) -> Vec<CodeUnit> {
43 let mut visited = HashSet::default();
44 let mut frontier = vec![receiver.clone()];
45 while !frontier.is_empty() && visited.len() < MAX_MEMBER_OWNER_FRONTIER {
46 let mut declarations = Vec::new();
47 let mut next = Vec::new();
48 for owner in frontier {
49 if !visited.insert(owner.clone()) {
50 continue;
51 }
52 let fqn = format!("{}.{member}", owner.fq_name());
53 declarations.extend(
54 graph
55 .index
56 .definitions(&fqn)
57 .filter(|candidate| graph.index.parent_of(candidate).as_ref() == Some(&owner)),
58 );
59 if let Some(hierarchy) = graph.hierarchy {
60 next.extend(hierarchy.get_direct_ancestors(&owner));
61 }
62 }
63 declarations.sort();
64 declarations.dedup();
65 if !declarations.is_empty() {
66 return declarations;
67 }
68 next.sort();
69 next.dedup();
70 frontier = next;
71 }
72 Vec::new()
73}
74
75pub fn infer_export_names(python: &dyn PythonUsageSource, target: &CodeUnit) -> BTreeSet<String> {
76 if target_owner_code_unit(python, target).is_some() {
77 let owner_name = top_level_identifier(python, target);
78 let owner_exports =
79 infer_export_names_for_local(python, target, target.source(), &owner_name);
80 if !owner_exports.is_empty() {
81 return owner_exports;
82 }
83 }
84
85 infer_export_names_for_local(python, target, target.source(), target.identifier())
86}
87
88pub fn infer_usage_seeds(
89 python: &dyn PythonUsageSource,
90 target: &CodeUnit,
91 seed_names: BTreeSet<String>,
92) -> BTreeSet<(ProjectFile, String)> {
93 let mut seeds = BTreeSet::new();
94 for seed_name in &seed_names {
95 seeds.extend(usage_seeds(python, target.source(), seed_name));
96 }
97 if seeds.is_empty()
98 && seed_names.contains(target.identifier())
99 && is_module_level_target_identifier(python, target, target.source(), target.identifier())
100 {
101 seeds.insert((target.source().clone(), target.identifier().to_string()));
102 }
103 seeds
104}
105
106fn infer_export_names_for_local(
107 python: &dyn PythonUsageSource,
108 target: &CodeUnit,
109 file: &ProjectFile,
110 local_name: &str,
111) -> BTreeSet<String> {
112 let index = python.export_index_of(file);
113 let mut export_names = BTreeSet::new();
114 if index.exports_by_name.contains_key(local_name) {
115 export_names.insert(local_name.to_string());
116 }
117 for (export_name, entry) in &index.exports_by_name {
118 if matches!(entry, ExportEntry::Local { local_name: name } if name == local_name) {
119 export_names.insert(export_name.clone());
120 }
121 }
122 if export_names.is_empty()
123 && is_module_level_target_identifier(python, target, file, local_name)
124 {
125 export_names.insert(local_name.to_string());
126 }
127 export_names
128}
129
130fn is_module_level_target_identifier(
131 python: &dyn PythonUsageSource,
132 target: &CodeUnit,
133 file: &ProjectFile,
134 local_name: &str,
135) -> bool {
136 target.source() == file
137 && target.identifier() == local_name
138 && python
139 .parent_of(target)
140 .is_some_and(|parent| parent.is_module() && parent.source() == file)
141}
142
143pub fn top_level_identifier(index: &dyn CodeUnitIndex, target: &CodeUnit) -> String {
150 let mut current = target.clone();
151 while let Some(parent) = index.parent_of(¤t) {
152 if !parent.is_class() {
153 break;
154 }
155 current = parent;
156 }
157 current.identifier().to_string()
158}
159
160pub fn member_name(index: &dyn CodeUnitIndex, target: &CodeUnit) -> Option<String> {
161 target_owner_code_unit(index, target).map(|_| target.identifier().to_string())
162}
163
164pub fn target_owner_code_unit(index: &dyn CodeUnitIndex, target: &CodeUnit) -> Option<CodeUnit> {
165 index
166 .parent_of(target)
167 .filter(|parent| parent.source() == target.source() && parent.is_class())
168}
169
170pub fn resolve_receiver_type(
171 graph: &PythonGraphSource<'_>,
172 python: &dyn PythonUsageSource,
173 file: &ProjectFile,
174 raw_type: &str,
175 target_self_file: bool,
176) -> Option<CodeUnit> {
177 let raw_type = raw_type.trim();
178 if raw_type.is_empty() || raw_type.contains('|') {
179 return None;
180 }
181 if raw_type.contains('.') {
182 let candidates = resolve_fqn_candidates(python, Some(file), raw_type, |name| {
183 graph.index.definitions(name).collect()
184 })
185 .into_iter()
186 .filter(CodeUnit::is_class)
187 .collect::<Vec<_>>();
188 return (candidates.len() == 1)
189 .then(|| candidates.into_iter().next())
190 .flatten();
191 }
192
193 if let Some(binding) = python.import_binder_of(file).bindings.get(raw_type)
194 && binding.kind == ImportKind::Named
195 && let Some(imported) = binding.imported_name.as_ref()
196 {
197 let fqn = format!("{}.{}", binding.module_specifier, imported);
198 if let Some(class) = resolve_fqn_candidates(python, Some(file), &fqn, |name| {
199 graph.index.definitions(name).collect()
200 })
201 .into_iter()
202 .find(CodeUnit::is_class)
203 {
204 return Some(class);
205 }
206 }
207
208 if let Some(provider) = graph.imports
209 && let Some(imported) = provider
210 .imported_code_units_of(file)
211 .iter()
212 .find(|code_unit| code_unit.identifier() == raw_type && code_unit.is_class())
213 {
214 return Some(imported.clone());
215 }
216
217 graph
218 .index
219 .declarations(file)
220 .into_iter()
221 .find(|code_unit| code_unit.identifier() == raw_type && code_unit.is_class())
222 .or_else(|| {
223 if !target_self_file {
224 return None;
225 }
226 resolve_indexed_receiver_type(graph, file, raw_type)
227 })
228}
229
230fn resolve_bare_annotation_symbol(
231 graph: &PythonGraphSource<'_>,
232 python: &dyn PythonUsageSource,
233 file: &ProjectFile,
234 source: &str,
235 node: Node<'_>,
236 raw_symbol: &str,
237) -> Vec<CodeUnit> {
238 let raw_symbol = raw_symbol.trim();
239 if raw_symbol.is_empty() {
240 return Vec::new();
241 }
242
243 if let Some(owner) = annotation_scope_owner_class(graph, file, source, node) {
244 let owner_candidates: Vec<_> = exact_owner_annotation_members(graph, &owner, raw_symbol)
245 .into_iter()
246 .filter(|candidate| !candidate.is_function())
247 .collect();
248 if !owner_candidates.is_empty() {
249 return owner_candidates;
250 }
251 }
252
253 let mut candidates = Vec::new();
254 if let Some(binding) = python.import_binder_of(file).bindings.get(raw_symbol)
255 && binding.kind == ImportKind::Named
256 && let Some(imported) = binding.imported_name.as_ref()
257 {
258 let fqn = format!("{}.{}", binding.module_specifier, imported);
259 let mut imported_candidates = resolve_fqn_candidates(python, Some(file), &fqn, |name| {
260 graph
261 .index
262 .definitions(name)
263 .filter(|candidate| candidate.source().language() == Language::Python)
264 .collect()
265 });
266 imported_candidates.retain(|candidate| {
267 !candidate.is_module() || candidate.fq_name() != binding.module_specifier
268 });
269 candidates.extend(imported_candidates);
270 }
271
272 candidates.extend(
273 graph
274 .index
275 .top_level_declarations(file)
276 .into_iter()
277 .filter(|code_unit| {
278 !code_unit.is_module()
279 && code_unit.identifier() == raw_symbol
280 && graph
281 .index
282 .parent_of(code_unit)
283 .is_some_and(|parent| parent.is_module())
284 }),
285 );
286
287 candidates.retain(|candidate| !candidate.is_function());
288 candidates.sort();
289 candidates.dedup();
290 candidates
291}
292
293pub fn annotation_reference_candidates(
300 graph: &PythonGraphSource<'_>,
301 python: &dyn PythonUsageSource,
302 file: &ProjectFile,
303 source: &str,
304 node: Node<'_>,
305 target_self_file: bool,
306) -> Option<Vec<CodeUnit>> {
307 if !is_annotation_reference_node(node, source) {
308 return None;
309 }
310
311 let mut candidates = match node.kind() {
312 "identifier" => resolve_bare_annotation_candidates(
313 graph,
314 python,
315 file,
316 source,
317 node,
318 node_text(node, source),
319 target_self_file,
320 ),
321 "string_content" => {
322 let Some(string) = node.parent() else {
323 return Some(Vec::new());
324 };
325 let Some(ranges) = python_deferred_annotation_identifier_ranges(string, source, None)
326 else {
327 return Some(Vec::new());
328 };
329 let mut candidates = Vec::new();
330 for range in ranges {
331 let Some(symbol) = source.get(range.start_byte..range.end_byte) else {
332 continue;
333 };
334 candidates.extend(resolve_bare_annotation_candidates(
335 graph,
336 python,
337 file,
338 source,
339 node,
340 symbol,
341 target_self_file,
342 ));
343 }
344 candidates
345 }
346 "attribute" => resolve_annotation_attribute_types(graph, python, file, source, node),
347 _ => Vec::new(),
348 };
349 candidates.sort();
350 candidates.dedup();
351 Some(candidates)
352}
353
354#[allow(clippy::too_many_arguments)]
358pub fn annotation_reference_candidates_at_focus(
359 graph: &PythonGraphSource<'_>,
360 python: &dyn PythonUsageSource,
361 file: &ProjectFile,
362 source: &str,
363 node: Node<'_>,
364 focus_start: usize,
365 focus_end: usize,
366 target_self_file: bool,
367) -> Option<Vec<CodeUnit>> {
368 if !is_annotation_reference_node(node, source) {
369 return None;
370 }
371
372 let deferred_tree = if node.kind() == "string_content" {
373 Some(python_deferred_annotation_tree(
374 node.parent()?,
375 source,
376 None,
377 )?)
378 } else {
379 None
380 };
381 let search_root = deferred_tree.as_ref().map_or(node, |tree| tree.root_node());
382 let focused = search_root.descendant_for_byte_range(focus_start, focus_end)?;
383 if focused.kind() != "identifier"
384 || focused.start_byte() != focus_start
385 || focused.end_byte() != focus_end
386 {
387 return Some(Vec::new());
388 }
389
390 let mut path = focused;
391 while let Some(parent) = path.parent() {
392 if parent.kind() != "attribute" {
393 break;
394 }
395 path = parent;
396 }
397 if path.kind() == "attribute" {
398 return Some(focused_annotation_attribute_candidates(
399 graph, python, file, source, node, path, focused,
400 ));
401 }
402
403 let symbol = node_text(focused, source);
404 let mut candidates = resolve_bare_annotation_candidates(
405 graph,
406 python,
407 file,
408 source,
409 node,
410 symbol,
411 target_self_file,
412 );
413 candidates.sort();
414 candidates.dedup();
415 Some(candidates)
416}
417
418fn focused_annotation_attribute_candidates(
419 graph: &PythonGraphSource<'_>,
420 python: &dyn PythonUsageSource,
421 file: &ProjectFile,
422 source: &str,
423 scope_node: Node<'_>,
424 path: Node<'_>,
425 focused: Node<'_>,
426) -> Vec<CodeUnit> {
427 let Some((root, attributes)) = attribute_chain(path) else {
428 return Vec::new();
429 };
430 if root.id() == focused.id() {
431 return resolve_bare_annotation_symbol(
432 graph,
433 python,
434 file,
435 source,
436 scope_node,
437 node_text(root, source),
438 );
439 }
440
441 if attributes
442 .last()
443 .is_some_and(|node| node.id() == focused.id())
444 {
445 let candidates = namespace_qualified_declarations(graph, python, file, source, path);
446 if !candidates.is_empty() {
447 return candidates;
448 }
449 }
450
451 let Some(mut owner) = unique_bare_annotation_class(
452 graph,
453 python,
454 file,
455 source,
456 scope_node,
457 node_text(root, source),
458 ) else {
459 return Vec::new();
460 };
461 for attribute in attributes {
462 let candidates = exact_nested_annotation_class(graph, &owner, node_text(attribute, source));
463 if attribute.id() == focused.id() {
464 return candidates;
465 }
466 let Some(next) =
467 unique_nested_annotation_class(graph, &owner, node_text(attribute, source))
468 else {
469 return Vec::new();
470 };
471 owner = next;
472 }
473 Vec::new()
474}
475
476pub fn annotation_class_qualifier_site<'tree>(
479 graph: &PythonGraphSource<'_>,
480 python: &dyn PythonUsageSource,
481 file: &ProjectFile,
482 source: &str,
483 node: Node<'tree>,
484 target: &CodeUnit,
485) -> Option<Node<'tree>> {
486 if !target.is_class() {
487 return None;
488 }
489 annotation_class_qualifier_sites(graph, python, file, source, node)
490 .find_map(|(site, owner)| (&owner == target).then_some(site))
491}
492
493pub fn annotation_class_qualifier_sites<'a, 'tree: 'a>(
496 graph: &'a PythonGraphSource<'_>,
497 python: &'a dyn PythonUsageSource,
498 file: &'a ProjectFile,
499 source: &'a str,
500 node: Node<'tree>,
501) -> impl Iterator<Item = (Node<'tree>, CodeUnit)> + 'a {
502 let resolved = (|| {
503 if node.kind() != "attribute" || !is_annotation_reference_node(node, source) {
504 return None;
505 }
506 let (root, attributes) = attribute_chain(node)?;
507 let owner = unique_bare_annotation_class(
508 graph,
509 python,
510 file,
511 source,
512 root,
513 node_text(root, source),
514 )?;
515 Some((root, owner, attributes))
516 })();
517 let (mut first, mut owner, attributes) = match resolved {
518 Some((root, owner, attributes)) => (Some((root, owner.clone())), Some(owner), attributes),
519 None => (None, None, Vec::new()),
520 };
521 let count = attributes.len().saturating_sub(1);
523 let mut attributes = attributes.into_iter().take(count);
524 std::iter::from_fn(move || {
525 if let Some(first) = first.take() {
526 return Some(first);
527 }
528 let previous = owner.take()?;
529 let attribute = attributes.next()?;
530 let next = unique_nested_annotation_class(graph, &previous, node_text(attribute, source))?;
531 owner = Some(next.clone());
532 Some((attribute, next))
533 })
534}
535
536fn resolve_annotation_attribute_types(
537 graph: &PythonGraphSource<'_>,
538 python: &dyn PythonUsageSource,
539 file: &ProjectFile,
540 source: &str,
541 node: Node<'_>,
542) -> Vec<CodeUnit> {
543 let mut candidates = namespace_qualified_declarations(graph, python, file, source, node);
548 let Some((root, attributes)) = attribute_chain(node) else {
549 return candidates;
550 };
551 let root_text = node_text(root, source);
552 let Some(mut owner) =
553 unique_bare_annotation_class(graph, python, file, source, root, root_text)
554 else {
555 return candidates;
556 };
557
558 for attribute in attributes {
559 let segment = node_text(attribute, source);
560 let Some(next) = unique_nested_annotation_class(graph, &owner, segment) else {
561 return candidates;
562 };
563 owner = next;
564 }
565
566 candidates.push(owner);
567 candidates
568}
569
570fn exact_nested_annotation_class(
571 graph: &PythonGraphSource<'_>,
572 owner: &CodeUnit,
573 segment: &str,
574) -> Vec<CodeUnit> {
575 let mut candidates: Vec<_> = exact_owner_annotation_members(graph, owner, segment)
576 .into_iter()
577 .filter(CodeUnit::is_class)
578 .collect();
579 candidates.sort();
580 candidates.dedup();
581 candidates
582}
583
584fn unique_bare_annotation_class(
585 graph: &PythonGraphSource<'_>,
586 python: &dyn PythonUsageSource,
587 file: &ProjectFile,
588 source: &str,
589 node: Node<'_>,
590 symbol: &str,
591) -> Option<CodeUnit> {
592 let owners: Vec<_> = resolve_bare_annotation_symbol(graph, python, file, source, node, symbol)
593 .into_iter()
594 .filter(CodeUnit::is_class)
595 .collect();
596 let [owner] = owners.as_slice() else {
597 return None;
598 };
599 Some(owner.clone())
600}
601
602fn unique_nested_annotation_class(
603 graph: &PythonGraphSource<'_>,
604 owner: &CodeUnit,
605 segment: &str,
606) -> Option<CodeUnit> {
607 let candidates = exact_nested_annotation_class(graph, owner, segment);
608 let [candidate] = candidates.as_slice() else {
609 return None;
610 };
611 Some(candidate.clone())
612}
613
614fn resolve_bare_annotation_candidates(
615 graph: &PythonGraphSource<'_>,
616 python: &dyn PythonUsageSource,
617 file: &ProjectFile,
618 source: &str,
619 node: Node<'_>,
620 symbol: &str,
621 target_self_file: bool,
622) -> Vec<CodeUnit> {
623 let mut candidates = resolve_bare_annotation_symbol(graph, python, file, source, node, symbol);
624 if candidates.is_empty() {
625 candidates.extend(resolve_receiver_type(
626 graph,
627 python,
628 file,
629 symbol,
630 target_self_file,
631 ));
632 }
633 candidates
634}
635
636fn exact_owner_annotation_members(
637 graph: &PythonGraphSource<'_>,
638 owner: &CodeUnit,
639 segment: &str,
640) -> Vec<CodeUnit> {
641 let mut candidates: Vec<_> = graph
642 .index
643 .declarations(owner.source())
644 .into_iter()
645 .filter(|unit| {
646 unit.identifier() == segment
647 && graph
648 .index
649 .parent_of(unit)
650 .is_some_and(|parent| parent.fq_name() == owner.fq_name())
651 })
652 .collect();
653 candidates.sort();
654 candidates.dedup();
655 candidates
656}
657
658fn annotation_scope_owner_class(
659 graph: &PythonGraphSource<'_>,
660 file: &ProjectFile,
661 source: &str,
662 node: Node<'_>,
663) -> Option<CodeUnit> {
664 if !annotation_expression_is_class_scoped(node) {
665 return None;
666 }
667 let range = Range {
668 start_byte: node.start_byte(),
669 end_byte: node.end_byte(),
670 start_line: 0,
671 end_line: 0,
672 };
673 if let Some(enclosing) = graph.index.enclosing_code_unit(file, &range) {
674 if enclosing.is_class() {
675 return Some(enclosing);
676 }
677 if let Some(owner) = target_owner_code_unit(graph.index, &enclosing) {
678 return Some(owner);
679 }
680 }
681 structural_annotation_owner_class(graph, file, source, node)
682}
683
684fn annotation_expression_is_class_scoped(node: Node<'_>) -> bool {
685 let site_start = node.start_byte();
686 let site_end = node.end_byte();
687 let mut current = node;
688 while let Some(parent) = current.parent() {
689 if matches!(parent.kind(), "function_definition" | "lambda")
690 && parent
691 .child_by_field_name("body")
692 .is_some_and(|body| body.start_byte() <= site_start && site_end <= body.end_byte())
693 {
694 return false;
695 }
696 if parent.kind() == "class_definition" {
697 return true;
698 }
699 current = parent;
700 }
701 false
702}
703
704fn structural_annotation_owner_class(
705 graph: &PythonGraphSource<'_>,
706 file: &ProjectFile,
707 source: &str,
708 node: Node<'_>,
709) -> Option<CodeUnit> {
710 let mut current = node;
711 while let Some(parent) = current.parent() {
712 if parent.kind() == "class_definition" {
713 let name = node_text(parent.child_by_field_name("name")?, source).trim();
714 if name.is_empty() {
715 return None;
716 }
717 let class_range = Range {
718 start_byte: parent.start_byte(),
719 end_byte: parent.end_byte(),
720 start_line: 0,
721 end_line: 0,
722 };
723 let mut matches: Vec<_> = graph
724 .index
725 .declarations(file)
726 .into_iter()
727 .filter(|unit| unit.is_class() && unit.identifier() == name)
728 .filter(|unit| {
729 graph
730 .index
731 .ranges(unit)
732 .into_iter()
733 .any(|range| range.contains(&class_range))
734 })
735 .collect();
736 matches.sort();
737 matches.dedup();
738 let [owner] = matches.as_slice() else {
739 return None;
740 };
741 return Some(owner.clone());
742 }
743 current = parent;
744 }
745 None
746}
747
748fn is_annotation_reference_node(node: Node<'_>, source: &str) -> bool {
749 if !matches!(node.kind(), "identifier" | "attribute" | "string_content") {
750 return false;
751 }
752 if !python_node_is_in_annotation(node) {
753 return false;
754 }
755
756 let mut current = node;
757 while let Some(parent) = current.parent() {
758 if parent.kind() == "call"
759 && parent
760 .child_by_field_name("arguments")
761 .is_some_and(|arguments| {
762 arguments.start_byte() <= node.start_byte()
763 && node.end_byte() <= arguments.end_byte()
764 })
765 {
766 return false;
767 }
768 if parent.kind() == "subscript"
769 && let Some(value) = parent.child_by_field_name("value")
770 && value.kind() == "identifier"
771 {
772 match node_text(value, source) {
773 "Literal" => return false,
777 "Annotated" => {
780 let Some(arguments) = parent.child_by_field_name("subscript") else {
781 return false;
782 };
783 let type_argument = if arguments.kind() == "expression_list" {
784 arguments.named_child(0)
785 } else {
786 Some(arguments)
787 };
788 if !type_argument.is_some_and(|argument| {
789 argument.start_byte() <= node.start_byte()
790 && node.end_byte() <= argument.end_byte()
791 }) {
792 return false;
793 }
794 }
795 _ => {}
796 }
797 }
798 if parent.kind() == "type_parameter"
799 && let Some(generic) = parent
800 .parent()
801 .filter(|parent| parent.kind() == "generic_type")
802 {
803 let generic_name = generic
804 .child_by_field_name("name")
805 .or_else(|| generic.named_child(0))
806 .map(|name| node_text(name, source));
807 match generic_name {
808 Some("Literal") => return false,
809 Some("Annotated") => {
810 let type_argument = parent.named_child(0);
811 if !type_argument.is_some_and(|argument| {
812 argument.start_byte() <= node.start_byte()
813 && node.end_byte() <= argument.end_byte()
814 }) {
815 return false;
816 }
817 }
818 _ => {}
819 }
820 }
821 current = parent;
822 }
823 true
824}
825
826pub fn resolve_visible_named_import_candidates(
834 graph: &PythonGraphSource<'_>,
835 python: &dyn PythonUsageSource,
836 file: &ProjectFile,
837 timeline: &ModuleBindingTimeline,
838 local: &str,
839 cutoff: usize,
840) -> Vec<CodeUnit> {
841 let Some(events) = timeline.get(local) else {
842 return Vec::new();
843 };
844 let mut candidates = Vec::new();
845 for event in visible_module_binding_suffix(events, cutoff) {
846 let ModuleBindingEventKind::FromImport {
847 module,
848 imported_name,
849 } = &event.kind
850 else {
851 continue;
852 };
853 let mut resolved_module = false;
854 for module_file in usage_resolve_module_files(python, file, module) {
855 let Some(module_fqn) = graph
856 .index
857 .declarations(&module_file)
858 .into_iter()
859 .find(CodeUnit::is_module)
860 .map(|unit| unit.fq_name())
861 else {
862 continue;
863 };
864 resolved_module = true;
865 let fqn = format!("{module_fqn}.{imported_name}");
866 candidates.extend(resolve_fqn_candidates(python, Some(file), &fqn, |name| {
867 graph.index.definitions(name).collect()
868 }));
869 }
870 if !resolved_module {
871 let fqn = if module.ends_with('.') {
872 format!("{module}{imported_name}")
873 } else {
874 format!("{module}.{imported_name}")
875 };
876 candidates.extend(resolve_fqn_candidates(python, Some(file), &fqn, |name| {
877 graph.index.definitions(name).collect()
878 }));
879 }
880 }
881 candidates.sort();
882 candidates.dedup();
883 candidates
884}
885
886pub fn resolve_constructor_types(
891 graph: &PythonGraphSource<'_>,
892 python: &dyn PythonUsageSource,
893 file: &ProjectFile,
894 source: &str,
895 function: Node<'_>,
896) -> Vec<CodeUnit> {
897 let candidates = match function.kind() {
898 "identifier" => {
899 let local = node_text(function, source);
900 if local.is_empty() {
901 return Vec::new();
902 }
903 let binder = python.import_binder_of(file);
904 let fqn = match binder.bindings.get(local) {
905 Some(binding) if binding.kind == ImportKind::Named => binding
906 .imported_name
907 .as_ref()
908 .map(|imported| format!("{}.{}", binding.module_specifier, imported)),
909 _ => graph
910 .index
911 .declarations(file)
912 .into_iter()
913 .find(|unit| unit.is_class() && unit.identifier() == local)
914 .map(|unit| unit.fq_name()),
915 };
916 let Some(fqn) = fqn else {
917 return Vec::new();
918 };
919 resolve_fqn_candidates(python, Some(file), &fqn, |name| {
920 graph.index.definitions(name).collect()
921 })
922 }
923 "attribute" => namespace_qualified_declarations(graph, python, file, source, function),
924 _ => Vec::new(),
925 };
926 let mut classes: Vec<CodeUnit> = candidates.into_iter().filter(CodeUnit::is_class).collect();
929 classes.sort();
930 classes.dedup();
931 classes
932}
933
934pub fn resolve_callable_parameter_default_types(
940 graph: &PythonGraphSource<'_>,
941 python: &dyn PythonUsageSource,
942 file: &ProjectFile,
943 source: &str,
944 reference: Node<'_>,
945 local_name: &str,
946) -> Vec<CodeUnit> {
947 let site_start = reference.start_byte();
948 let site_end = reference.end_byte();
949 let mut current = reference;
950 while let Some(parent) = current.parent() {
951 current = parent;
952 if !matches!(current.kind(), "function_definition" | "lambda") {
953 continue;
954 }
955 if current
956 .child_by_field_name("body")
957 .is_none_or(|body| !(body.start_byte() <= site_start && site_end <= body.end_byte()))
958 {
959 continue;
960 }
961 let Some(parameters) = current.child_by_field_name("parameters") else {
962 return Vec::new();
963 };
964 let mut cursor = parameters.walk();
965 for parameter in parameters.named_children(&mut cursor) {
966 let name = if parameter.kind() == "identifier" {
967 Some(parameter)
968 } else {
969 parameter.child_by_field_name("name")
970 };
971 if name.is_none_or(|name| node_text(name, source) != local_name) {
972 continue;
973 }
974 let Some(value) = parameter.child_by_field_name("value") else {
975 return Vec::new();
976 };
977 return resolve_constructor_types(graph, python, file, source, value);
978 }
979 return Vec::new();
980 }
981 Vec::new()
982}
983
984fn namespace_qualified_declarations(
991 graph: &PythonGraphSource<'_>,
992 python: &dyn PythonUsageSource,
993 file: &ProjectFile,
994 source: &str,
995 node: Node<'_>,
996) -> Vec<CodeUnit> {
997 let binder = python.import_binder_of(file);
998 let Some(fqn) = namespace_constructor_fqn(&binder, source, node) else {
999 return Vec::new();
1000 };
1001 let mut candidates = resolve_fqn_candidates(python, Some(file), &fqn, |name| {
1002 graph.index.definitions(name).collect()
1003 });
1004 candidates.sort();
1005 candidates.dedup();
1006 candidates
1007}
1008
1009fn namespace_constructor_fqn(
1010 binder: &ImportBinder,
1011 source: &str,
1012 function: Node<'_>,
1013) -> Option<String> {
1014 let mut attributes = Vec::new();
1015 let mut current = function;
1016 while current.kind() == "attribute" {
1017 let attribute = current.child_by_field_name("attribute")?;
1018 let text = node_text(attribute, source);
1019 if text.is_empty() {
1020 return None;
1021 }
1022 attributes.push(text);
1023 current = current.child_by_field_name("object")?;
1024 }
1025 if current.kind() != "identifier" {
1026 return None;
1027 }
1028 let root = node_text(current, source);
1029 let binding = binder.bindings.get(root)?;
1030 if binding.kind != ImportKind::Namespace {
1031 return None;
1032 }
1033 let mut fqn = binding.module_specifier.clone();
1034 for attribute in attributes.into_iter().rev() {
1035 fqn.push('.');
1036 fqn.push_str(attribute);
1037 }
1038 Some(fqn)
1039}
1040
1041fn resolve_indexed_receiver_type(
1042 graph: &PythonGraphSource<'_>,
1043 file: &ProjectFile,
1044 raw_type: &str,
1045) -> Option<CodeUnit> {
1046 let package_matches = module_fq_for_file(graph.index, file).map_or_else(Vec::new, |module| {
1047 relational_definitions(
1048 graph,
1049 RelationalName::stable(module),
1050 RelationalDefinitionQuery::PackageTypes {
1051 simple_name: raw_type.to_string(),
1052 },
1053 )
1054 });
1055 let mut root_name = FqName::new();
1056 root_name.push(segment_interner().intern(raw_type, SegmentKind::Type));
1057 let exact_matches = relational_definitions(
1058 graph,
1059 RelationalName::stable(root_name.clone()),
1060 RelationalDefinitionQuery::ExactName,
1061 );
1062 let normalized_matches = relational_definitions(
1063 graph,
1064 RelationalName::stable(root_name),
1065 RelationalDefinitionQuery::NormalizedName,
1066 );
1067 package_matches
1068 .into_iter()
1069 .chain(exact_matches)
1070 .chain(normalized_matches)
1071 .find(|code_unit| code_unit.identifier() == raw_type && code_unit.is_class())
1072}
1073
1074fn relational_definitions(
1075 graph: &PythonGraphSource<'_>,
1076 name: RelationalName,
1077 query: RelationalDefinitionQuery,
1078) -> Vec<CodeUnit> {
1079 let question = RelationalDefinitionQuestion {
1080 language_scope: DefinitionLanguageScope::Language(Language::Python),
1086 name,
1087 query,
1088 };
1089 match graph.definitions.ask(&question) {
1090 RelationalDefinitionValue::Definitions(units) => units,
1091 _ => panic!("definition question returned the wrong result shape"),
1092 }
1093}
1094
1095fn module_fq_for_file(index: &dyn CodeUnitIndex, file: &ProjectFile) -> Option<FqName> {
1096 index
1097 .declarations(file)
1098 .into_iter()
1099 .find(|code_unit| code_unit.is_module())
1100 .map(|code_unit| code_unit.fq().clone())
1101 .or_else(|| {
1102 index
1103 .declarations(file)
1104 .into_iter()
1105 .find(|code_unit| !code_unit.package_name().is_empty())
1106 .map(|code_unit| code_unit.package_fq())
1107 })
1108}
1109
1110pub fn normalized_receiver_type(annotation: &str) -> Option<String> {
1111 let annotation = unwrap_python_string_annotation(annotation.trim());
1112 let annotation = unwrap_supported_receiver_wrapper(annotation);
1113 if annotation.is_empty()
1114 || annotation.contains('|')
1115 || annotation.contains('[')
1116 || annotation.contains(']')
1117 || annotation.contains(',')
1118 || annotation.contains('(')
1119 || annotation.contains(')')
1120 || annotation.contains('{')
1121 || annotation.contains('}')
1122 || annotation.contains(':')
1123 {
1124 return None;
1125 }
1126 Some(annotation.to_string())
1127}
1128
1129fn unwrap_python_string_annotation(annotation: &str) -> &str {
1130 if annotation.len() >= 2 {
1131 let bytes = annotation.as_bytes();
1132 let first = bytes[0];
1133 let last = bytes[annotation.len() - 1];
1134 if (first == b'\'' || first == b'"') && first == last {
1135 return annotation[1..annotation.len() - 1].trim();
1136 }
1137 }
1138 annotation
1139}
1140
1141fn unwrap_supported_receiver_wrapper(annotation: &str) -> &str {
1142 let mut current = annotation.trim();
1143 loop {
1144 let next = current
1145 .strip_prefix("Optional[")
1146 .or_else(|| current.strip_prefix("typing.Optional["))
1147 .and_then(|inner| inner.strip_suffix(']'))
1148 .map(str::trim);
1149 let Some(unwrapped) = next else {
1150 return current;
1151 };
1152 current = unwrapped;
1153 }
1154}
1155
1156pub fn receiver_annotation_matches_target(
1157 annotation: &str,
1158 edges: &[ImportEdge],
1159 target_short: &str,
1160 target_self_file: bool,
1161) -> bool {
1162 let annotation = annotation.trim();
1163 if annotation.is_empty() {
1164 return false;
1165 }
1166 if annotation.contains('|')
1167 || annotation.contains('[')
1168 || annotation.contains(']')
1169 || annotation.contains(',')
1170 || annotation.contains('(')
1171 || annotation.contains(')')
1172 {
1173 return false;
1174 }
1175 if annotation == target_short {
1176 return target_self_file || edges.iter().any(|edge| edge.local_name == target_short);
1177 }
1178
1179 let segments = parse_symbol_path(Language::Python, annotation);
1185 let Some((member, qualifier_parts)) = segments.split_last() else {
1186 return false;
1187 };
1188 if qualifier_parts.is_empty() {
1189 return false;
1190 }
1191 let qualifier = qualifier_parts.join(".");
1192 let member = member.as_str();
1193 if member != target_short {
1194 return false;
1195 }
1196 edges.iter().any(|edge| {
1197 matches!(edge.kind, ImportEdgeKind::Namespace)
1198 && (edge.local_name == qualifier
1199 || qualifier.ends_with(&format!(".{}", edge.local_name)))
1200 })
1201}
1202
1203