1use crate::graph::PythonGraphSource;
9use crate::graph_support::PythonUsageSource;
10use crate::imports::resolve_fqn_candidates;
11use crate::syntax::{
12 python_deferred_annotation_identifier_ranges, python_deferred_annotation_tree,
13 python_node_is_in_annotation,
14};
15use crate::usage_index::usage_seeds;
16use brokk_bifrost_core::analyzer::symbol_path::parse_symbol_path;
17use brokk_bifrost_core::analyzer::usages::model::{ExportEntry, ImportBinder, ImportKind};
18use brokk_bifrost_core::analyzer::usages::{ImportEdge, ImportEdgeKind};
19use brokk_bifrost_core::analyzer::{
20 BoundedDefinitionLookup, CodeUnit, CodeUnitIndex, Language, ProjectFile, Range,
21};
22use std::collections::BTreeSet;
23use tree_sitter::Node;
24
25pub fn infer_export_names(python: &dyn PythonUsageSource, target: &CodeUnit) -> BTreeSet<String> {
26 if target_owner_code_unit(python, target).is_some() {
27 let owner_name = top_level_identifier(python, target);
28 let owner_exports =
29 infer_export_names_for_local(python, target, target.source(), &owner_name);
30 if !owner_exports.is_empty() {
31 return owner_exports;
32 }
33 }
34
35 infer_export_names_for_local(python, target, target.source(), target.identifier())
36}
37
38pub fn infer_usage_seeds(
39 python: &dyn PythonUsageSource,
40 target: &CodeUnit,
41 seed_names: BTreeSet<String>,
42) -> BTreeSet<(ProjectFile, String)> {
43 let mut seeds = BTreeSet::new();
44 for seed_name in &seed_names {
45 seeds.extend(usage_seeds(python, target.source(), seed_name));
46 }
47 if seeds.is_empty()
48 && seed_names.contains(target.identifier())
49 && is_module_level_target_identifier(python, target, target.source(), target.identifier())
50 {
51 seeds.insert((target.source().clone(), target.identifier().to_string()));
52 }
53 seeds
54}
55
56fn infer_export_names_for_local(
57 python: &dyn PythonUsageSource,
58 target: &CodeUnit,
59 file: &ProjectFile,
60 local_name: &str,
61) -> BTreeSet<String> {
62 let index = python.export_index_of(file);
63 let mut export_names = BTreeSet::new();
64 if index.exports_by_name.contains_key(local_name) {
65 export_names.insert(local_name.to_string());
66 }
67 for (export_name, entry) in &index.exports_by_name {
68 if matches!(entry, ExportEntry::Local { local_name: name } if name == local_name) {
69 export_names.insert(export_name.clone());
70 }
71 }
72 if export_names.is_empty()
73 && is_module_level_target_identifier(python, target, file, local_name)
74 {
75 export_names.insert(local_name.to_string());
76 }
77 export_names
78}
79
80fn is_module_level_target_identifier(
81 python: &dyn PythonUsageSource,
82 target: &CodeUnit,
83 file: &ProjectFile,
84 local_name: &str,
85) -> bool {
86 target.source() == file
87 && target.identifier() == local_name
88 && python
89 .parent_of(target)
90 .is_some_and(|parent| parent.is_module() && parent.source() == file)
91}
92
93pub fn top_level_identifier(index: &dyn CodeUnitIndex, target: &CodeUnit) -> String {
94 let mut current = target.clone();
95 while let Some(parent) = index.parent_of(¤t) {
96 if parent.is_module() {
97 break;
98 }
99 current = parent;
100 }
101 current.identifier().to_string()
102}
103
104pub fn member_name(index: &dyn CodeUnitIndex, target: &CodeUnit) -> Option<String> {
105 target_owner_code_unit(index, target).map(|_| target.identifier().to_string())
106}
107
108pub fn target_owner_code_unit(index: &dyn CodeUnitIndex, target: &CodeUnit) -> Option<CodeUnit> {
109 index
110 .parent_of(target)
111 .filter(|parent| parent.source() == target.source() && parent.is_class())
112}
113
114pub fn resolve_receiver_type(
115 graph: &PythonGraphSource<'_>,
116 python: &dyn PythonUsageSource,
117 file: &ProjectFile,
118 raw_type: &str,
119 target_self_file: bool,
120) -> Option<CodeUnit> {
121 let raw_type = raw_type.trim();
122 if raw_type.is_empty() || raw_type.contains('.') || raw_type.contains('|') {
123 return None;
124 }
125
126 if let Some(binding) = python.import_binder_of(file).bindings.get(raw_type)
127 && binding.kind == ImportKind::Named
128 && let Some(imported) = binding.imported_name.as_ref()
129 {
130 let fqn = format!("{}.{}", binding.module_specifier, imported);
131 if let Some(class) =
132 resolve_fqn_candidates(python, &fqn, |name| graph.index.definitions(name).collect())
133 .into_iter()
134 .find(CodeUnit::is_class)
135 {
136 return Some(class);
137 }
138 }
139
140 if let Some(provider) = graph.imports
141 && let Some(imported) = provider
142 .imported_code_units_of(file)
143 .iter()
144 .find(|code_unit| code_unit.identifier() == raw_type && code_unit.is_class())
145 {
146 return Some(imported.clone());
147 }
148
149 graph
150 .index
151 .declarations(file)
152 .into_iter()
153 .find(|code_unit| code_unit.identifier() == raw_type && code_unit.is_class())
154 .or_else(|| {
155 if !target_self_file {
156 return None;
157 }
158 let mut resolved = None;
163 (graph.definitions)(&mut |support| {
164 resolved = resolve_indexed_receiver_type(graph.index, support, file, raw_type);
165 });
166 resolved
167 })
168}
169
170fn resolve_bare_annotation_symbol(
171 graph: &PythonGraphSource<'_>,
172 python: &dyn PythonUsageSource,
173 file: &ProjectFile,
174 source: &str,
175 node: Node<'_>,
176 raw_symbol: &str,
177) -> Vec<CodeUnit> {
178 let raw_symbol = raw_symbol.trim();
179 if raw_symbol.is_empty() {
180 return Vec::new();
181 }
182
183 if let Some(owner) = annotation_scope_owner_class(graph, file, source, node) {
184 let owner_candidates: Vec<_> = exact_owner_annotation_members(graph, &owner, raw_symbol)
185 .into_iter()
186 .filter(|candidate| !candidate.is_function())
187 .collect();
188 if !owner_candidates.is_empty() {
189 return owner_candidates;
190 }
191 }
192
193 let mut candidates = Vec::new();
194 if let Some(binding) = python.import_binder_of(file).bindings.get(raw_symbol)
195 && binding.kind == ImportKind::Named
196 && let Some(imported) = binding.imported_name.as_ref()
197 {
198 let fqn = format!("{}.{}", binding.module_specifier, imported);
199 let mut imported_candidates =
200 resolve_fqn_candidates(python, &fqn, |name| graph.index.definitions(name).collect());
201 imported_candidates.retain(|candidate| {
202 !candidate.is_module() || candidate.fq_name() != binding.module_specifier
203 });
204 candidates.extend(imported_candidates);
205 }
206
207 candidates.extend(
208 graph
209 .index
210 .top_level_declarations(file)
211 .into_iter()
212 .filter(|code_unit| {
213 !code_unit.is_module()
214 && code_unit.identifier() == raw_symbol
215 && graph
216 .index
217 .parent_of(code_unit)
218 .is_some_and(|parent| parent.is_module())
219 }),
220 );
221
222 candidates.retain(|candidate| !candidate.is_function());
223 candidates.sort();
224 candidates.dedup();
225 candidates
226}
227
228pub fn annotation_reference_candidates(
235 graph: &PythonGraphSource<'_>,
236 python: &dyn PythonUsageSource,
237 file: &ProjectFile,
238 source: &str,
239 node: Node<'_>,
240 target_self_file: bool,
241) -> Option<Vec<CodeUnit>> {
242 if !is_annotation_reference_node(node) {
243 return None;
244 }
245
246 let mut candidates = match node.kind() {
247 "identifier" => {
248 let mut candidates = resolve_bare_annotation_symbol(
249 graph,
250 python,
251 file,
252 source,
253 node,
254 node_text(node, source),
255 );
256 if candidates.is_empty() {
257 candidates.extend(resolve_receiver_type(
258 graph,
259 python,
260 file,
261 node_text(node, source),
262 target_self_file,
263 ));
264 }
265 candidates
266 }
267 "string_content" => {
268 let Some(string) = node.parent() else {
269 return Some(Vec::new());
270 };
271 let Some(ranges) = python_deferred_annotation_identifier_ranges(string, source, None)
272 else {
273 return Some(Vec::new());
274 };
275 let mut candidates = Vec::new();
276 for range in ranges {
277 let Some(symbol) = source.get(range.start_byte..range.end_byte) else {
278 continue;
279 };
280 let mut symbol_candidates =
281 resolve_bare_annotation_symbol(graph, python, file, source, node, symbol);
282 if symbol_candidates.is_empty() {
283 symbol_candidates.extend(resolve_receiver_type(
284 graph,
285 python,
286 file,
287 symbol,
288 target_self_file,
289 ));
290 }
291 candidates.extend(symbol_candidates);
292 }
293 candidates
294 }
295 "attribute" => resolve_annotation_attribute_types(graph, python, file, source, node),
296 _ => Vec::new(),
297 };
298 candidates.sort();
299 candidates.dedup();
300 Some(candidates)
301}
302
303#[allow(clippy::too_many_arguments)]
307pub fn annotation_reference_candidates_at_focus(
308 graph: &PythonGraphSource<'_>,
309 python: &dyn PythonUsageSource,
310 file: &ProjectFile,
311 source: &str,
312 node: Node<'_>,
313 focus_start: usize,
314 focus_end: usize,
315 target_self_file: bool,
316) -> Option<Vec<CodeUnit>> {
317 if !is_annotation_reference_node(node) {
318 return None;
319 }
320
321 let deferred_tree = if node.kind() == "string_content" {
322 Some(python_deferred_annotation_tree(
323 node.parent()?,
324 source,
325 None,
326 )?)
327 } else {
328 None
329 };
330 let search_root = deferred_tree.as_ref().map_or(node, |tree| tree.root_node());
331 let focused = search_root.descendant_for_byte_range(focus_start, focus_end)?;
332 if focused.kind() != "identifier"
333 || focused.start_byte() != focus_start
334 || focused.end_byte() != focus_end
335 {
336 return Some(Vec::new());
337 }
338
339 let mut path = focused;
340 while let Some(parent) = path.parent() {
341 if parent.kind() != "attribute" {
342 break;
343 }
344 path = parent;
345 }
346 if path.kind() == "attribute" {
347 return Some(focused_annotation_attribute_candidates(
348 graph, python, file, source, node, path, focused,
349 ));
350 }
351
352 let symbol = node_text(focused, source);
353 let mut candidates = resolve_bare_annotation_symbol(graph, python, file, source, node, symbol);
354 if candidates.is_empty() {
355 candidates.extend(resolve_receiver_type(
356 graph,
357 python,
358 file,
359 symbol,
360 target_self_file,
361 ));
362 }
363 candidates.sort();
364 candidates.dedup();
365 Some(candidates)
366}
367
368fn focused_annotation_attribute_candidates(
369 graph: &PythonGraphSource<'_>,
370 python: &dyn PythonUsageSource,
371 file: &ProjectFile,
372 source: &str,
373 scope_node: Node<'_>,
374 path: Node<'_>,
375 focused: Node<'_>,
376) -> Vec<CodeUnit> {
377 let Some((root, attributes)) = annotation_attribute_chain(path) else {
378 return Vec::new();
379 };
380 if root.id() == focused.id() {
381 return resolve_bare_annotation_symbol(
382 graph,
383 python,
384 file,
385 source,
386 scope_node,
387 node_text(root, source),
388 );
389 }
390
391 if attributes
392 .last()
393 .is_some_and(|node| node.id() == focused.id())
394 {
395 let candidates = namespace_qualified_declarations(graph, python, file, source, path);
396 if !candidates.is_empty() {
397 return candidates;
398 }
399 }
400
401 let owners: Vec<_> = resolve_bare_annotation_symbol(
402 graph,
403 python,
404 file,
405 source,
406 scope_node,
407 node_text(root, source),
408 )
409 .into_iter()
410 .filter(CodeUnit::is_class)
411 .collect();
412 let [owner] = owners.as_slice() else {
413 return Vec::new();
414 };
415 let mut owner = owner.clone();
416 for attribute in attributes {
417 let candidates = exact_nested_annotation_class(graph, &owner, node_text(attribute, source));
418 if attribute.id() == focused.id() {
419 return candidates;
420 }
421 let [next] = candidates.as_slice() else {
422 return Vec::new();
423 };
424 owner = next.clone();
425 }
426 Vec::new()
427}
428
429pub fn annotation_class_qualifier_site<'tree>(
432 graph: &PythonGraphSource<'_>,
433 python: &dyn PythonUsageSource,
434 file: &ProjectFile,
435 source: &str,
436 node: Node<'tree>,
437 target: &CodeUnit,
438) -> Option<Node<'tree>> {
439 if node.kind() != "attribute" || !target.is_class() || !is_annotation_reference_node(node) {
440 return None;
441 }
442
443 let (root, attributes) = annotation_attribute_chain(node)?;
444 let owners: Vec<_> =
445 resolve_bare_annotation_symbol(graph, python, file, source, root, node_text(root, source))
446 .into_iter()
447 .filter(CodeUnit::is_class)
448 .collect();
449 let [owner] = owners.as_slice() else {
450 return None;
451 };
452 let mut owner = owner.clone();
453 if &owner == target {
454 return Some(root);
455 }
456
457 let qualifier_count = attributes.len().saturating_sub(1);
460 for attribute in attributes.into_iter().take(qualifier_count) {
461 let next_candidates =
462 exact_nested_annotation_class(graph, &owner, node_text(attribute, source));
463 let [next] = next_candidates.as_slice() else {
464 return None;
465 };
466 owner = next.clone();
467 if &owner == target {
468 return Some(attribute);
469 }
470 }
471
472 None
473}
474
475fn resolve_annotation_attribute_types(
476 graph: &PythonGraphSource<'_>,
477 python: &dyn PythonUsageSource,
478 file: &ProjectFile,
479 source: &str,
480 node: Node<'_>,
481) -> Vec<CodeUnit> {
482 let mut candidates = namespace_qualified_declarations(graph, python, file, source, node);
487 let Some((root, attributes)) = annotation_attribute_chain(node) else {
488 return candidates;
489 };
490 let root_text = node_text(root, source);
491 let owners: Vec<_> =
492 resolve_bare_annotation_symbol(graph, python, file, source, root, root_text)
493 .into_iter()
494 .filter(CodeUnit::is_class)
495 .collect();
496 let [owner] = owners.as_slice() else {
497 return candidates;
498 };
499 let mut owner = owner.clone();
500
501 for attribute in attributes {
502 let segment = node_text(attribute, source);
503 let next_candidates = exact_nested_annotation_class(graph, &owner, segment);
504 let [next] = next_candidates.as_slice() else {
505 return candidates;
506 };
507 owner = next.clone();
508 }
509
510 candidates.push(owner);
511 candidates
512}
513
514fn annotation_attribute_chain(node: Node<'_>) -> Option<(Node<'_>, Vec<Node<'_>>)> {
515 let mut attributes = Vec::new();
516 let mut current = node;
517 while current.kind() == "attribute" {
518 attributes.push(current.child_by_field_name("attribute")?);
519 current = current.child_by_field_name("object")?;
520 }
521 if current.kind() != "identifier" || attributes.is_empty() {
522 return None;
523 }
524 attributes.reverse();
525 Some((current, attributes))
526}
527
528fn exact_nested_annotation_class(
529 graph: &PythonGraphSource<'_>,
530 owner: &CodeUnit,
531 segment: &str,
532) -> Vec<CodeUnit> {
533 let mut candidates: Vec<_> = exact_owner_annotation_members(graph, owner, segment)
534 .into_iter()
535 .filter(CodeUnit::is_class)
536 .collect();
537 candidates.sort();
538 candidates.dedup();
539 candidates
540}
541
542fn exact_owner_annotation_members(
543 graph: &PythonGraphSource<'_>,
544 owner: &CodeUnit,
545 segment: &str,
546) -> Vec<CodeUnit> {
547 let mut candidates: Vec<_> = graph
548 .index
549 .declarations(owner.source())
550 .into_iter()
551 .filter(|unit| {
552 unit.identifier() == segment
553 && graph
554 .index
555 .parent_of(unit)
556 .is_some_and(|parent| parent.fq_name() == owner.fq_name())
557 })
558 .collect();
559 candidates.sort();
560 candidates.dedup();
561 candidates
562}
563
564fn annotation_scope_owner_class(
565 graph: &PythonGraphSource<'_>,
566 file: &ProjectFile,
567 source: &str,
568 node: Node<'_>,
569) -> Option<CodeUnit> {
570 if !annotation_expression_is_class_scoped(node) {
571 return None;
572 }
573 let range = Range {
574 start_byte: node.start_byte(),
575 end_byte: node.end_byte(),
576 start_line: 0,
577 end_line: 0,
578 };
579 if let Some(enclosing) = graph.index.enclosing_code_unit(file, &range) {
580 if enclosing.is_class() {
581 return Some(enclosing);
582 }
583 if let Some(owner) = target_owner_code_unit(graph.index, &enclosing) {
584 return Some(owner);
585 }
586 }
587 structural_annotation_owner_class(graph, file, source, node)
588}
589
590fn annotation_expression_is_class_scoped(node: Node<'_>) -> bool {
591 let site_start = node.start_byte();
592 let site_end = node.end_byte();
593 let mut current = node;
594 while let Some(parent) = current.parent() {
595 if matches!(parent.kind(), "function_definition" | "lambda")
596 && parent
597 .child_by_field_name("body")
598 .is_some_and(|body| body.start_byte() <= site_start && site_end <= body.end_byte())
599 {
600 return false;
601 }
602 if parent.kind() == "class_definition" {
603 return true;
604 }
605 current = parent;
606 }
607 false
608}
609
610fn structural_annotation_owner_class(
611 graph: &PythonGraphSource<'_>,
612 file: &ProjectFile,
613 source: &str,
614 node: Node<'_>,
615) -> Option<CodeUnit> {
616 let mut current = node;
617 while let Some(parent) = current.parent() {
618 if parent.kind() == "class_definition" {
619 let name = node_text(parent.child_by_field_name("name")?, source).trim();
620 if name.is_empty() {
621 return None;
622 }
623 let class_range = Range {
624 start_byte: parent.start_byte(),
625 end_byte: parent.end_byte(),
626 start_line: 0,
627 end_line: 0,
628 };
629 let mut matches: Vec<_> = graph
630 .index
631 .declarations(file)
632 .into_iter()
633 .filter(|unit| unit.is_class() && unit.identifier() == name)
634 .filter(|unit| {
635 graph
636 .index
637 .ranges(unit)
638 .into_iter()
639 .any(|range| range.contains(&class_range))
640 })
641 .collect();
642 matches.sort();
643 matches.dedup();
644 let [owner] = matches.as_slice() else {
645 return None;
646 };
647 return Some(owner.clone());
648 }
649 current = parent;
650 }
651 None
652}
653
654fn is_annotation_reference_node(node: Node<'_>) -> bool {
655 if !matches!(node.kind(), "identifier" | "attribute" | "string_content") {
656 return false;
657 }
658 python_node_is_in_annotation(node)
659}
660
661pub fn resolve_constructor_types(
666 graph: &PythonGraphSource<'_>,
667 python: &dyn PythonUsageSource,
668 file: &ProjectFile,
669 source: &str,
670 function: Node<'_>,
671) -> Vec<CodeUnit> {
672 let candidates = match function.kind() {
673 "identifier" => {
674 let local = node_text(function, source);
675 if local.is_empty() {
676 return Vec::new();
677 }
678 let binder = python.import_binder_of(file);
679 let fqn = match binder.bindings.get(local) {
680 Some(binding) if binding.kind == ImportKind::Named => binding
681 .imported_name
682 .as_ref()
683 .map(|imported| format!("{}.{}", binding.module_specifier, imported)),
684 _ => graph
685 .index
686 .declarations(file)
687 .into_iter()
688 .find(|unit| unit.is_class() && unit.identifier() == local)
689 .map(|unit| unit.fq_name()),
690 };
691 let Some(fqn) = fqn else {
692 return Vec::new();
693 };
694 resolve_fqn_candidates(python, &fqn, |name| graph.index.definitions(name).collect())
695 }
696 "attribute" => namespace_qualified_declarations(graph, python, file, source, function),
697 _ => Vec::new(),
698 };
699 let mut classes: Vec<CodeUnit> = candidates.into_iter().filter(CodeUnit::is_class).collect();
702 classes.sort();
703 classes.dedup();
704 classes
705}
706
707pub fn resolve_callable_parameter_default_types(
713 graph: &PythonGraphSource<'_>,
714 python: &dyn PythonUsageSource,
715 file: &ProjectFile,
716 source: &str,
717 reference: Node<'_>,
718 local_name: &str,
719) -> Vec<CodeUnit> {
720 let site_start = reference.start_byte();
721 let site_end = reference.end_byte();
722 let mut current = reference;
723 while let Some(parent) = current.parent() {
724 current = parent;
725 if !matches!(current.kind(), "function_definition" | "lambda") {
726 continue;
727 }
728 if current
729 .child_by_field_name("body")
730 .is_none_or(|body| !(body.start_byte() <= site_start && site_end <= body.end_byte()))
731 {
732 continue;
733 }
734 let Some(parameters) = current.child_by_field_name("parameters") else {
735 return Vec::new();
736 };
737 let mut cursor = parameters.walk();
738 for parameter in parameters.named_children(&mut cursor) {
739 let name = if parameter.kind() == "identifier" {
740 Some(parameter)
741 } else {
742 parameter.child_by_field_name("name")
743 };
744 if name.is_none_or(|name| node_text(name, source) != local_name) {
745 continue;
746 }
747 let Some(value) = parameter.child_by_field_name("value") else {
748 return Vec::new();
749 };
750 return resolve_constructor_types(graph, python, file, source, value);
751 }
752 return Vec::new();
753 }
754 Vec::new()
755}
756
757fn namespace_qualified_declarations(
764 graph: &PythonGraphSource<'_>,
765 python: &dyn PythonUsageSource,
766 file: &ProjectFile,
767 source: &str,
768 node: Node<'_>,
769) -> Vec<CodeUnit> {
770 let binder = python.import_binder_of(file);
771 let Some(fqn) = namespace_constructor_fqn(&binder, source, node) else {
772 return Vec::new();
773 };
774 let mut candidates =
775 resolve_fqn_candidates(python, &fqn, |name| graph.index.definitions(name).collect());
776 candidates.sort();
777 candidates.dedup();
778 candidates
779}
780
781fn namespace_constructor_fqn(
782 binder: &ImportBinder,
783 source: &str,
784 function: Node<'_>,
785) -> Option<String> {
786 let mut attributes = Vec::new();
787 let mut current = function;
788 while current.kind() == "attribute" {
789 let attribute = current.child_by_field_name("attribute")?;
790 let text = node_text(attribute, source);
791 if text.is_empty() {
792 return None;
793 }
794 attributes.push(text);
795 current = current.child_by_field_name("object")?;
796 }
797 if current.kind() != "identifier" {
798 return None;
799 }
800 let root = node_text(current, source);
801 let binding = binder.bindings.get(root)?;
802 if binding.kind != ImportKind::Namespace {
803 return None;
804 }
805 let mut fqn = binding.module_specifier.clone();
806 for attribute in attributes.into_iter().rev() {
807 fqn.push('.');
808 fqn.push_str(attribute);
809 }
810 Some(fqn)
811}
812
813fn node_text<'a>(node: Node<'_>, source: &'a str) -> &'a str {
814 brokk_bifrost_core::analyzer::common::node_source_text(node, source)
815}
816
817fn resolve_indexed_receiver_type(
818 index: &dyn CodeUnitIndex,
819 lookup: &dyn BoundedDefinitionLookup,
820 file: &ProjectFile,
821 raw_type: &str,
822) -> Option<CodeUnit> {
823 module_fqn_for_file(index, file)
824 .into_iter()
825 .flat_map(|module| lookup.types_in_package(&module, raw_type))
826 .chain(lookup.fqn(raw_type))
827 .chain(lookup.by_normalized_fqn(raw_type))
828 .find(|code_unit| code_unit.identifier() == raw_type && code_unit.is_class())
829}
830
831fn module_fqn_for_file(index: &dyn CodeUnitIndex, file: &ProjectFile) -> Option<String> {
832 index
833 .declarations(file)
834 .into_iter()
835 .find(|code_unit| code_unit.is_module())
836 .map(|code_unit| code_unit.fq_name())
837 .or_else(|| {
838 index
839 .declarations(file)
840 .into_iter()
841 .find(|code_unit| !code_unit.package_name().is_empty())
842 .map(|code_unit| code_unit.package_name().to_string())
843 })
844}
845
846pub fn normalized_receiver_type(annotation: &str) -> Option<String> {
847 let annotation = unwrap_python_string_annotation(annotation.trim());
848 let annotation = unwrap_supported_receiver_wrapper(annotation);
849 if annotation.is_empty()
850 || annotation.contains('|')
851 || annotation.contains('[')
852 || annotation.contains(']')
853 || annotation.contains(',')
854 || annotation.contains('(')
855 || annotation.contains(')')
856 || annotation.contains('{')
857 || annotation.contains('}')
858 || annotation.contains(':')
859 {
860 return None;
861 }
862 Some(annotation.to_string())
863}
864
865fn unwrap_python_string_annotation(annotation: &str) -> &str {
866 if annotation.len() >= 2 {
867 let bytes = annotation.as_bytes();
868 let first = bytes[0];
869 let last = bytes[annotation.len() - 1];
870 if (first == b'\'' || first == b'"') && first == last {
871 return annotation[1..annotation.len() - 1].trim();
872 }
873 }
874 annotation
875}
876
877fn unwrap_supported_receiver_wrapper(annotation: &str) -> &str {
878 let mut current = annotation.trim();
879 loop {
880 let next = current
881 .strip_prefix("Optional[")
882 .or_else(|| current.strip_prefix("typing.Optional["))
883 .and_then(|inner| inner.strip_suffix(']'))
884 .map(str::trim);
885 let Some(unwrapped) = next else {
886 return current;
887 };
888 current = unwrapped;
889 }
890}
891
892pub fn receiver_annotation_matches_target(
893 annotation: &str,
894 edges: &[ImportEdge],
895 target_short: &str,
896 target_self_file: bool,
897) -> bool {
898 let annotation = annotation.trim();
899 if annotation.is_empty() {
900 return false;
901 }
902 if annotation.contains('|')
903 || annotation.contains('[')
904 || annotation.contains(']')
905 || annotation.contains(',')
906 || annotation.contains('(')
907 || annotation.contains(')')
908 {
909 return false;
910 }
911 if annotation == target_short {
912 return target_self_file || edges.iter().any(|edge| edge.local_name == target_short);
913 }
914
915 let segments = parse_symbol_path(Language::Python, annotation);
921 let Some((member, qualifier_parts)) = segments.split_last() else {
922 return false;
923 };
924 if qualifier_parts.is_empty() {
925 return false;
926 }
927 let qualifier = qualifier_parts.join(".");
928 let member = member.as_str();
929 if member != target_short {
930 return false;
931 }
932 edges.iter().any(|edge| {
933 matches!(edge.kind, ImportEdgeKind::Namespace)
934 && (edge.local_name == qualifier
935 || qualifier.ends_with(&format!(".{}", edge.local_name)))
936 })
937}
938
939