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 node.kind() != "attribute"
487 || !target.is_class()
488 || !is_annotation_reference_node(node, source)
489 {
490 return None;
491 }
492
493 let (root, attributes) = attribute_chain(node)?;
494 let mut owner =
495 unique_bare_annotation_class(graph, python, file, source, root, node_text(root, source))?;
496 if &owner == target {
497 return Some(root);
498 }
499
500 let qualifier_count = attributes.len().saturating_sub(1);
503 for attribute in attributes.into_iter().take(qualifier_count) {
504 owner = unique_nested_annotation_class(graph, &owner, node_text(attribute, source))?;
505 if &owner == target {
506 return Some(attribute);
507 }
508 }
509
510 None
511}
512
513fn resolve_annotation_attribute_types(
514 graph: &PythonGraphSource<'_>,
515 python: &dyn PythonUsageSource,
516 file: &ProjectFile,
517 source: &str,
518 node: Node<'_>,
519) -> Vec<CodeUnit> {
520 let mut candidates = namespace_qualified_declarations(graph, python, file, source, node);
525 let Some((root, attributes)) = attribute_chain(node) else {
526 return candidates;
527 };
528 let root_text = node_text(root, source);
529 let Some(mut owner) =
530 unique_bare_annotation_class(graph, python, file, source, root, root_text)
531 else {
532 return candidates;
533 };
534
535 for attribute in attributes {
536 let segment = node_text(attribute, source);
537 let Some(next) = unique_nested_annotation_class(graph, &owner, segment) else {
538 return candidates;
539 };
540 owner = next;
541 }
542
543 candidates.push(owner);
544 candidates
545}
546
547fn exact_nested_annotation_class(
548 graph: &PythonGraphSource<'_>,
549 owner: &CodeUnit,
550 segment: &str,
551) -> Vec<CodeUnit> {
552 let mut candidates: Vec<_> = exact_owner_annotation_members(graph, owner, segment)
553 .into_iter()
554 .filter(CodeUnit::is_class)
555 .collect();
556 candidates.sort();
557 candidates.dedup();
558 candidates
559}
560
561fn unique_bare_annotation_class(
562 graph: &PythonGraphSource<'_>,
563 python: &dyn PythonUsageSource,
564 file: &ProjectFile,
565 source: &str,
566 node: Node<'_>,
567 symbol: &str,
568) -> Option<CodeUnit> {
569 let owners: Vec<_> = resolve_bare_annotation_symbol(graph, python, file, source, node, symbol)
570 .into_iter()
571 .filter(CodeUnit::is_class)
572 .collect();
573 let [owner] = owners.as_slice() else {
574 return None;
575 };
576 Some(owner.clone())
577}
578
579fn unique_nested_annotation_class(
580 graph: &PythonGraphSource<'_>,
581 owner: &CodeUnit,
582 segment: &str,
583) -> Option<CodeUnit> {
584 let candidates = exact_nested_annotation_class(graph, owner, segment);
585 let [candidate] = candidates.as_slice() else {
586 return None;
587 };
588 Some(candidate.clone())
589}
590
591fn resolve_bare_annotation_candidates(
592 graph: &PythonGraphSource<'_>,
593 python: &dyn PythonUsageSource,
594 file: &ProjectFile,
595 source: &str,
596 node: Node<'_>,
597 symbol: &str,
598 target_self_file: bool,
599) -> Vec<CodeUnit> {
600 let mut candidates = resolve_bare_annotation_symbol(graph, python, file, source, node, symbol);
601 if candidates.is_empty() {
602 candidates.extend(resolve_receiver_type(
603 graph,
604 python,
605 file,
606 symbol,
607 target_self_file,
608 ));
609 }
610 candidates
611}
612
613fn exact_owner_annotation_members(
614 graph: &PythonGraphSource<'_>,
615 owner: &CodeUnit,
616 segment: &str,
617) -> Vec<CodeUnit> {
618 let mut candidates: Vec<_> = graph
619 .index
620 .declarations(owner.source())
621 .into_iter()
622 .filter(|unit| {
623 unit.identifier() == segment
624 && graph
625 .index
626 .parent_of(unit)
627 .is_some_and(|parent| parent.fq_name() == owner.fq_name())
628 })
629 .collect();
630 candidates.sort();
631 candidates.dedup();
632 candidates
633}
634
635fn annotation_scope_owner_class(
636 graph: &PythonGraphSource<'_>,
637 file: &ProjectFile,
638 source: &str,
639 node: Node<'_>,
640) -> Option<CodeUnit> {
641 if !annotation_expression_is_class_scoped(node) {
642 return None;
643 }
644 let range = Range {
645 start_byte: node.start_byte(),
646 end_byte: node.end_byte(),
647 start_line: 0,
648 end_line: 0,
649 };
650 if let Some(enclosing) = graph.index.enclosing_code_unit(file, &range) {
651 if enclosing.is_class() {
652 return Some(enclosing);
653 }
654 if let Some(owner) = target_owner_code_unit(graph.index, &enclosing) {
655 return Some(owner);
656 }
657 }
658 structural_annotation_owner_class(graph, file, source, node)
659}
660
661fn annotation_expression_is_class_scoped(node: Node<'_>) -> bool {
662 let site_start = node.start_byte();
663 let site_end = node.end_byte();
664 let mut current = node;
665 while let Some(parent) = current.parent() {
666 if matches!(parent.kind(), "function_definition" | "lambda")
667 && parent
668 .child_by_field_name("body")
669 .is_some_and(|body| body.start_byte() <= site_start && site_end <= body.end_byte())
670 {
671 return false;
672 }
673 if parent.kind() == "class_definition" {
674 return true;
675 }
676 current = parent;
677 }
678 false
679}
680
681fn structural_annotation_owner_class(
682 graph: &PythonGraphSource<'_>,
683 file: &ProjectFile,
684 source: &str,
685 node: Node<'_>,
686) -> Option<CodeUnit> {
687 let mut current = node;
688 while let Some(parent) = current.parent() {
689 if parent.kind() == "class_definition" {
690 let name = node_text(parent.child_by_field_name("name")?, source).trim();
691 if name.is_empty() {
692 return None;
693 }
694 let class_range = Range {
695 start_byte: parent.start_byte(),
696 end_byte: parent.end_byte(),
697 start_line: 0,
698 end_line: 0,
699 };
700 let mut matches: Vec<_> = graph
701 .index
702 .declarations(file)
703 .into_iter()
704 .filter(|unit| unit.is_class() && unit.identifier() == name)
705 .filter(|unit| {
706 graph
707 .index
708 .ranges(unit)
709 .into_iter()
710 .any(|range| range.contains(&class_range))
711 })
712 .collect();
713 matches.sort();
714 matches.dedup();
715 let [owner] = matches.as_slice() else {
716 return None;
717 };
718 return Some(owner.clone());
719 }
720 current = parent;
721 }
722 None
723}
724
725fn is_annotation_reference_node(node: Node<'_>, source: &str) -> bool {
726 if !matches!(node.kind(), "identifier" | "attribute" | "string_content") {
727 return false;
728 }
729 if !python_node_is_in_annotation(node) {
730 return false;
731 }
732
733 let mut current = node;
734 while let Some(parent) = current.parent() {
735 if parent.kind() == "call"
736 && parent
737 .child_by_field_name("arguments")
738 .is_some_and(|arguments| {
739 arguments.start_byte() <= node.start_byte()
740 && node.end_byte() <= arguments.end_byte()
741 })
742 {
743 return false;
744 }
745 if parent.kind() == "subscript"
746 && let Some(value) = parent.child_by_field_name("value")
747 && value.kind() == "identifier"
748 {
749 match node_text(value, source) {
750 "Literal" => return false,
754 "Annotated" => {
757 let Some(arguments) = parent.child_by_field_name("subscript") else {
758 return false;
759 };
760 let type_argument = if arguments.kind() == "expression_list" {
761 arguments.named_child(0)
762 } else {
763 Some(arguments)
764 };
765 if !type_argument.is_some_and(|argument| {
766 argument.start_byte() <= node.start_byte()
767 && node.end_byte() <= argument.end_byte()
768 }) {
769 return false;
770 }
771 }
772 _ => {}
773 }
774 }
775 if parent.kind() == "type_parameter"
776 && let Some(generic) = parent
777 .parent()
778 .filter(|parent| parent.kind() == "generic_type")
779 {
780 let generic_name = generic
781 .child_by_field_name("name")
782 .or_else(|| generic.named_child(0))
783 .map(|name| node_text(name, source));
784 match generic_name {
785 Some("Literal") => return false,
786 Some("Annotated") => {
787 let type_argument = parent.named_child(0);
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 current = parent;
799 }
800 true
801}
802
803pub fn resolve_visible_named_import_candidates(
811 graph: &PythonGraphSource<'_>,
812 python: &dyn PythonUsageSource,
813 file: &ProjectFile,
814 timeline: &ModuleBindingTimeline,
815 local: &str,
816 cutoff: usize,
817) -> Vec<CodeUnit> {
818 let Some(events) = timeline.get(local) else {
819 return Vec::new();
820 };
821 let mut candidates = Vec::new();
822 for event in visible_module_binding_suffix(events, cutoff) {
823 let ModuleBindingEventKind::FromImport {
824 module,
825 imported_name,
826 } = &event.kind
827 else {
828 continue;
829 };
830 let mut resolved_module = false;
831 for module_file in usage_resolve_module_files(python, file, module) {
832 let Some(module_fqn) = graph
833 .index
834 .declarations(&module_file)
835 .into_iter()
836 .find(CodeUnit::is_module)
837 .map(|unit| unit.fq_name())
838 else {
839 continue;
840 };
841 resolved_module = true;
842 let fqn = format!("{module_fqn}.{imported_name}");
843 candidates.extend(resolve_fqn_candidates(python, Some(file), &fqn, |name| {
844 graph.index.definitions(name).collect()
845 }));
846 }
847 if !resolved_module {
848 let fqn = if module.ends_with('.') {
849 format!("{module}{imported_name}")
850 } else {
851 format!("{module}.{imported_name}")
852 };
853 candidates.extend(resolve_fqn_candidates(python, Some(file), &fqn, |name| {
854 graph.index.definitions(name).collect()
855 }));
856 }
857 }
858 candidates.sort();
859 candidates.dedup();
860 candidates
861}
862
863pub fn resolve_constructor_types(
868 graph: &PythonGraphSource<'_>,
869 python: &dyn PythonUsageSource,
870 file: &ProjectFile,
871 source: &str,
872 function: Node<'_>,
873) -> Vec<CodeUnit> {
874 let candidates = match function.kind() {
875 "identifier" => {
876 let local = node_text(function, source);
877 if local.is_empty() {
878 return Vec::new();
879 }
880 let binder = python.import_binder_of(file);
881 let fqn = match binder.bindings.get(local) {
882 Some(binding) if binding.kind == ImportKind::Named => binding
883 .imported_name
884 .as_ref()
885 .map(|imported| format!("{}.{}", binding.module_specifier, imported)),
886 _ => graph
887 .index
888 .declarations(file)
889 .into_iter()
890 .find(|unit| unit.is_class() && unit.identifier() == local)
891 .map(|unit| unit.fq_name()),
892 };
893 let Some(fqn) = fqn else {
894 return Vec::new();
895 };
896 resolve_fqn_candidates(python, Some(file), &fqn, |name| {
897 graph.index.definitions(name).collect()
898 })
899 }
900 "attribute" => namespace_qualified_declarations(graph, python, file, source, function),
901 _ => Vec::new(),
902 };
903 let mut classes: Vec<CodeUnit> = candidates.into_iter().filter(CodeUnit::is_class).collect();
906 classes.sort();
907 classes.dedup();
908 classes
909}
910
911pub fn resolve_callable_parameter_default_types(
917 graph: &PythonGraphSource<'_>,
918 python: &dyn PythonUsageSource,
919 file: &ProjectFile,
920 source: &str,
921 reference: Node<'_>,
922 local_name: &str,
923) -> Vec<CodeUnit> {
924 let site_start = reference.start_byte();
925 let site_end = reference.end_byte();
926 let mut current = reference;
927 while let Some(parent) = current.parent() {
928 current = parent;
929 if !matches!(current.kind(), "function_definition" | "lambda") {
930 continue;
931 }
932 if current
933 .child_by_field_name("body")
934 .is_none_or(|body| !(body.start_byte() <= site_start && site_end <= body.end_byte()))
935 {
936 continue;
937 }
938 let Some(parameters) = current.child_by_field_name("parameters") else {
939 return Vec::new();
940 };
941 let mut cursor = parameters.walk();
942 for parameter in parameters.named_children(&mut cursor) {
943 let name = if parameter.kind() == "identifier" {
944 Some(parameter)
945 } else {
946 parameter.child_by_field_name("name")
947 };
948 if name.is_none_or(|name| node_text(name, source) != local_name) {
949 continue;
950 }
951 let Some(value) = parameter.child_by_field_name("value") else {
952 return Vec::new();
953 };
954 return resolve_constructor_types(graph, python, file, source, value);
955 }
956 return Vec::new();
957 }
958 Vec::new()
959}
960
961fn namespace_qualified_declarations(
968 graph: &PythonGraphSource<'_>,
969 python: &dyn PythonUsageSource,
970 file: &ProjectFile,
971 source: &str,
972 node: Node<'_>,
973) -> Vec<CodeUnit> {
974 let binder = python.import_binder_of(file);
975 let Some(fqn) = namespace_constructor_fqn(&binder, source, node) else {
976 return Vec::new();
977 };
978 let mut candidates = resolve_fqn_candidates(python, Some(file), &fqn, |name| {
979 graph.index.definitions(name).collect()
980 });
981 candidates.sort();
982 candidates.dedup();
983 candidates
984}
985
986fn namespace_constructor_fqn(
987 binder: &ImportBinder,
988 source: &str,
989 function: Node<'_>,
990) -> Option<String> {
991 let mut attributes = Vec::new();
992 let mut current = function;
993 while current.kind() == "attribute" {
994 let attribute = current.child_by_field_name("attribute")?;
995 let text = node_text(attribute, source);
996 if text.is_empty() {
997 return None;
998 }
999 attributes.push(text);
1000 current = current.child_by_field_name("object")?;
1001 }
1002 if current.kind() != "identifier" {
1003 return None;
1004 }
1005 let root = node_text(current, source);
1006 let binding = binder.bindings.get(root)?;
1007 if binding.kind != ImportKind::Namespace {
1008 return None;
1009 }
1010 let mut fqn = binding.module_specifier.clone();
1011 for attribute in attributes.into_iter().rev() {
1012 fqn.push('.');
1013 fqn.push_str(attribute);
1014 }
1015 Some(fqn)
1016}
1017
1018fn resolve_indexed_receiver_type(
1019 graph: &PythonGraphSource<'_>,
1020 file: &ProjectFile,
1021 raw_type: &str,
1022) -> Option<CodeUnit> {
1023 let package_matches = module_fq_for_file(graph.index, file).map_or_else(Vec::new, |module| {
1024 relational_definitions(
1025 graph,
1026 RelationalName::stable(module),
1027 RelationalDefinitionQuery::PackageTypes {
1028 simple_name: raw_type.to_string(),
1029 },
1030 )
1031 });
1032 let mut root_name = FqName::new();
1033 root_name.push(segment_interner().intern(raw_type, SegmentKind::Type));
1034 let exact_matches = relational_definitions(
1035 graph,
1036 RelationalName::stable(root_name.clone()),
1037 RelationalDefinitionQuery::ExactName,
1038 );
1039 let normalized_matches = relational_definitions(
1040 graph,
1041 RelationalName::stable(root_name),
1042 RelationalDefinitionQuery::NormalizedName,
1043 );
1044 package_matches
1045 .into_iter()
1046 .chain(exact_matches)
1047 .chain(normalized_matches)
1048 .find(|code_unit| code_unit.identifier() == raw_type && code_unit.is_class())
1049}
1050
1051fn relational_definitions(
1052 graph: &PythonGraphSource<'_>,
1053 name: RelationalName,
1054 query: RelationalDefinitionQuery,
1055) -> Vec<CodeUnit> {
1056 let question = RelationalDefinitionQuestion {
1057 language_scope: DefinitionLanguageScope::Language(Language::Python),
1063 name,
1064 query,
1065 };
1066 match graph.definitions.ask(&question) {
1067 RelationalDefinitionValue::Definitions(units) => units,
1068 _ => panic!("definition question returned the wrong result shape"),
1069 }
1070}
1071
1072fn module_fq_for_file(index: &dyn CodeUnitIndex, file: &ProjectFile) -> Option<FqName> {
1073 index
1074 .declarations(file)
1075 .into_iter()
1076 .find(|code_unit| code_unit.is_module())
1077 .map(|code_unit| code_unit.fq().clone())
1078 .or_else(|| {
1079 index
1080 .declarations(file)
1081 .into_iter()
1082 .find(|code_unit| !code_unit.package_name().is_empty())
1083 .map(|code_unit| code_unit.package_fq())
1084 })
1085}
1086
1087pub fn normalized_receiver_type(annotation: &str) -> Option<String> {
1088 let annotation = unwrap_python_string_annotation(annotation.trim());
1089 let annotation = unwrap_supported_receiver_wrapper(annotation);
1090 if annotation.is_empty()
1091 || annotation.contains('|')
1092 || annotation.contains('[')
1093 || annotation.contains(']')
1094 || annotation.contains(',')
1095 || annotation.contains('(')
1096 || annotation.contains(')')
1097 || annotation.contains('{')
1098 || annotation.contains('}')
1099 || annotation.contains(':')
1100 {
1101 return None;
1102 }
1103 Some(annotation.to_string())
1104}
1105
1106fn unwrap_python_string_annotation(annotation: &str) -> &str {
1107 if annotation.len() >= 2 {
1108 let bytes = annotation.as_bytes();
1109 let first = bytes[0];
1110 let last = bytes[annotation.len() - 1];
1111 if (first == b'\'' || first == b'"') && first == last {
1112 return annotation[1..annotation.len() - 1].trim();
1113 }
1114 }
1115 annotation
1116}
1117
1118fn unwrap_supported_receiver_wrapper(annotation: &str) -> &str {
1119 let mut current = annotation.trim();
1120 loop {
1121 let next = current
1122 .strip_prefix("Optional[")
1123 .or_else(|| current.strip_prefix("typing.Optional["))
1124 .and_then(|inner| inner.strip_suffix(']'))
1125 .map(str::trim);
1126 let Some(unwrapped) = next else {
1127 return current;
1128 };
1129 current = unwrapped;
1130 }
1131}
1132
1133pub fn receiver_annotation_matches_target(
1134 annotation: &str,
1135 edges: &[ImportEdge],
1136 target_short: &str,
1137 target_self_file: bool,
1138) -> bool {
1139 let annotation = annotation.trim();
1140 if annotation.is_empty() {
1141 return false;
1142 }
1143 if annotation.contains('|')
1144 || annotation.contains('[')
1145 || annotation.contains(']')
1146 || annotation.contains(',')
1147 || annotation.contains('(')
1148 || annotation.contains(')')
1149 {
1150 return false;
1151 }
1152 if annotation == target_short {
1153 return target_self_file || edges.iter().any(|edge| edge.local_name == target_short);
1154 }
1155
1156 let segments = parse_symbol_path(Language::Python, annotation);
1162 let Some((member, qualifier_parts)) = segments.split_last() else {
1163 return false;
1164 };
1165 if qualifier_parts.is_empty() {
1166 return false;
1167 }
1168 let qualifier = qualifier_parts.join(".");
1169 let member = member.as_str();
1170 if member != target_short {
1171 return false;
1172 }
1173 edges.iter().any(|edge| {
1174 matches!(edge.kind, ImportEdgeKind::Namespace)
1175 && (edge.local_name == qualifier
1176 || qualifier.ends_with(&format!(".{}", edge.local_name)))
1177 })
1178}
1179
1180