1use crate::bindings::python_direct_scope_bindings_bounded;
8use crate::declarations::{parse_python_tree, py_node_text, python_module_name};
9use crate::graph_support::{
10 PythonSource, import_binder_from_imports, public_declarations_in_module,
11 resolve_module_code_unit, resolve_module_code_units_batch,
12};
13use brokk_bifrost_core::analyzer::common::node_source_text;
14use brokk_bifrost_core::analyzer::model::{
15 ImportInfo, StructuredImportPath, StructuredImportPathKind,
16};
17use brokk_bifrost_core::analyzer::query_token::QueryToken;
18use brokk_bifrost_core::analyzer::usages::model::{ExportEntry, ImportBinding, ImportKind};
19use brokk_bifrost_core::analyzer::{CodeUnit, CodeUnitIndex, ProjectFile};
20use brokk_bifrost_core::hash::{HashMap, HashSet};
21use std::collections::VecDeque;
22use tree_sitter::Node;
23
24#[derive(Debug, Clone, PartialEq, Eq)]
25pub struct PythonModuleReplacement {
26 pub target_module: String,
27}
28
29pub(crate) fn imported_module_assignment_at(
39 reference: Node<'_>,
40 local: &str,
41 source: &str,
42 mut is_importlib_binding: impl FnMut(&str) -> bool,
43) -> Option<String> {
44 let suite = enclosing_execution_suite(reference)?;
45 let mut resolved = None;
46 let mut cursor = suite.walk();
47 for statement in suite.named_children(&mut cursor) {
48 if statement.start_byte() >= reference.start_byte()
49 || (statement.start_byte() <= reference.start_byte()
50 && reference.end_byte() <= statement.end_byte())
51 {
52 break;
53 }
54
55 let expression = if statement.kind() == "expression_statement" {
56 statement.named_child(0).unwrap_or(statement)
57 } else {
58 statement
59 };
60 if expression.kind() == "assignment"
61 && expression.child_by_field_name("left").is_some_and(|left| {
62 left.kind() == "identifier" && node_source_text(left, source) == local
63 })
64 {
65 resolved = expression
66 .child_by_field_name("right")
67 .and_then(|right| importlib_call_module(right, source, &mut is_importlib_binding));
68 continue;
69 }
70
71 let binds_local = python_direct_scope_bindings_bounded(statement, source, || true)
72 .expect("unbounded binding collection cannot be cancelled")
73 .into_iter()
74 .any(|binding| node_source_text(binding.declaration, source) == local);
75 if binds_local {
76 resolved = None;
77 }
78 }
79 resolved
80}
81
82fn enclosing_execution_suite(mut node: Node<'_>) -> Option<Node<'_>> {
83 while let Some(parent) = node.parent() {
84 if matches!(parent.kind(), "function_definition" | "lambda") {
85 return parent.child_by_field_name("body");
86 }
87 if parent.kind() == "module" {
88 return Some(parent);
89 }
90 node = parent;
91 }
92 None
93}
94
95fn importlib_call_module(
96 call: Node<'_>,
97 source: &str,
98 is_importlib_binding: &mut impl FnMut(&str) -> bool,
99) -> Option<String> {
100 if call.kind() != "call" {
101 return None;
102 }
103 let function = call.child_by_field_name("function")?;
104 if function.kind() != "attribute" {
105 return None;
106 }
107 let (object, attribute) = (
108 function.child_by_field_name("object")?,
109 function.child_by_field_name("attribute")?,
110 );
111 if object.kind() != "identifier"
112 || attribute.kind() != "identifier"
113 || node_source_text(attribute, source) != "import_module"
114 || !is_importlib_binding(node_source_text(object, source))
115 {
116 return None;
117 }
118
119 let arguments = call.child_by_field_name("arguments")?;
120 if arguments.named_child_count() != 1 {
121 return None;
122 }
123 let literal = arguments.named_child(0)?;
124 if literal.kind() != "string"
125 || literal
126 .parent()
127 .is_some_and(|parent| parent.kind() == "concatenated_string")
128 {
129 return None;
130 }
131 let mut content = None;
132 for index in 0..literal.named_child_count() {
133 let child = literal.named_child(index)?;
134 match child.kind() {
135 "string_start" | "string_end" => {}
136 "string_content" if content.is_none() && child.named_child_count() == 0 => {
137 content = Some(child);
138 }
139 _ => return None,
140 }
141 }
142 let content = content?;
143 let module = node_source_text(content, source);
144 (!module.is_empty() && !module.starts_with('.')).then(|| module.to_string())
145}
146
147pub(crate) fn literal_importlib_modules(
152 source: &str,
153 bindings: &HashMap<String, ImportBinding>,
154) -> HashSet<String> {
155 let Some(tree) = parse_python_tree(source) else {
156 return HashSet::default();
157 };
158 let mut modules = HashSet::default();
159 let mut stack = vec![tree.root_node()];
160 while let Some(node) = stack.pop() {
161 if node.kind() == "assignment"
162 && node
163 .child_by_field_name("left")
164 .is_some_and(|left| left.kind() == "identifier")
165 && let Some(module) = node.child_by_field_name("right").and_then(|right| {
166 importlib_call_module(right, source, &mut |local| {
167 bindings.get(local).is_some_and(|binding| {
168 binding.kind == ImportKind::Namespace
169 && binding
170 .namespace_imported_module
171 .as_deref()
172 .unwrap_or(&binding.module_specifier)
173 == "importlib"
174 })
175 })
176 })
177 {
178 modules.insert(module);
179 }
180 let mut cursor = node.walk();
181 stack.extend(node.named_children(&mut cursor));
182 }
183 modules
184}
185
186fn module_replacement_from_assignment(
194 assignment: Node<'_>,
195 source: &str,
196 bindings: &HashMap<String, ImportBinding>,
197) -> Option<PythonModuleReplacement> {
198 let (left, right) = (
199 assignment.child_by_field_name("left")?,
200 assignment.child_by_field_name("right")?,
201 );
202 if left.kind() != "subscript" || right.kind() != "identifier" {
203 return None;
204 }
205 let (value, subscript) = (
206 left.child_by_field_name("value")?,
207 left.child_by_field_name("subscript")?,
208 );
209 if value.kind() != "attribute"
210 || subscript.kind() != "identifier"
211 || node_source_text(subscript, source) != "__name__"
212 {
213 return None;
214 }
215 let (sys_local, modules) = (
216 value.child_by_field_name("object")?,
217 value.child_by_field_name("attribute")?,
218 );
219 if sys_local.kind() != "identifier"
220 || modules.kind() != "identifier"
221 || node_source_text(modules, source) != "modules"
222 {
223 return None;
224 }
225 let sys_binding = bindings.get(node_source_text(sys_local, source))?;
226 let sys_module = sys_binding
227 .namespace_imported_module
228 .as_deref()
229 .unwrap_or(&sys_binding.module_specifier);
230 if sys_binding.kind != ImportKind::Namespace || sys_module != "sys" {
231 return None;
232 }
233
234 let target_binding = bindings.get(node_source_text(right, source))?;
235 if target_binding.kind != ImportKind::Namespace {
236 return None;
237 }
238 Some(PythonModuleReplacement {
239 target_module: target_binding
240 .namespace_imported_module
241 .clone()
242 .unwrap_or_else(|| target_binding.module_specifier.clone()),
243 })
244}
245
246fn remove_direct_scope_bindings(
247 statement: Node<'_>,
248 source: &str,
249 bindings: &mut HashMap<String, ImportBinding>,
250) {
251 let mut stack = vec![statement];
252 while let Some(node) = stack.pop() {
253 for binding in python_direct_scope_bindings_bounded(node, source, || true)
254 .expect("unbounded binding collection cannot be cancelled")
255 {
256 bindings.remove(node_source_text(binding.declaration, source));
257 }
258 if matches!(
259 node.kind(),
260 "function_definition" | "class_definition" | "lambda"
261 ) {
262 continue;
263 }
264 let mut cursor = node.walk();
265 stack.extend(node.named_children(&mut cursor));
266 }
267}
268
269pub fn parse_python_import_infos(source: &str) -> Vec<ImportInfo> {
273 let mut parser = tree_sitter::Parser::new();
274 parser
275 .set_language(&tree_sitter_python::LANGUAGE.into())
276 .expect("failed to load Python parser");
277 let Some(tree) = parser.parse(source, None) else {
278 return Vec::new();
279 };
280 let mut pending = vec![tree.root_node()];
281 let mut imports = Vec::new();
282 while let Some(node) = pending.pop() {
283 if matches!(node.kind(), "import_statement" | "import_from_statement") {
284 imports.extend(python_import_infos_from_node(node, source));
285 continue;
286 }
287 let mut cursor = node.walk();
288 pending.extend(node.named_children(&mut cursor));
289 }
290 imports
291}
292
293#[derive(Debug, Clone, PartialEq, Eq)]
294pub struct PythonImportBinding {
295 pub start_byte: usize,
296 pub scope_start_byte: usize,
297 pub scope_end_byte: usize,
298 function_scoped: bool,
299 pub local_name: String,
300 pub qualified_name: String,
301 pub consumed_attributes: usize,
302}
303
304impl PythonImportBinding {
305 pub fn is_function_scoped(&self) -> bool {
306 self.function_scoped
307 }
308}
309
310pub fn parse_python_import_bindings(source: &str) -> Vec<PythonImportBinding> {
314 let mut parser = tree_sitter::Parser::new();
315 parser
316 .set_language(&tree_sitter_python::LANGUAGE.into())
317 .expect("failed to load Python parser");
318 let Some(tree) = parser.parse(source, None) else {
319 return Vec::new();
320 };
321 python_import_bindings_from_tree(tree.root_node(), source)
322}
323
324pub fn python_import_bindings_from_tree(root: Node<'_>, source: &str) -> Vec<PythonImportBinding> {
328 let mut pending = vec![root];
329 let mut nodes = Vec::new();
330 while let Some(node) = pending.pop() {
331 if matches!(node.kind(), "import_statement" | "import_from_statement") {
332 nodes.push(node);
333 continue;
334 }
335 let mut cursor = node.walk();
336 pending.extend(node.named_children(&mut cursor));
337 }
338 nodes.sort_by_key(Node::start_byte);
339 nodes
340 .into_iter()
341 .flat_map(|node| {
342 let (scope_start_byte, scope_end_byte, function_scoped) =
343 python_import_binding_scope(node, source.len());
344 python_import_infos_from_node(node, source)
345 .into_iter()
346 .filter_map(move |import| {
347 let path = import.path.as_ref()?;
348 let details = python_import_details(&import)?;
349 match details {
350 PythonImportDetails::Import { module, alias } => {
351 let consumed_attributes = if alias.is_some() {
352 0
353 } else {
354 path.segments.len().saturating_sub(1)
355 };
356 Some(PythonImportBinding {
357 start_byte: node.start_byte(),
358 scope_start_byte,
359 scope_end_byte,
360 function_scoped,
361 local_name: alias.or_else(|| path.segments.first().cloned())?,
362 qualified_name: module,
363 consumed_attributes,
364 })
365 }
366 PythonImportDetails::FromImport {
367 module,
368 name,
369 alias,
370 wildcard: false,
371 } => Some(PythonImportBinding {
372 start_byte: node.start_byte(),
373 scope_start_byte,
374 scope_end_byte,
375 function_scoped,
376 local_name: alias.unwrap_or(name.clone()),
377 qualified_name: format!("{module}.{name}"),
378 consumed_attributes: 0,
379 }),
380 PythonImportDetails::FromImport { wildcard: true, .. } => None,
381 }
382 })
383 })
384 .collect()
385}
386
387fn python_import_binding_scope(node: Node<'_>, source_len: usize) -> (usize, usize, bool) {
388 let mut parent = node.parent();
389 while let Some(scope) = parent {
390 if matches!(scope.kind(), "function_definition" | "lambda") {
391 return (scope.start_byte(), scope.end_byte(), true);
392 }
393 parent = scope.parent();
394 }
395 (0, source_len, false)
396}
397
398#[cfg(test)]
399mod tests {
400 use super::*;
401 use brokk_bifrost_core::analyzer::usages::model::ImportBinder;
402
403 fn importlib_assignment_for(source: &str, importlib_local: &str) -> Option<String> {
404 let tree = parse_python_tree(source).expect("valid Python fixture");
405 let reference_start = source.rfind("canonical.setup").expect("reference site");
406 let mut stack = vec![tree.root_node()];
407 let mut reference = None;
408 while let Some(node) = stack.pop() {
409 if node.kind() == "identifier"
410 && node.start_byte() == reference_start
411 && node_source_text(node, source) == "canonical"
412 {
413 reference = Some(node);
414 break;
415 }
416 let mut cursor = node.walk();
417 stack.extend(node.named_children(&mut cursor));
418 }
419 imported_module_assignment_at(
420 reference.expect("canonical reference node"),
421 "canonical",
422 source,
423 |local| local == importlib_local,
424 )
425 }
426
427 #[test]
428 fn existing_tree_import_bindings_preserve_function_scope_and_relative_identity() {
429 let source = r#"def register():
430 from .workers.special import Worker
431 return Worker()
432"#;
433 let tree = parse_python_tree(source).expect("Python tree");
434 let bindings = python_import_bindings_from_tree(tree.root_node(), source);
435
436 assert_eq!(bindings.len(), 1, "{bindings:#?}");
437 assert!(bindings[0].is_function_scoped(), "{bindings:#?}");
438 assert_eq!(bindings[0].local_name, "Worker");
439 assert_eq!(bindings[0].qualified_name, ".workers.special.Worker");
440 assert!(bindings[0].scope_start_byte <= bindings[0].start_byte);
441 assert!(bindings[0].scope_end_byte >= source.rfind("Worker()").unwrap());
442 }
443
444 #[test]
445 fn importlib_assignment_requires_a_direct_literal_dominating_write() {
446 let valid = r#"import importlib as loader
447
448def test():
449 canonical = loader.import_module("routes.contacts.contacts_routes")
450 assert canonical.setup_contacts_routes
451"#;
452 assert_eq!(
453 importlib_assignment_for(valid, "loader"),
454 Some("routes.contacts.contacts_routes".to_string())
455 );
456
457 for near_miss in [
458 r#"import importlib as loader
459def test(module_name):
460 canonical = loader.import_module(module_name)
461 assert canonical.setup_contacts_routes
462"#,
463 r#"import importlib as loader
464def test():
465 if enabled:
466 canonical = loader.import_module("routes.contacts.contacts_routes")
467 assert canonical.setup_contacts_routes
468"#,
469 r#"import importlib as loader
470def test():
471 canonical = loader.import_module("routes.contacts.contacts_routes")
472 canonical = object()
473 assert canonical.setup_contacts_routes
474"#,
475 ] {
476 assert_eq!(
477 importlib_assignment_for(near_miss, "loader"),
478 None,
479 "near miss must not infer a module receiver: {near_miss:?}"
480 );
481 }
482 assert_eq!(
483 importlib_assignment_for(valid, "shadowed_loader"),
484 None,
485 "a similarly shaped non-importlib receiver must not resolve"
486 );
487 }
488
489 fn replacement_for(source: &str, binder: &ImportBinder) -> Option<PythonModuleReplacement> {
490 let tree = parse_python_tree(source).expect("valid Python fixture");
491 let root = tree.root_node();
492 let mut replacement = None;
493 let mut cursor = root.walk();
494 for statement in root.named_children(&mut cursor) {
495 let statement = if statement.kind() == "expression_statement" {
496 statement.named_child(0).expect("fixture expression")
497 } else {
498 statement
499 };
500 if statement.kind() == "assignment"
501 && let Some(next) =
502 module_replacement_from_assignment(statement, source, &binder.bindings)
503 && replacement.replace(next).is_some()
504 {
505 return None;
506 }
507 }
508 replacement
509 }
510
511 #[test]
512 fn module_replacement_requires_exact_structured_import_bindings() {
513 let source = r#"import sys as _sys
514from routes.contacts import contacts_routes as _canonical
515
516_sys.modules[__name__] = _canonical
517"#;
518 let mut binder = ImportBinder::empty();
519 binder.bindings.insert(
520 "_sys".to_string(),
521 ImportBinding {
522 module_specifier: "sys".to_string(),
523 namespace_imported_module: Some("sys".to_string()),
524 kind: ImportKind::Namespace,
525 imported_name: None,
526 },
527 );
528 binder.bindings.insert(
529 "_canonical".to_string(),
530 ImportBinding {
531 module_specifier: "routes.contacts.contacts_routes".to_string(),
532 namespace_imported_module: None,
533 kind: ImportKind::Namespace,
534 imported_name: None,
535 },
536 );
537
538 assert_eq!(
539 replacement_for(source, &binder),
540 Some(PythonModuleReplacement {
541 target_module: "routes.contacts.contacts_routes".to_string(),
542 })
543 );
544
545 for near_miss in [
546 "cache.modules[__name__] = _canonical\n",
547 "_sys.modules[module_name] = _canonical\n",
548 "_sys.modules[__name__] = build()\n",
549 "def replace():\n _sys.modules[__name__] = _canonical\n",
550 ] {
551 assert_eq!(
552 replacement_for(near_miss, &binder),
553 None,
554 "near miss must not replace module identity: {near_miss:?}"
555 );
556 }
557 }
558}
559
560pub fn module_replacement_of(
561 python: &dyn PythonSource,
562 file: &ProjectFile,
563 source: &str,
564) -> Option<PythonModuleReplacement> {
565 let tree = parse_python_tree(source)?;
566 let root = tree.root_node();
567 let mut bindings: HashMap<String, ImportBinding> = HashMap::default();
568 let mut replacement = None;
569 let mut cursor = root.walk();
570 for statement in root.named_children(&mut cursor) {
571 let statement = if statement.kind() == "expression_statement" {
572 let Some(expression) = statement.named_child(0) else {
573 continue;
574 };
575 expression
576 } else {
577 statement
578 };
579 match statement.kind() {
580 "import_statement" | "import_from_statement" => {
581 let imports = python_import_infos_from_node(statement, source);
582 bindings.extend(import_binder_from_imports(python, file, &imports).bindings);
583 }
584 "assignment" => {
585 if let Some(next) = module_replacement_from_assignment(statement, source, &bindings)
586 && replacement.replace(next).is_some()
587 {
588 return None;
589 }
590 remove_direct_scope_bindings(statement, source, &mut bindings);
591 }
592 _ => remove_direct_scope_bindings(statement, source, &mut bindings),
593 }
594 }
595 replacement
596}
597
598pub fn resolve_import_bindings(
599 python: &dyn PythonSource,
600 token: QueryToken<'_>,
601 file: &ProjectFile,
602) -> HashMap<String, CodeUnit> {
603 let imports = python.import_info_of(token, file);
604 let mut bindings = HashMap::default();
605 for resolved in resolve_imports_batched(python, file, &imports) {
606 for (binding, code_unit) in resolved {
607 bindings.insert(binding, code_unit);
608 }
609 }
610 bindings
611}
612
613pub fn resolve_imports_batched(
620 python: &dyn PythonSource,
621 file: &ProjectFile,
622 imports: &[ImportInfo],
623) -> Vec<Vec<(String, CodeUnit)>> {
624 let primary_fqns: Vec<Option<String>> = imports
625 .iter()
626 .map(|import| primary_module_fqn(file, import))
627 .collect();
628 let to_resolve: Vec<String> = primary_fqns.iter().flatten().cloned().collect();
629 let mut batch_results = resolve_module_code_units_batch(python, &to_resolve).into_iter();
630
631 imports
632 .iter()
633 .zip(primary_fqns.iter())
634 .map(|(import, primary_fqn)| {
635 let hint = primary_fqn.as_ref().map(|_| batch_results.next().unwrap());
636 resolve_import_with_hint(python, file, import, hint.as_ref())
637 })
638 .collect()
639}
640
641fn primary_module_fqn(file: &ProjectFile, import: &ImportInfo) -> Option<String> {
645 match python_import_details(import)? {
646 PythonImportDetails::Import { module, alias } => Some(python_namespace_binding_module(
647 import,
648 alias.as_deref(),
649 &module,
650 )),
651 PythonImportDetails::FromImport {
652 module,
653 name,
654 wildcard,
655 ..
656 } => {
657 if wildcard {
658 return None;
659 }
660 let resolved_module = if module.starts_with('.') {
661 resolve_python_relative_module(file, &module)
662 } else {
663 Some(module)
664 };
665 resolved_module.map(|resolved_module| format!("{resolved_module}.{name}"))
666 }
667 }
668}
669
670pub fn resolve_import(
671 python: &dyn PythonSource,
672 file: &ProjectFile,
673 import: &ImportInfo,
674) -> Vec<(String, CodeUnit)> {
675 resolve_import_with_hint(python, file, import, None)
676}
677
678fn resolve_import_with_hint(
681 python: &dyn PythonSource,
682 file: &ProjectFile,
683 import: &ImportInfo,
684 primary_hint: Option<&Option<CodeUnit>>,
685) -> Vec<(String, CodeUnit)> {
686 if let Some(details) = python_import_details(import) {
687 match details {
688 PythonImportDetails::Import { module, alias } => {
689 let binding = python_namespace_binding_name(import, alias.as_deref(), &module);
690 let bound_module =
691 python_namespace_binding_module(import, alias.as_deref(), &module);
692 let resolved = match primary_hint {
693 Some(hint) => hint.clone(),
694 None => resolve_module_code_unit(python, &bound_module),
695 };
696 if let Some(module_code_unit) = resolved {
697 return vec![(binding, module_code_unit)];
698 }
699 }
700 PythonImportDetails::FromImport {
701 module,
702 name,
703 alias,
704 wildcard,
705 } => {
706 let resolved_module = if module.starts_with('.') {
707 resolve_python_relative_module(file, &module)
708 } else {
709 Some(module)
710 };
711 let Some(resolved_module) = resolved_module else {
712 return Vec::new();
713 };
714 if wildcard {
715 return public_declarations_in_module(python, &resolved_module)
716 .into_iter()
717 .map(|code_unit| (code_unit.identifier().to_string(), code_unit))
718 .collect();
719 }
720
721 let binding = alias.clone().unwrap_or_else(|| name.clone());
722 let module_candidate = format!("{resolved_module}.{name}");
723 let resolved = match primary_hint {
724 Some(hint) => hint.clone(),
725 None => resolve_module_code_unit(python, &module_candidate),
726 };
727 if let Some(code_unit) = resolved {
728 return vec![(binding, code_unit)];
729 }
730 let exported = resolve_exported_name_from_module(python, &resolved_module, &name);
731 if !exported.is_empty() {
732 return exported
733 .into_iter()
734 .map(|code_unit| (binding.clone(), code_unit))
735 .collect();
736 }
737 let definitions: Vec<_> = python.definitions(&module_candidate).collect();
738 if !definitions.is_empty() {
739 return definitions
740 .into_iter()
741 .map(|code_unit| (binding.clone(), code_unit))
742 .collect();
743 }
744 let package_candidate: Vec<_> = python
745 .definitions(&format!("{resolved_module}.{name}"))
746 .collect();
747 if !package_candidate.is_empty() {
748 return package_candidate
749 .into_iter()
750 .map(|code_unit| (binding.clone(), code_unit))
751 .collect();
752 }
753 }
754 }
755 }
756 Vec::new()
757}
758
759pub fn resolve_exported_fqn(python: &dyn PythonSource, fqn: &str) -> Vec<CodeUnit> {
760 let Some((module, name)) = fqn.rsplit_once('.') else {
761 return Vec::new();
762 };
763 resolve_exported_name_from_module(python, module, name)
764}
765
766fn resolve_direct_named_exported_fqn(
771 python: &dyn PythonSource,
772 fqn: &str,
773) -> Option<Vec<CodeUnit>> {
774 let (module, name) = fqn.rsplit_once('.')?;
775 let mut results = Vec::new();
776 let mut queue = VecDeque::from([(module.to_string(), name.to_string())]);
777 let mut visited = HashSet::default();
778
779 while let Some((module, export_name)) = queue.pop_front() {
780 if !visited.insert((module.clone(), export_name.clone())) {
781 continue;
782 }
783 let module_unit = resolve_module_code_unit(python, &module)?;
784 let file = module_unit.source();
785 let local = local_export_declarations(python, file, &export_name);
786 let binder = python.import_binder_of(file);
787 let binding = binder.bindings.get(&export_name);
788 if !local.is_empty() && binding.is_some() {
789 return None;
790 }
791 if !local.is_empty() {
792 results.extend(local);
793 continue;
794 }
795 let binding = binding?;
796 if binding.kind != ImportKind::Named {
797 return None;
798 }
799 let imported_name = binding.imported_name.as_ref()?;
800 queue.push_back((binding.module_specifier.clone(), imported_name.clone()));
801 }
802
803 results.sort_by(|left, right| {
804 left.source()
805 .cmp(right.source())
806 .then_with(|| left.fq_name().cmp(&right.fq_name()))
807 });
808 results.dedup();
809 (!results.is_empty()).then_some(results)
810}
811
812pub fn resolve_fqn_candidates(
817 python: &dyn PythonSource,
818 fqn: &str,
819 exact: impl FnOnce(&str) -> Vec<CodeUnit>,
820) -> Vec<CodeUnit> {
821 if let Some(candidates) = resolve_direct_named_exported_fqn(python, fqn) {
822 return candidates;
823 }
824 let candidates = resolve_exported_fqn(python, fqn);
825 if !candidates.is_empty() {
826 return candidates;
827 }
828 exact(fqn)
829}
830
831fn resolve_exported_name_from_module(
832 python: &dyn PythonSource,
833 module: &str,
834 name: &str,
835) -> Vec<CodeUnit> {
836 let Some(module_unit) = resolve_module_code_unit(python, module) else {
837 return Vec::new();
838 };
839 resolve_exported_name(python, module_unit.source(), name)
840}
841
842fn resolve_exported_name(
843 python: &dyn PythonSource,
844 module_file: &ProjectFile,
845 name: &str,
846) -> Vec<CodeUnit> {
847 let mut results = Vec::new();
848 let mut queue = VecDeque::from([(module_file.clone(), name.to_string())]);
849 let mut visited = HashSet::default();
850
851 while let Some((file, export_name)) = queue.pop_front() {
852 if !visited.insert((file.clone(), export_name.clone())) {
853 continue;
854 }
855
856 let index = python.export_index_of(&file);
857 if let Some(entry) = index.exports_by_name.get(&export_name) {
858 match entry {
859 ExportEntry::Local { local_name } => {
860 results.extend(local_export_declarations(python, &file, local_name));
861 }
862 ExportEntry::ReexportedNamed {
863 module_specifier,
864 imported_name,
865 } => {
866 for target_file in
867 resolve_module_files_for_export(python, &file, module_specifier)
868 {
869 queue.push_back((target_file, imported_name.clone()));
870 }
871 }
872 ExportEntry::ReexportedModule { module_specifier } => {
873 results.extend(resolve_module_code_unit(python, module_specifier));
876 }
877 ExportEntry::Default { local_name } => {
878 if let Some(local_name) = local_name {
879 results.extend(local_export_declarations(python, &file, local_name));
880 }
881 }
882 }
883 continue;
884 }
885
886 if !export_name.starts_with('_') {
887 for star in &index.reexport_stars {
888 for target_file in
889 resolve_module_files_for_export(python, &file, &star.module_specifier)
890 {
891 queue.push_back((target_file, export_name.clone()));
892 }
893 }
894 }
895 }
896
897 results.sort_by(|left, right| {
898 left.source()
899 .cmp(right.source())
900 .then_with(|| left.fq_name().cmp(&right.fq_name()))
901 });
902 results.dedup();
903 results
904}
905
906fn local_export_declarations(
907 index: &dyn CodeUnitIndex,
908 file: &ProjectFile,
909 local_name: &str,
910) -> Vec<CodeUnit> {
911 index
912 .top_level_declarations(file)
913 .into_iter()
914 .filter(|unit| {
915 unit.identifier() == local_name
916 && index
917 .parent_of(unit)
918 .is_some_and(|parent| parent.is_module() && parent.source() == file)
919 })
920 .collect()
921}
922
923fn resolve_module_files_for_export(
924 python: &dyn PythonSource,
925 importing_file: &ProjectFile,
926 module_specifier: &str,
927) -> Vec<ProjectFile> {
928 let resolved_module = if module_specifier.starts_with('.') {
929 resolve_python_relative_module(importing_file, module_specifier)
930 } else {
931 Some(module_specifier.to_string())
932 };
933 let Some(resolved_module) = resolved_module else {
934 return Vec::new();
935 };
936 resolve_module_code_unit(python, &resolved_module)
940 .map(|unit| vec![unit.source().clone()])
941 .unwrap_or_default()
942}
943
944pub fn extract_package_from_python_wildcard(import: &ImportInfo) -> Option<String> {
945 let details = python_import_details(import)?;
946 match details {
947 PythonImportDetails::FromImport {
948 module, wildcard, ..
949 } if wildcard => Some(module),
950 _ => None,
951 }
952}
953
954#[derive(Debug, Clone)]
955pub enum PythonImportDetails {
956 Import {
957 module: String,
958 alias: Option<String>,
959 },
960 FromImport {
961 module: String,
962 name: String,
963 alias: Option<String>,
964 wildcard: bool,
965 },
966}
967
968pub fn python_import_infos_from_node(node: Node<'_>, source: &str) -> Vec<ImportInfo> {
969 match node.kind() {
970 "import_statement" => python_namespace_import_infos(node, source),
971 "import_from_statement" => python_from_import_infos(node, source),
972 _ => Vec::new(),
973 }
974}
975
976pub fn python_import_details(import: &ImportInfo) -> Option<PythonImportDetails> {
977 let path = import.path.as_ref()?;
978 match path.kind? {
979 StructuredImportPathKind::Namespace => Some(PythonImportDetails::Import {
980 module: join_python_import_segments(&path.segments),
981 alias: import.alias.clone(),
982 }),
983 StructuredImportPathKind::StaticMember => None,
985 StructuredImportPathKind::ImportFrom => {
986 let (name, module_segments) = if import.is_wildcard {
987 ("*".to_string(), path.segments.as_slice())
988 } else {
989 let (name, module_segments) = path.segments.split_last()?;
990 (name.clone(), module_segments)
991 };
992 Some(PythonImportDetails::FromImport {
993 module: join_python_import_segments(module_segments),
994 name,
995 alias: import.alias.clone(),
996 wildcard: import.is_wildcard,
997 })
998 }
999 }
1000}
1001
1002fn python_namespace_import_infos(node: Node<'_>, source: &str) -> Vec<ImportInfo> {
1003 let mut infos = Vec::new();
1004 let mut cursor = node.walk();
1005 for imported in node.children_by_field_name("name", &mut cursor) {
1006 let (module_node, alias_node) = if imported.kind() == "aliased_import" {
1007 let Some(name) = imported.child_by_field_name("name") else {
1008 continue;
1009 };
1010 (name, imported.child_by_field_name("alias"))
1011 } else {
1012 (imported, None)
1013 };
1014 let alias = alias_node
1015 .map(|alias| py_node_text(alias, source).trim().to_string())
1016 .filter(|alias| !alias.is_empty());
1017 let segments = python_path_segments(module_node, source);
1018 if segments.is_empty() {
1019 continue;
1020 }
1021 let binder_span = alias
1024 .is_some()
1025 .then_some(alias_node)
1026 .flatten()
1027 .or_else(|| python_first_segment_node(module_node))
1028 .map(brokk_bifrost_core::analyzer::common::node_span);
1029 let module = join_python_import_segments(&segments);
1030 let identifier = alias.clone().or_else(|| segments.first().cloned());
1031 infos.push(ImportInfo {
1032 raw_snippet: if let Some(alias) = &alias {
1033 format!("import {module} as {alias}")
1034 } else {
1035 format!("import {module}")
1036 },
1037 is_wildcard: false,
1038 is_global: false,
1039 identifier,
1040 alias,
1041 path: Some(StructuredImportPath {
1042 segments,
1043 kind: Some(StructuredImportPathKind::Namespace),
1044 lexical_prefixes: Vec::new(),
1045 lexical_scopes: Vec::new(),
1046 declaration_start_byte: node.start_byte(),
1047 }),
1048 binder_span,
1049 });
1050 }
1051 infos
1052}
1053
1054fn python_from_import_infos(node: Node<'_>, source: &str) -> Vec<ImportInfo> {
1055 let Some(module_node) = node.child_by_field_name("module_name") else {
1056 return Vec::new();
1057 };
1058 let module_segments = python_module_segments(module_node, source);
1059 if module_segments.is_empty() {
1060 return Vec::new();
1061 }
1062
1063 let mut infos = Vec::new();
1064 let has_wildcard_import = {
1065 let mut cursor = node.walk();
1066 node.named_children(&mut cursor)
1067 .any(|child| child.kind() == "wildcard_import")
1068 };
1069 let mut cursor = node.walk();
1070 let imported_names: Vec<_> = node.children_by_field_name("name", &mut cursor).collect();
1071 if has_wildcard_import {
1072 let module = join_python_import_segments(&module_segments);
1073 infos.push(ImportInfo {
1074 raw_snippet: format!("from {module} import *"),
1075 is_wildcard: true,
1076 is_global: false,
1077 identifier: None,
1078 alias: None,
1079 path: Some(StructuredImportPath {
1080 segments: module_segments,
1081 kind: Some(StructuredImportPathKind::ImportFrom),
1082 lexical_prefixes: Vec::new(),
1083 lexical_scopes: Vec::new(),
1084 declaration_start_byte: node.start_byte(),
1085 }),
1086 binder_span: None,
1087 });
1088 return infos;
1089 }
1090 if imported_names.is_empty() {
1091 return infos;
1092 }
1093
1094 for imported in imported_names {
1095 let (name_node, alias_node) = if imported.kind() == "aliased_import" {
1096 let Some(name) = imported.child_by_field_name("name") else {
1097 continue;
1098 };
1099 (name, imported.child_by_field_name("alias"))
1100 } else {
1101 (imported, None)
1102 };
1103 let alias = alias_node
1104 .map(|alias| py_node_text(alias, source).trim().to_string())
1105 .filter(|alias| !alias.is_empty());
1106 let name_segments = python_path_segments(name_node, source);
1107 if name_segments.is_empty() {
1108 continue;
1109 }
1110 let binder_span = alias
1113 .is_some()
1114 .then_some(alias_node)
1115 .flatten()
1116 .or_else(|| {
1117 (name_segments.len() == 1)
1118 .then(|| python_first_segment_node(name_node))
1119 .flatten()
1120 })
1121 .map(brokk_bifrost_core::analyzer::common::node_span);
1122 let imported_name = join_python_import_segments(&name_segments);
1123 let mut segments = module_segments.clone();
1124 segments.extend(name_segments);
1125 let module = join_python_import_segments(&module_segments);
1126 infos.push(ImportInfo {
1127 raw_snippet: if let Some(alias) = &alias {
1128 format!("from {module} import {imported_name} as {alias}")
1129 } else {
1130 format!("from {module} import {imported_name}")
1131 },
1132 is_wildcard: false,
1133 is_global: false,
1134 identifier: Some(alias.clone().unwrap_or_else(|| imported_name.clone())),
1135 alias,
1136 path: Some(StructuredImportPath {
1137 segments,
1138 kind: Some(StructuredImportPathKind::ImportFrom),
1139 lexical_prefixes: Vec::new(),
1140 lexical_scopes: Vec::new(),
1141 declaration_start_byte: node.start_byte(),
1142 }),
1143 binder_span,
1144 });
1145 }
1146 infos
1147}
1148
1149fn python_module_segments(module: Node<'_>, source: &str) -> Vec<String> {
1150 if module.kind() == "relative_import" {
1151 let mut cursor = module.walk();
1152 let mut prefix = String::new();
1153 let mut path_node = None;
1154 for child in module.named_children(&mut cursor) {
1155 match child.kind() {
1156 "import_prefix" if prefix.is_empty() => {
1157 prefix = py_node_text(child, source).trim().to_string();
1158 }
1159 "dotted_name" if path_node.is_none() => {
1160 path_node = Some(child);
1161 }
1162 _ => {}
1163 }
1164 }
1165 let mut segments = path_node
1166 .map(|path| python_path_segments(path, source))
1167 .unwrap_or_default();
1168 if !prefix.is_empty() {
1169 if let Some(first) = segments.first_mut() {
1170 first.insert_str(0, &prefix);
1171 } else {
1172 segments.push(prefix);
1173 }
1174 }
1175 return segments;
1176 }
1177 python_path_segments(module, source)
1178}
1179
1180fn python_first_segment_node(node: Node<'_>) -> Option<Node<'_>> {
1184 match node.kind() {
1185 "identifier" => Some(node),
1186 "dotted_name" => {
1187 let mut cursor = node.walk();
1188 node.named_children(&mut cursor)
1189 .find(|child| child.kind() == "identifier")
1190 }
1191 _ => None,
1192 }
1193}
1194
1195fn python_path_segments(node: Node<'_>, source: &str) -> Vec<String> {
1196 match node.kind() {
1197 "identifier" => vec![py_node_text(node, source).trim().to_string()],
1198 "dotted_name" => {
1199 let mut segments = Vec::new();
1200 let mut cursor = node.walk();
1201 for child in node.named_children(&mut cursor) {
1202 segments.extend(python_path_segments(child, source));
1203 }
1204 segments
1205 }
1206 _ => {
1207 let mut segments = Vec::new();
1208 let mut cursor = node.walk();
1209 for child in node.named_children(&mut cursor) {
1210 segments.extend(python_path_segments(child, source));
1211 }
1212 segments
1213 }
1214 }
1215}
1216
1217fn join_python_import_segments(segments: &[String]) -> String {
1218 let Some((first, rest)) = segments.split_first() else {
1219 return String::new();
1220 };
1221 if first.starts_with('.') && !rest.is_empty() {
1222 format!("{first}.{}", rest.join("."))
1223 } else {
1224 segments.join(".")
1225 }
1226}
1227
1228pub fn python_namespace_binding_name(
1229 import: &ImportInfo,
1230 alias: Option<&str>,
1231 module: &str,
1232) -> String {
1233 import
1234 .identifier
1235 .clone()
1236 .or_else(|| alias.map(str::to_string))
1237 .unwrap_or_else(|| module.to_string())
1238}
1239
1240pub fn python_namespace_binding_module(
1241 import: &ImportInfo,
1242 alias: Option<&str>,
1243 module: &str,
1244) -> String {
1245 if alias.is_some() {
1246 return module.to_string();
1247 }
1248 import
1249 .path
1250 .as_ref()
1251 .and_then(|path| path.segments.first().cloned())
1252 .unwrap_or_else(|| module.to_string())
1253}
1254
1255pub fn resolve_python_relative_module(
1256 source_file: &ProjectFile,
1257 module_expr: &str,
1258) -> Option<String> {
1259 resolve_python_relative_module_from_package(&python_current_package(source_file), module_expr)
1260}
1261
1262pub fn resolve_python_relative_module_from_package(
1266 current_package: &str,
1267 module_expr: &str,
1268) -> Option<String> {
1269 let level = module_expr.chars().take_while(|ch| *ch == '.').count();
1270 let suffix = module_expr[level..].trim_matches('.');
1271 let mut parts: Vec<_> = current_package
1272 .split('.')
1273 .filter(|part| !part.is_empty())
1274 .map(str::to_string)
1275 .collect();
1276 if level == 0 {
1277 return Some(module_expr.to_string());
1278 }
1279 if level > 0 {
1280 if level - 1 > parts.len() {
1281 return None;
1282 }
1283 parts.truncate(parts.len() - (level - 1));
1284 }
1285 if !suffix.is_empty() {
1286 parts.extend(suffix.split('.').map(str::to_string));
1287 }
1288 Some(parts.join("."))
1289}
1290
1291fn python_current_package(source_file: &ProjectFile) -> String {
1292 let module = python_module_name(source_file);
1293 if source_file
1294 .rel_path()
1295 .file_name()
1296 .and_then(|name| name.to_str())
1297 == Some("__init__.py")
1298 {
1299 module
1300 } else {
1301 module
1302 .rsplit_once('.')
1303 .map(|(package, _)| package.to_string())
1304 .unwrap_or_default()
1305 }
1306}