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