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 crate::runtime_values::enclosing_execution_context;
14use crate::syntax::python_static_attribute_path;
15use brokk_bifrost_core::analyzer::common::node_source_text;
16use brokk_bifrost_core::analyzer::model::{
17 ImportInfo, StructuredImportPath, StructuredImportPathKind,
18};
19use brokk_bifrost_core::analyzer::query_token::QueryToken;
20use brokk_bifrost_core::analyzer::usages::model::{ExportEntry, ImportBinding, ImportKind};
21use brokk_bifrost_core::analyzer::{CodeUnit, CodeUnitIndex, ProjectFile};
22use brokk_bifrost_core::hash::{HashMap, HashSet};
23use std::collections::VecDeque;
24use tree_sitter::Node;
25
26#[derive(Debug, Clone, PartialEq, Eq)]
27pub struct PythonModuleReplacement {
28 pub target_module: String,
29}
30
31pub(crate) fn imported_module_assignment_at(
41 reference: Node<'_>,
42 local: &str,
43 source: &str,
44 mut is_importlib_binding: impl FnMut(&str) -> bool,
45) -> Option<String> {
46 let suite = enclosing_execution_suite(reference)?;
47 let mut resolved = None;
48 let mut cursor = suite.walk();
49 for statement in suite.named_children(&mut cursor) {
50 if statement.start_byte() >= reference.start_byte()
51 || (statement.start_byte() <= reference.start_byte()
52 && reference.end_byte() <= statement.end_byte())
53 {
54 break;
55 }
56
57 let expression = if statement.kind() == "expression_statement" {
58 statement.named_child(0).unwrap_or(statement)
59 } else {
60 statement
61 };
62 if expression.kind() == "assignment"
63 && expression.child_by_field_name("left").is_some_and(|left| {
64 left.kind() == "identifier" && node_source_text(left, source) == local
65 })
66 {
67 resolved = expression
68 .child_by_field_name("right")
69 .and_then(|right| importlib_call_module(right, source, &mut is_importlib_binding));
70 continue;
71 }
72
73 let binds_local = python_direct_scope_bindings_bounded(statement, source, || true)
74 .expect("unbounded binding collection cannot be cancelled")
75 .into_iter()
76 .any(|binding| node_source_text(binding.declaration, source) == local);
77 if binds_local {
78 resolved = None;
79 }
80 }
81 resolved
82}
83
84pub fn python_local_module_alias_path_before<E, C>(
92 reference: Node<'_>,
93 local_name: &str,
94 source: &str,
95 context: &mut C,
96 mut imported_module_path: impl FnMut(Node<'_>, &mut C) -> Result<Option<Box<[Box<str>]>>, E>,
97 mut charge_step: impl FnMut(&mut C) -> Result<(), E>,
98) -> Result<Option<Box<[Box<str>]>>, E> {
99 let mut current = reference;
100 let suite = loop {
101 let Some(parent) = current.parent() else {
102 return Ok(None);
103 };
104 match parent.kind() {
105 "function_definition" | "lambda" => break parent.child_by_field_name("body"),
106 "module" => break Some(parent),
107 "class_definition" => return Ok(None),
108 _ => current = parent,
109 }
110 };
111 let Some(suite) = suite else {
112 return Ok(None);
113 };
114
115 let mut aliases: HashMap<Box<str>, Box<[Box<str>]>> = HashMap::default();
116 let mut cursor = suite.walk();
117 for statement in suite.named_children(&mut cursor) {
118 if statement.start_byte() >= reference.start_byte()
119 || (statement.start_byte() <= reference.start_byte()
120 && reference.end_byte() <= statement.end_byte())
121 {
122 break;
123 }
124 charge_step(context)?;
125 let expression = if statement.kind() == "expression_statement" {
126 statement.named_child(0).unwrap_or(statement)
127 } else {
128 statement
129 };
130 if expression.kind() == "assignment" {
131 let (targets, value) = python_assignment_chain(expression);
132 if !targets.is_empty() && targets.iter().all(|target| target.kind() == "identifier") {
133 let module_path = value
134 .map(|value| {
135 python_module_alias_expression_path(
136 value,
137 source,
138 &aliases,
139 context,
140 &mut imported_module_path,
141 )
142 })
143 .transpose()?
144 .flatten();
145 for target in targets {
146 let name = node_source_text(target, source);
147 if let Some(module_path) = module_path.as_ref() {
148 aliases.insert(name.into(), module_path.clone());
149 } else {
150 aliases.remove(name);
151 }
152 }
153 continue;
154 }
155 }
156
157 if !matches!(
158 expression.kind(),
159 "assignment"
160 | "augmented_assignment"
161 | "delete_statement"
162 | "for_statement"
163 | "async_for_statement"
164 | "while_statement"
165 | "with_statement"
166 | "try_statement"
167 | "match_statement"
168 | "import_statement"
169 | "import_from_statement"
170 | "function_definition"
171 | "class_definition"
172 | "global_statement"
173 | "nonlocal_statement"
174 | "named_expression"
175 | "if_statement"
176 ) {
177 continue;
178 }
179 let mut binding_error = None;
180 let bindings =
181 python_direct_scope_bindings_bounded(statement, source, || {
182 match charge_step(context) {
183 Ok(()) => true,
184 Err(error) => {
185 binding_error = Some(error);
186 false
187 }
188 }
189 });
190 if let Some(error) = binding_error {
191 return Err(error);
192 }
193 let bindings = bindings.expect("bounded Python alias scan stopped without a cause");
194 for binding in bindings {
195 aliases.remove(node_source_text(binding.declaration, source));
196 }
197 }
198 Ok(aliases.get(local_name).cloned())
199}
200
201fn python_module_alias_expression_path<E, C>(
202 expression: Node<'_>,
203 source: &str,
204 aliases: &HashMap<Box<str>, Box<[Box<str>]>>,
205 context: &mut C,
206 imported_module_path: &mut impl FnMut(Node<'_>, &mut C) -> Result<Option<Box<[Box<str>]>>, E>,
207) -> Result<Option<Box<[Box<str>]>>, E> {
208 if let Some(path) = python_static_attribute_path(expression)
209 && let Some(name) = path
210 .first()
211 .map(|segment| node_source_text(*segment, source))
212 && let Some(module_path) = aliases.get(name)
213 {
214 let mut canonical = module_path.to_vec();
215 canonical.extend(path.iter().skip(1).map(|segment| {
216 node_source_text(*segment, source)
217 .to_owned()
218 .into_boxed_str()
219 }));
220 return Ok(Some(canonical.into_boxed_slice()));
221 }
222 imported_module_path(expression, context)
223}
224
225pub fn python_assignment_chain<'tree>(
229 node: Node<'tree>,
230) -> (Vec<Node<'tree>>, Option<Node<'tree>>) {
231 let mut targets = Vec::new();
232 let mut assignment = node;
233 loop {
234 let Some(target) = assignment.child_by_field_name("left") else {
235 return (targets, None);
236 };
237 targets.push(target);
238 let Some(value) = assignment.child_by_field_name("right") else {
239 return (targets, None);
240 };
241 if value.kind() != "assignment" {
242 return (targets, Some(value));
243 }
244 assignment = value;
245 }
246}
247
248fn enclosing_execution_suite(node: Node<'_>) -> Option<Node<'_>> {
249 if let Some(function) = enclosing_execution_context(node) {
250 return function.child_by_field_name("body");
251 }
252
253 let mut current = node;
254 while let Some(parent) = current.parent() {
255 if parent.kind() == "class_definition" {
256 return parent.child_by_field_name("body");
257 }
258 if parent.kind() == "module" {
259 return Some(parent);
260 }
261 current = parent;
262 }
263 None
264}
265
266fn importlib_call_module(
267 call: Node<'_>,
268 source: &str,
269 is_importlib_binding: &mut impl FnMut(&str) -> bool,
270) -> Option<String> {
271 if call.kind() != "call" {
272 return None;
273 }
274 let function = call.child_by_field_name("function")?;
275 if function.kind() != "attribute" {
276 return None;
277 }
278 let (object, attribute) = (
279 function.child_by_field_name("object")?,
280 function.child_by_field_name("attribute")?,
281 );
282 if object.kind() != "identifier"
283 || attribute.kind() != "identifier"
284 || node_source_text(attribute, source) != "import_module"
285 || !is_importlib_binding(node_source_text(object, source))
286 {
287 return None;
288 }
289
290 let arguments = call.child_by_field_name("arguments")?;
291 if arguments.named_child_count() != 1 {
292 return None;
293 }
294 let literal = arguments.named_child(0)?;
295 if literal.kind() != "string"
296 || literal
297 .parent()
298 .is_some_and(|parent| parent.kind() == "concatenated_string")
299 {
300 return None;
301 }
302 let mut content = None;
303 for index in 0..literal.named_child_count() {
304 let child = literal.named_child(index)?;
305 match child.kind() {
306 "string_start" | "string_end" => {}
307 "string_content" if content.is_none() && child.named_child_count() == 0 => {
308 content = Some(child);
309 }
310 _ => return None,
311 }
312 }
313 let content = content?;
314 let module = node_source_text(content, source);
315 (!module.is_empty() && !module.starts_with('.')).then(|| module.to_string())
316}
317
318pub(crate) fn literal_importlib_modules(
323 source: &str,
324 bindings: &HashMap<String, ImportBinding>,
325) -> HashSet<String> {
326 let Some(tree) = parse_python_tree(source) else {
327 return HashSet::default();
328 };
329 let mut modules = HashSet::default();
330 let mut stack = vec![tree.root_node()];
331 while let Some(node) = stack.pop() {
332 if node.kind() == "assignment"
333 && node
334 .child_by_field_name("left")
335 .is_some_and(|left| left.kind() == "identifier")
336 && let Some(module) = node.child_by_field_name("right").and_then(|right| {
337 importlib_call_module(right, source, &mut |local| {
338 bindings.get(local).is_some_and(|binding| {
339 binding.kind == ImportKind::Namespace
340 && binding
341 .namespace_imported_module
342 .as_deref()
343 .unwrap_or(&binding.module_specifier)
344 == "importlib"
345 })
346 })
347 })
348 {
349 modules.insert(module);
350 }
351 let mut cursor = node.walk();
352 stack.extend(node.named_children(&mut cursor));
353 }
354 modules
355}
356
357fn module_replacement_from_assignment(
365 assignment: Node<'_>,
366 source: &str,
367 bindings: &HashMap<String, ImportBinding>,
368) -> Option<PythonModuleReplacement> {
369 let (left, right) = (
370 assignment.child_by_field_name("left")?,
371 assignment.child_by_field_name("right")?,
372 );
373 if left.kind() != "subscript" || right.kind() != "identifier" {
374 return None;
375 }
376 let (value, subscript) = (
377 left.child_by_field_name("value")?,
378 left.child_by_field_name("subscript")?,
379 );
380 if value.kind() != "attribute"
381 || subscript.kind() != "identifier"
382 || node_source_text(subscript, source) != "__name__"
383 {
384 return None;
385 }
386 let (sys_local, modules) = (
387 value.child_by_field_name("object")?,
388 value.child_by_field_name("attribute")?,
389 );
390 if sys_local.kind() != "identifier"
391 || modules.kind() != "identifier"
392 || node_source_text(modules, source) != "modules"
393 {
394 return None;
395 }
396 let sys_binding = bindings.get(node_source_text(sys_local, source))?;
397 let sys_module = sys_binding
398 .namespace_imported_module
399 .as_deref()
400 .unwrap_or(&sys_binding.module_specifier);
401 if sys_binding.kind != ImportKind::Namespace || sys_module != "sys" {
402 return None;
403 }
404
405 let target_binding = bindings.get(node_source_text(right, source))?;
406 if target_binding.kind != ImportKind::Namespace {
407 return None;
408 }
409 Some(PythonModuleReplacement {
410 target_module: target_binding
411 .namespace_imported_module
412 .clone()
413 .unwrap_or_else(|| target_binding.module_specifier.clone()),
414 })
415}
416
417fn remove_direct_scope_bindings(
418 statement: Node<'_>,
419 source: &str,
420 bindings: &mut HashMap<String, ImportBinding>,
421) {
422 let mut stack = vec![statement];
423 while let Some(node) = stack.pop() {
424 for binding in python_direct_scope_bindings_bounded(node, source, || true)
425 .expect("unbounded binding collection cannot be cancelled")
426 {
427 bindings.remove(node_source_text(binding.declaration, source));
428 }
429 if matches!(
430 node.kind(),
431 "function_definition" | "class_definition" | "lambda"
432 ) {
433 continue;
434 }
435 let mut cursor = node.walk();
436 stack.extend(node.named_children(&mut cursor));
437 }
438}
439
440pub fn parse_python_import_infos(source: &str) -> Vec<ImportInfo> {
444 let Some(tree) = parse_python_tree(source) else {
445 return Vec::new();
446 };
447 let mut imports = Vec::new();
448 for node in python_import_statement_nodes(tree.root_node()) {
449 imports.extend(python_import_infos_from_node(node, source));
450 }
451 imports
452}
453
454fn python_import_statement_nodes<'tree>(root: Node<'tree>) -> Vec<Node<'tree>> {
455 let mut pending = vec![root];
456 let mut nodes = Vec::new();
457 while let Some(node) = pending.pop() {
458 if matches!(node.kind(), "import_statement" | "import_from_statement") {
459 nodes.push(node);
460 continue;
461 }
462 let mut cursor = node.walk();
463 pending.extend(node.named_children(&mut cursor));
464 }
465 nodes
466}
467
468#[derive(Debug, Clone, PartialEq, Eq)]
469pub struct PythonImportBinding {
470 pub start_byte: usize,
471 pub scope_start_byte: usize,
472 pub scope_end_byte: usize,
473 function_scoped: bool,
474 pub local_name: String,
475 pub qualified_name: String,
476 pub consumed_attributes: usize,
477}
478
479impl PythonImportBinding {
480 pub fn is_function_scoped(&self) -> bool {
481 self.function_scoped
482 }
483}
484
485pub fn python_function_scoped_import_binding<'a>(
487 bindings: &'a [PythonImportBinding],
488 name: &str,
489 reference: Node<'_>,
490) -> Option<&'a PythonImportBinding> {
491 bindings.iter().rev().find(|binding| {
492 binding.is_function_scoped()
493 && binding.start_byte <= reference.start_byte()
494 && binding.scope_start_byte <= reference.start_byte()
495 && reference.end_byte() <= binding.scope_end_byte
496 && binding.local_name == name
497 })
498}
499
500pub fn parse_python_import_bindings(source: &str) -> Vec<PythonImportBinding> {
504 let Some(tree) = parse_python_tree(source) else {
505 return Vec::new();
506 };
507 python_import_bindings_from_tree(tree.root_node(), source)
508}
509
510pub fn python_import_bindings_from_tree(root: Node<'_>, source: &str) -> Vec<PythonImportBinding> {
514 let mut nodes = python_import_statement_nodes(root);
515 nodes.sort_by_key(Node::start_byte);
516 nodes
517 .into_iter()
518 .flat_map(|node| {
519 let (scope_start_byte, scope_end_byte, function_scoped) =
520 python_import_binding_scope(node, source.len());
521 python_import_infos_from_node(node, source)
522 .into_iter()
523 .filter_map(move |import| {
524 let path = import.path.as_ref()?;
525 let details = python_import_details(&import)?;
526 match details {
527 PythonImportDetails::Import { module, alias } => {
528 let consumed_attributes = if alias.is_some() {
529 0
530 } else {
531 path.segments.len().saturating_sub(1)
532 };
533 Some(PythonImportBinding {
534 start_byte: node.start_byte(),
535 scope_start_byte,
536 scope_end_byte,
537 function_scoped,
538 local_name: import.identifier.clone()?,
539 qualified_name: module,
540 consumed_attributes,
541 })
542 }
543 PythonImportDetails::FromImport {
544 module,
545 name,
546 alias: _,
547 wildcard: false,
548 } => Some(PythonImportBinding {
549 start_byte: node.start_byte(),
550 scope_start_byte,
551 scope_end_byte,
552 function_scoped,
553 local_name: import.identifier.clone()?,
554 qualified_name: format!("{module}.{name}"),
555 consumed_attributes: 0,
556 }),
557 PythonImportDetails::FromImport { wildcard: true, .. } => None,
558 }
559 })
560 })
561 .collect()
562}
563
564pub fn python_import_bindings_from_imports(
571 imports: &[ImportInfo],
572 root: Node<'_>,
573 source_len: usize,
574) -> Vec<PythonImportBinding> {
575 python_import_bindings_from_imports_bounded(imports, root, source_len, || true)
576 .expect("unbounded import projection completes")
577}
578
579pub fn python_import_bindings_from_imports_bounded(
582 imports: &[ImportInfo],
583 root: Node<'_>,
584 source_len: usize,
585 mut scope_step: impl FnMut() -> bool,
586) -> Option<Vec<PythonImportBinding>> {
587 let mut bindings = Vec::new();
588 for import in imports {
589 if !scope_step() {
590 return None;
591 }
592 let Some(path) = import.path.as_ref() else {
593 continue;
594 };
595 for _segment in &path.segments {
596 if !scope_step() {
597 return None;
598 }
599 }
600 let Some(details) = python_import_details(import) else {
601 continue;
602 };
603 let mut binding_scope = None;
604 let mut valid = true;
605 for scope in path.lexical_scopes.iter().rev() {
606 if !scope_step() {
607 return None;
608 }
609 let Some(node) = root.descendant_for_byte_range(scope.start_byte, scope.end_byte)
610 else {
611 valid = false;
612 break;
613 };
614 if node.start_byte() != scope.start_byte || node.end_byte() != scope.end_byte {
615 valid = false;
616 break;
617 }
618 if matches!(node.kind(), "function_definition" | "lambda") {
619 binding_scope = Some((scope.start_byte, scope.end_byte, true));
620 break;
621 }
622 }
623 if !valid {
624 continue;
625 }
626 let (scope_start_byte, scope_end_byte, function_scoped) =
627 binding_scope.unwrap_or((0, source_len, false));
628 let Some(local_name) = import.identifier.clone() else {
629 continue;
630 };
631 let (qualified_name, consumed_attributes) = match details {
632 PythonImportDetails::Import { module, alias } => (
633 module,
634 if alias.is_some() {
635 0
636 } else {
637 path.segments.len().saturating_sub(1)
638 },
639 ),
640 PythonImportDetails::FromImport {
641 module,
642 name,
643 wildcard: false,
644 ..
645 } => (format!("{module}.{name}"), 0),
646 PythonImportDetails::FromImport { wildcard: true, .. } => continue,
647 };
648 bindings.push(PythonImportBinding {
649 start_byte: path.declaration_start_byte,
650 scope_start_byte,
651 scope_end_byte,
652 function_scoped,
653 local_name,
654 qualified_name,
655 consumed_attributes,
656 });
657 }
658 bindings.sort_by_key(|binding| binding.start_byte);
659 Some(bindings)
660}
661
662fn python_import_binding_scope(node: Node<'_>, source_len: usize) -> (usize, usize, bool) {
663 let mut parent = node.parent();
664 while let Some(scope) = parent {
665 if matches!(scope.kind(), "function_definition" | "lambda") {
666 return (scope.start_byte(), scope.end_byte(), true);
667 }
668 parent = scope.parent();
669 }
670 (0, source_len, false)
671}
672
673#[cfg(test)]
674mod tests {
675 use super::*;
676 use brokk_bifrost_core::analyzer::usages::model::ImportBinder;
677 use brokk_bifrost_test_macros::hardwon;
678
679 fn importlib_assignment_for(source: &str, importlib_local: &str) -> Option<String> {
680 let tree = parse_python_tree(source).expect("valid Python fixture");
681 let reference_start = source.rfind("canonical.setup").expect("reference site");
682 let mut stack = vec![tree.root_node()];
683 let mut reference = None;
684 while let Some(node) = stack.pop() {
685 if node.kind() == "identifier"
686 && node.start_byte() == reference_start
687 && node_source_text(node, source) == "canonical"
688 {
689 reference = Some(node);
690 break;
691 }
692 let mut cursor = node.walk();
693 stack.extend(node.named_children(&mut cursor));
694 }
695 imported_module_assignment_at(
696 reference.expect("canonical reference node"),
697 "canonical",
698 source,
699 |local| local == importlib_local,
700 )
701 }
702
703 #[test]
704 fn importlib_assignment_requires_a_direct_literal_dominating_write() {
705 let valid = r#"import importlib as loader
706
707def test():
708 canonical = loader.import_module("routes.contacts.contacts_routes")
709 assert canonical.setup_contacts_routes
710"#;
711 assert_eq!(
712 importlib_assignment_for(valid, "loader"),
713 Some("routes.contacts.contacts_routes".to_string())
714 );
715
716 for near_miss in [
717 r#"import importlib as loader
718def test(module_name):
719 canonical = loader.import_module(module_name)
720 assert canonical.setup_contacts_routes
721"#,
722 r#"import importlib as loader
723def test():
724 if enabled:
725 canonical = loader.import_module("routes.contacts.contacts_routes")
726 assert canonical.setup_contacts_routes
727"#,
728 r#"import importlib as loader
729def test():
730 canonical = loader.import_module("routes.contacts.contacts_routes")
731 canonical = object()
732 assert canonical.setup_contacts_routes
733"#,
734 ] {
735 assert_eq!(
736 importlib_assignment_for(near_miss, "loader"),
737 None,
738 "near miss must not infer a module receiver: {near_miss:?}"
739 );
740 }
741 assert_eq!(
742 importlib_assignment_for(valid, "shadowed_loader"),
743 None,
744 "a similarly shaped non-importlib receiver must not resolve"
745 );
746 }
747
748 fn replacement_for(source: &str, binder: &ImportBinder) -> Option<PythonModuleReplacement> {
749 let tree = parse_python_tree(source).expect("valid Python fixture");
750 let root = tree.root_node();
751 let mut replacement = None;
752 let mut cursor = root.walk();
753 for statement in root.named_children(&mut cursor) {
754 let statement = if statement.kind() == "expression_statement" {
755 statement.named_child(0).expect("fixture expression")
756 } else {
757 statement
758 };
759 if statement.kind() == "assignment"
760 && let Some(next) =
761 module_replacement_from_assignment(statement, source, &binder.bindings)
762 && replacement.replace(next).is_some()
763 {
764 return None;
765 }
766 }
767 replacement
768 }
769
770 #[test]
771 #[hardwon("59dfb80c3bb8: module replacement accepted malformed or shadowed import shapes")]
772 fn module_replacement_requires_exact_structured_import_bindings() {
773 let source = r#"import sys as _sys
774from routes.contacts import contacts_routes as _canonical
775
776_sys.modules[__name__] = _canonical
777"#;
778 let mut binder = ImportBinder::empty();
779 binder.bindings.insert(
780 "_sys".to_string(),
781 ImportBinding {
782 module_specifier: "sys".to_string(),
783 namespace_imported_module: Some("sys".to_string()),
784 kind: ImportKind::Namespace,
785 imported_name: None,
786 },
787 );
788 binder.bindings.insert(
789 "_canonical".to_string(),
790 ImportBinding {
791 module_specifier: "routes.contacts.contacts_routes".to_string(),
792 namespace_imported_module: None,
793 kind: ImportKind::Namespace,
794 imported_name: None,
795 },
796 );
797
798 assert_eq!(
799 replacement_for(source, &binder),
800 Some(PythonModuleReplacement {
801 target_module: "routes.contacts.contacts_routes".to_string(),
802 })
803 );
804
805 for near_miss in [
806 "cache.modules[__name__] = _canonical\n",
807 "_sys.modules[module_name] = _canonical\n",
808 "_sys.modules[__name__] = build()\n",
809 "def replace():\n _sys.modules[__name__] = _canonical\n",
810 ] {
811 assert_eq!(
812 replacement_for(near_miss, &binder),
813 None,
814 "near miss must not replace module identity: {near_miss:?}"
815 );
816 }
817 }
818
819 #[test]
820 #[hardwon(
821 "b58235b247fd: Python import bindings lost function scope or relative module identity"
822 )]
823 fn existing_tree_import_bindings_preserve_function_scope_and_relative_identity() {
824 let source = r#"def register():
825 from .workers.special import Worker
826 return Worker()
827"#;
828 let tree = parse_python_tree(source).expect("Python tree");
829 let bindings = python_import_bindings_from_tree(tree.root_node(), source);
830
831 assert_eq!(bindings.len(), 1, "{bindings:#?}");
832 assert!(bindings[0].is_function_scoped(), "{bindings:#?}");
833 assert_eq!(bindings[0].local_name, "Worker");
834 assert_eq!(bindings[0].qualified_name, ".workers.special.Worker");
835 assert!(bindings[0].scope_start_byte <= bindings[0].start_byte);
836 assert!(bindings[0].scope_end_byte >= source.rfind("Worker()").unwrap());
837 }
838}
839
840pub(crate) fn python_imports_by_statement(imports: &[ImportInfo]) -> HashMap<usize, &[ImportInfo]> {
843 let mut statement_imports = HashMap::default();
844 for group in imports.chunk_by(|left, right| {
845 left.path.as_ref().map(|path| path.declaration_start_byte)
846 == right.path.as_ref().map(|path| path.declaration_start_byte)
847 }) {
848 if let Some(path) = &group[0].path {
849 assert!(
850 statement_imports
851 .insert(path.declaration_start_byte, group)
852 .is_none(),
853 "canonical Python import leaves must remain in statement order"
854 );
855 }
856 }
857 statement_imports
858}
859
860pub fn module_replacement_of(
861 python: &dyn PythonSource,
862 file: &ProjectFile,
863 source: &str,
864 imports: &[ImportInfo],
865) -> Option<PythonModuleReplacement> {
866 let tree = parse_python_tree(source)?;
867 let root = tree.root_node();
868 let mut bindings: HashMap<String, ImportBinding> = HashMap::default();
869 let statement_imports = python_imports_by_statement(imports);
870 let mut replacement = None;
871 let mut cursor = root.walk();
872 for statement in root.named_children(&mut cursor) {
873 let statement = if statement.kind() == "expression_statement" {
874 let Some(expression) = statement.named_child(0) else {
875 continue;
876 };
877 expression
878 } else {
879 statement
880 };
881 match statement.kind() {
882 "import_statement" | "import_from_statement" => {
883 if let Some(imports) = statement_imports.get(&statement.start_byte()) {
884 bindings.extend(import_binder_from_imports(python, file, imports).bindings);
885 }
886 }
887 "assignment" => {
888 if let Some(next) = module_replacement_from_assignment(statement, source, &bindings)
889 && replacement.replace(next).is_some()
890 {
891 return None;
892 }
893 remove_direct_scope_bindings(statement, source, &mut bindings);
894 }
895 _ => remove_direct_scope_bindings(statement, source, &mut bindings),
896 }
897 }
898 replacement
899}
900
901pub fn resolve_import_bindings(
902 python: &dyn PythonSource,
903 token: QueryToken<'_>,
904 file: &ProjectFile,
905) -> HashMap<String, CodeUnit> {
906 let imports = python.import_info_of(token, file);
907 let mut bindings = HashMap::default();
908 for resolved in resolve_imports_batched(python, file, &imports) {
909 for (binding, code_unit) in resolved {
910 bindings.insert(binding, code_unit);
911 }
912 }
913 bindings
914}
915
916pub fn resolve_imports_batched(
923 python: &dyn PythonSource,
924 file: &ProjectFile,
925 imports: &[ImportInfo],
926) -> Vec<Vec<(String, CodeUnit)>> {
927 let primary_fqns: Vec<Option<String>> = imports
928 .iter()
929 .map(|import| primary_module_fqn(python, file, import))
930 .collect();
931 let to_resolve: Vec<String> = primary_fqns.iter().flatten().cloned().collect();
932 let mut batch_results =
933 resolve_module_code_units_batch(python, Some(file), &to_resolve).into_iter();
934
935 imports
936 .iter()
937 .zip(primary_fqns.iter())
938 .map(|(import, primary_fqn)| {
939 let hint = primary_fqn.as_ref().map(|_| batch_results.next().unwrap());
940 resolve_import_with_hint(python, file, import, hint.as_ref())
941 })
942 .collect()
943}
944
945pub fn resolve_import_files_batched(
955 python: &dyn PythonSource,
956 file: &ProjectFile,
957 imports: &[ImportInfo],
958) -> HashSet<ProjectFile> {
959 let mut module_fqns = Vec::new();
960 for import in imports {
961 match python_import_details(import) {
962 Some(PythonImportDetails::Import { module, .. }) => module_fqns.push(module),
963 Some(PythonImportDetails::FromImport {
964 module,
965 name,
966 wildcard,
967 ..
968 }) => {
969 let resolved_module = if module.starts_with('.') {
970 resolve_python_relative_module(python.project(), file, &module)
971 } else {
972 Some(module)
973 };
974 let Some(resolved_module) = resolved_module else {
975 continue;
976 };
977 if !wildcard {
978 module_fqns.push(format!("{resolved_module}.{name}"));
979 }
980 module_fqns.push(resolved_module);
981 }
982 None => {}
983 }
984 }
985 module_fqns.sort();
986 module_fqns.dedup();
987
988 python
989 .path_module_fqns_batch(&module_fqns)
990 .into_iter()
991 .flatten()
992 .flatten()
993 .filter(CodeUnit::is_module)
994 .map(|unit| unit.source().clone())
995 .collect()
996}
997
998fn primary_module_fqn(
1002 python: &dyn PythonSource,
1003 file: &ProjectFile,
1004 import: &ImportInfo,
1005) -> Option<String> {
1006 match python_import_details(import)? {
1007 PythonImportDetails::Import { module, alias } => Some(python_namespace_binding_module(
1008 import,
1009 alias.as_deref(),
1010 &module,
1011 )),
1012 PythonImportDetails::FromImport {
1013 module,
1014 name,
1015 wildcard,
1016 ..
1017 } => {
1018 if wildcard {
1019 return None;
1020 }
1021 let resolved_module = resolve_python_module_specifier(python.project(), file, &module);
1022 resolved_module.map(|resolved_module| format!("{resolved_module}.{name}"))
1023 }
1024 }
1025}
1026
1027pub fn resolve_import(
1028 python: &dyn PythonSource,
1029 file: &ProjectFile,
1030 import: &ImportInfo,
1031) -> Vec<(String, CodeUnit)> {
1032 resolve_import_with_hint(python, file, import, None)
1033}
1034
1035fn resolve_import_with_hint(
1038 python: &dyn PythonSource,
1039 file: &ProjectFile,
1040 import: &ImportInfo,
1041 primary_hint: Option<&Option<CodeUnit>>,
1042) -> Vec<(String, CodeUnit)> {
1043 if let Some(details) = python_import_details(import) {
1044 match details {
1045 PythonImportDetails::Import { module, alias } => {
1046 let binding = python_namespace_binding_name(import, alias.as_deref(), &module);
1047 let bound_module =
1048 python_namespace_binding_module(import, alias.as_deref(), &module);
1049 let resolved = match primary_hint {
1050 Some(hint) => hint.clone(),
1051 None => resolve_module_code_unit(python, Some(file), &bound_module),
1052 };
1053 if let Some(module_code_unit) = resolved {
1054 return vec![(binding, module_code_unit)];
1055 }
1056 }
1057 PythonImportDetails::FromImport {
1058 module,
1059 name,
1060 alias,
1061 wildcard,
1062 } => {
1063 let resolved_module =
1064 resolve_python_module_specifier(python.project(), file, &module);
1065 let Some(resolved_module) = resolved_module else {
1066 return Vec::new();
1067 };
1068 if wildcard {
1069 return public_declarations_in_module(python, file, &resolved_module)
1070 .into_iter()
1071 .map(|code_unit| (code_unit.identifier().to_string(), code_unit))
1072 .collect();
1073 }
1074
1075 let binding = alias.clone().unwrap_or_else(|| name.clone());
1076 let module_candidate = format!("{resolved_module}.{name}");
1077 let exported =
1078 resolve_exported_name_from_module(python, Some(file), &resolved_module, &name);
1079 if !exported.is_empty() {
1080 return exported
1081 .into_iter()
1082 .map(|code_unit| (binding.clone(), code_unit))
1083 .collect();
1084 }
1085 let resolved = match primary_hint {
1086 Some(hint) => hint.clone(),
1087 None => resolve_module_code_unit(python, Some(file), &module_candidate),
1088 };
1089 if let Some(code_unit) = resolved {
1090 return vec![(binding, code_unit)];
1091 }
1092 let definitions: Vec<_> = python.definitions(&module_candidate).collect();
1093 if !definitions.is_empty() {
1094 return definitions
1095 .into_iter()
1096 .map(|code_unit| (binding.clone(), code_unit))
1097 .collect();
1098 }
1099 let package_candidate: Vec<_> = python
1100 .definitions(&format!("{resolved_module}.{name}"))
1101 .collect();
1102 if !package_candidate.is_empty() {
1103 return package_candidate
1104 .into_iter()
1105 .map(|code_unit| (binding.clone(), code_unit))
1106 .collect();
1107 }
1108 }
1109 }
1110 }
1111 Vec::new()
1112}
1113
1114pub fn resolve_exported_fqn(
1118 python: &dyn PythonSource,
1119 importer: Option<&ProjectFile>,
1120 fqn: &str,
1121) -> Vec<CodeUnit> {
1122 let Some((module, name)) = fqn.rsplit_once('.') else {
1123 return Vec::new();
1124 };
1125 resolve_exported_name_from_module(python, importer, module, name)
1126}
1127
1128fn resolve_direct_named_exported_fqn(
1137 python: &dyn PythonSource,
1138 importer: Option<&ProjectFile>,
1139 fqn: &str,
1140) -> Option<Vec<CodeUnit>> {
1141 let (module, name) = fqn.rsplit_once('.')?;
1142 let mut results = Vec::new();
1143 let mut queue = VecDeque::from([(importer.cloned(), module.to_string(), name.to_string())]);
1144 let mut visited = HashSet::default();
1145
1146 while let Some((importer, module, export_name)) = queue.pop_front() {
1147 if !visited.insert((module.clone(), export_name.clone())) {
1148 continue;
1149 }
1150 let module_unit = resolve_module_code_unit(python, importer.as_ref(), &module)?;
1151 let file = module_unit.source();
1152 let local = local_export_declarations(python, file, &export_name);
1153 let binder = python.import_binder_of(file);
1154 let binding = binder.bindings.get(&export_name);
1155 if !local.is_empty() && binding.is_some() {
1156 return None;
1157 }
1158 if !local.is_empty() {
1159 results.extend(local);
1160 continue;
1161 }
1162 let binding = binding?;
1163 if binding.kind != ImportKind::Named {
1164 return None;
1165 }
1166 let imported_name = binding.imported_name.as_ref()?;
1167 queue.push_back((
1168 Some(file.clone()),
1169 binding.module_specifier.clone(),
1170 imported_name.clone(),
1171 ));
1172 }
1173
1174 results.sort_by(|left, right| {
1175 left.source()
1176 .cmp(right.source())
1177 .then_with(|| left.fq_name().cmp(&right.fq_name()))
1178 });
1179 results.dedup();
1180 (!results.is_empty()).then_some(results)
1181}
1182
1183pub fn resolve_fqn_candidates(
1192 python: &dyn PythonSource,
1193 importer: Option<&ProjectFile>,
1194 fqn: &str,
1195 exact: impl FnOnce(&str) -> Vec<CodeUnit>,
1196) -> Vec<CodeUnit> {
1197 if let Some(candidates) = resolve_direct_named_exported_fqn(python, importer, fqn) {
1198 return candidates;
1199 }
1200 let candidates = resolve_exported_fqn(python, importer, fqn);
1201 if !candidates.is_empty() {
1202 return candidates;
1203 }
1204 exact(fqn)
1205}
1206
1207fn resolve_exported_name_from_module(
1208 python: &dyn PythonSource,
1209 importer: Option<&ProjectFile>,
1210 module: &str,
1211 name: &str,
1212) -> Vec<CodeUnit> {
1213 let Some(module_unit) = resolve_module_code_unit(python, importer, module) else {
1214 return Vec::new();
1215 };
1216 resolve_exported_name(python, module_unit.source(), name)
1217}
1218
1219fn resolve_exported_name(
1220 python: &dyn PythonSource,
1221 module_file: &ProjectFile,
1222 name: &str,
1223) -> Vec<CodeUnit> {
1224 let mut results = Vec::new();
1225 let mut queue = VecDeque::from([(module_file.clone(), name.to_string())]);
1226 let mut visited = HashSet::default();
1227
1228 while let Some((file, export_name)) = queue.pop_front() {
1229 if !visited.insert((file.clone(), export_name.clone())) {
1230 continue;
1231 }
1232
1233 let index = python.export_index_of(&file);
1234 if let Some(entry) = index.exports_by_name.get(&export_name) {
1235 match entry {
1236 ExportEntry::Local { local_name } => {
1237 results.extend(local_export_declarations(python, &file, local_name));
1238 }
1239 ExportEntry::ReexportedNamed {
1240 module_specifier,
1241 imported_name,
1242 } => {
1243 for target_file in
1244 resolve_module_files_for_export(python, &file, module_specifier)
1245 {
1246 queue.push_back((target_file, imported_name.clone()));
1247 }
1248 }
1249 ExportEntry::ReexportedModule { module_specifier } => {
1250 results.extend(resolve_module_code_unit(
1253 python,
1254 Some(&file),
1255 module_specifier,
1256 ));
1257 }
1258 ExportEntry::Default { local_name } => {
1259 if let Some(local_name) = local_name {
1260 results.extend(local_export_declarations(python, &file, local_name));
1261 }
1262 }
1263 }
1264 continue;
1265 }
1266
1267 if !export_name.starts_with('_') {
1268 for star in &index.reexport_stars {
1269 for target_file in
1270 resolve_module_files_for_export(python, &file, &star.module_specifier)
1271 {
1272 queue.push_back((target_file, export_name.clone()));
1273 }
1274 }
1275 }
1276 }
1277
1278 results.sort_by(|left, right| {
1279 left.source()
1280 .cmp(right.source())
1281 .then_with(|| left.fq_name().cmp(&right.fq_name()))
1282 });
1283 results.dedup();
1284 results
1285}
1286
1287fn local_export_declarations(
1288 index: &dyn CodeUnitIndex,
1289 file: &ProjectFile,
1290 local_name: &str,
1291) -> Vec<CodeUnit> {
1292 index
1293 .top_level_declarations(file)
1294 .into_iter()
1295 .filter(|unit| {
1296 unit.identifier() == local_name
1297 && index
1298 .parent_of(unit)
1299 .is_some_and(|parent| parent.is_module() && parent.source() == file)
1300 })
1301 .collect()
1302}
1303
1304fn resolve_module_files_for_export(
1305 python: &dyn PythonSource,
1306 importing_file: &ProjectFile,
1307 module_specifier: &str,
1308) -> Vec<ProjectFile> {
1309 let resolved_module =
1310 resolve_python_module_specifier(python.project(), importing_file, module_specifier);
1311 let Some(resolved_module) = resolved_module else {
1312 return Vec::new();
1313 };
1314 resolve_module_code_unit(python, Some(importing_file), &resolved_module)
1318 .map(|unit| vec![unit.source().clone()])
1319 .unwrap_or_default()
1320}
1321
1322pub fn extract_package_from_python_wildcard(import: &ImportInfo) -> Option<String> {
1323 let details = python_import_details(import)?;
1324 match details {
1325 PythonImportDetails::FromImport {
1326 module, wildcard, ..
1327 } if wildcard => Some(module),
1328 _ => None,
1329 }
1330}
1331
1332#[derive(Debug, Clone)]
1333pub enum PythonImportDetails {
1334 Import {
1335 module: String,
1336 alias: Option<String>,
1337 },
1338 FromImport {
1339 module: String,
1340 name: String,
1341 alias: Option<String>,
1342 wildcard: bool,
1343 },
1344}
1345
1346#[derive(Debug, Clone)]
1353pub struct PythonImportSyntax<'tree> {
1354 pub info: ImportInfo,
1355 pub module: Node<'tree>,
1356 pub member: Option<Node<'tree>>,
1357 pub leaf: Node<'tree>,
1358 pub declaration: Node<'tree>,
1359 pub target: Option<Node<'tree>>,
1360 pub alias: Option<Node<'tree>>,
1361}
1362
1363pub fn python_import_infos_from_node(node: Node<'_>, source: &str) -> Vec<ImportInfo> {
1364 python_import_syntaxes_from_node(node, source)
1365 .into_iter()
1366 .map(|syntax| syntax.info)
1367 .collect()
1368}
1369
1370pub fn python_import_syntaxes_from_node<'tree>(
1374 node: Node<'tree>,
1375 source: &str,
1376) -> Vec<PythonImportSyntax<'tree>> {
1377 match node.kind() {
1378 "import_statement" => python_namespace_import_syntaxes(node, source),
1379 "import_from_statement" => python_from_import_syntaxes(node, source),
1380 _ => Vec::new(),
1381 }
1382}
1383
1384pub fn python_import_details(import: &ImportInfo) -> Option<PythonImportDetails> {
1385 let path = import.path.as_ref()?;
1386 match path.kind? {
1387 StructuredImportPathKind::Namespace => Some(PythonImportDetails::Import {
1388 module: join_python_import_segments(&path.segments),
1389 alias: import.alias.clone(),
1390 }),
1391 StructuredImportPathKind::StaticMember | StructuredImportPathKind::ExternCrate => None,
1393 StructuredImportPathKind::FileInclude | StructuredImportPathKind::DynamicInclude => None,
1395 StructuredImportPathKind::ImportFrom => {
1396 let (name, module_segments) = if import.is_wildcard {
1397 ("*".to_string(), path.segments.as_slice())
1398 } else {
1399 let (name, module_segments) = path.segments.split_last()?;
1400 (name.clone(), module_segments)
1401 };
1402 Some(PythonImportDetails::FromImport {
1403 module: join_python_import_segments(module_segments),
1404 name,
1405 alias: import.alias.clone(),
1406 wildcard: import.is_wildcard,
1407 })
1408 }
1409 }
1410}
1411
1412fn python_namespace_import_syntaxes<'tree>(
1413 node: Node<'tree>,
1414 source: &str,
1415) -> Vec<PythonImportSyntax<'tree>> {
1416 let mut infos = Vec::new();
1417 let mut cursor = node.walk();
1418 for imported in node.children_by_field_name("name", &mut cursor) {
1419 let (module_node, alias_node) = if imported.kind() == "aliased_import" {
1420 let Some(name) = imported.child_by_field_name("name") else {
1421 continue;
1422 };
1423 (name, imported.child_by_field_name("alias"))
1424 } else {
1425 (imported, None)
1426 };
1427 let alias = alias_node
1428 .map(|alias| py_node_text(alias, source).trim().to_string())
1429 .filter(|alias| !alias.is_empty());
1430 let segments = python_path_segments(module_node, source);
1431 if segments.is_empty() {
1432 continue;
1433 }
1434 let target = alias
1437 .is_none()
1438 .then(|| python_first_segment_node(module_node))
1439 .flatten();
1440 let binder_span = alias_node
1441 .or(target)
1442 .map(brokk_bifrost_core::analyzer::common::node_span);
1443 let module = join_python_import_segments(&segments);
1444 let identifier = alias.clone().or_else(|| segments.first().cloned());
1445 let info = ImportInfo {
1446 raw_snippet: if let Some(alias) = &alias {
1447 format!("import {module} as {alias}")
1448 } else {
1449 format!("import {module}")
1450 },
1451 is_wildcard: false,
1452 is_global: false,
1453 identifier,
1454 alias,
1455 path: Some(StructuredImportPath {
1456 segments,
1457 kind: Some(StructuredImportPathKind::Namespace),
1458 lexical_prefixes: Vec::new(),
1459 lexical_scopes: Vec::new(),
1460 declaration_start_byte: node.start_byte(),
1461 }),
1462 binder_span,
1463 };
1464 infos.push(PythonImportSyntax {
1465 info,
1466 module: module_node,
1467 member: None,
1468 leaf: imported,
1469 declaration: node,
1470 target,
1471 alias: alias_node,
1472 });
1473 }
1474 infos
1475}
1476
1477fn python_from_import_syntaxes<'tree>(
1478 node: Node<'tree>,
1479 source: &str,
1480) -> Vec<PythonImportSyntax<'tree>> {
1481 let Some(module_node) = node.child_by_field_name("module_name") else {
1482 return Vec::new();
1483 };
1484 let module_segments = python_module_segments(module_node, source);
1485 if module_segments.is_empty() {
1486 return Vec::new();
1487 }
1488
1489 let mut infos = Vec::new();
1490 let has_wildcard_import = {
1491 let mut cursor = node.walk();
1492 node.named_children(&mut cursor)
1493 .any(|child| child.kind() == "wildcard_import")
1494 };
1495 let mut cursor = node.walk();
1496 let imported_names: Vec<_> = node.children_by_field_name("name", &mut cursor).collect();
1497 if has_wildcard_import {
1498 let module = join_python_import_segments(&module_segments);
1499 let info = ImportInfo {
1500 raw_snippet: format!("from {module} import *"),
1501 is_wildcard: true,
1502 is_global: false,
1503 identifier: None,
1504 alias: None,
1505 path: Some(StructuredImportPath {
1506 segments: module_segments,
1507 kind: Some(StructuredImportPathKind::ImportFrom),
1508 lexical_prefixes: Vec::new(),
1509 lexical_scopes: Vec::new(),
1510 declaration_start_byte: node.start_byte(),
1511 }),
1512 binder_span: None,
1513 };
1514 infos.push(PythonImportSyntax {
1515 info,
1516 module: module_node,
1517 member: None,
1518 leaf: node,
1519 declaration: node,
1520 target: None,
1521 alias: None,
1522 });
1523 return infos;
1524 }
1525 if imported_names.is_empty() {
1526 return infos;
1527 }
1528
1529 for imported in imported_names {
1530 let (name_node, alias_node) = if imported.kind() == "aliased_import" {
1531 let Some(name) = imported.child_by_field_name("name") else {
1532 continue;
1533 };
1534 (name, imported.child_by_field_name("alias"))
1535 } else {
1536 (imported, None)
1537 };
1538 let alias = alias_node
1539 .map(|alias| py_node_text(alias, source).trim().to_string())
1540 .filter(|alias| !alias.is_empty());
1541 let name_segments = python_path_segments(name_node, source);
1542 if name_segments.is_empty() {
1543 continue;
1544 }
1545 let target = alias
1548 .is_none()
1549 .then(|| {
1550 (name_segments.len() == 1)
1551 .then(|| python_first_segment_node(name_node))
1552 .flatten()
1553 })
1554 .flatten();
1555 let binder_span = alias_node
1556 .or(target)
1557 .map(brokk_bifrost_core::analyzer::common::node_span);
1558 let imported_name = join_python_import_segments(&name_segments);
1559 let mut segments = module_segments.clone();
1560 segments.extend(name_segments);
1561 let module = join_python_import_segments(&module_segments);
1562 let info = ImportInfo {
1563 raw_snippet: if let Some(alias) = &alias {
1564 format!("from {module} import {imported_name} as {alias}")
1565 } else {
1566 format!("from {module} import {imported_name}")
1567 },
1568 is_wildcard: false,
1569 is_global: false,
1570 identifier: Some(alias.clone().unwrap_or_else(|| imported_name.clone())),
1571 alias,
1572 path: Some(StructuredImportPath {
1573 segments,
1574 kind: Some(StructuredImportPathKind::ImportFrom),
1575 lexical_prefixes: Vec::new(),
1576 lexical_scopes: Vec::new(),
1577 declaration_start_byte: node.start_byte(),
1578 }),
1579 binder_span,
1580 };
1581 infos.push(PythonImportSyntax {
1582 info,
1583 module: module_node,
1584 member: Some(name_node),
1585 leaf: imported,
1586 declaration: node,
1587 target,
1588 alias: alias_node,
1589 });
1590 }
1591 infos
1592}
1593
1594fn python_module_segments(module: Node<'_>, source: &str) -> Vec<String> {
1595 if module.kind() == "relative_import" {
1596 let mut cursor = module.walk();
1597 let mut prefix = String::new();
1598 let mut path_node = None;
1599 for child in module.named_children(&mut cursor) {
1600 match child.kind() {
1601 "import_prefix" if prefix.is_empty() => {
1602 prefix = py_node_text(child, source).trim().to_string();
1603 }
1604 "dotted_name" if path_node.is_none() => {
1605 path_node = Some(child);
1606 }
1607 _ => {}
1608 }
1609 }
1610 let mut segments = path_node
1611 .map(|path| python_path_segments(path, source))
1612 .unwrap_or_default();
1613 if !prefix.is_empty() {
1614 if let Some(first) = segments.first_mut() {
1615 first.insert_str(0, &prefix);
1616 } else {
1617 segments.push(prefix);
1618 }
1619 }
1620 return segments;
1621 }
1622 python_path_segments(module, source)
1623}
1624
1625fn python_first_segment_node(node: Node<'_>) -> Option<Node<'_>> {
1629 match node.kind() {
1630 "identifier" => Some(node),
1631 "dotted_name" => {
1632 let mut cursor = node.walk();
1633 node.named_children(&mut cursor)
1634 .find(|child| child.kind() == "identifier")
1635 }
1636 _ => None,
1637 }
1638}
1639
1640pub(crate) fn python_import_identifier_nodes(node: Node<'_>) -> Vec<Node<'_>> {
1643 use brokk_bifrost_core::analyzer::tree_walk::{WalkControl, walk_named_tree_preorder};
1644 let mut segments = Vec::new();
1645 walk_named_tree_preorder(node, true, |node| {
1646 if node.kind() == "identifier" {
1647 segments.push(node);
1648 }
1649 WalkControl::Continue
1650 });
1651 segments
1652}
1653
1654fn python_path_segments(node: Node<'_>, source: &str) -> Vec<String> {
1655 python_import_identifier_nodes(node)
1656 .into_iter()
1657 .map(|node| py_node_text(node, source).trim().to_string())
1658 .collect()
1659}
1660
1661fn join_python_import_segments(segments: &[String]) -> String {
1662 let Some((first, rest)) = segments.split_first() else {
1663 return String::new();
1664 };
1665 if first.starts_with('.') && !rest.is_empty() {
1666 format!("{first}.{}", rest.join("."))
1667 } else {
1668 segments.join(".")
1669 }
1670}
1671
1672pub fn python_namespace_binding_name(
1673 import: &ImportInfo,
1674 alias: Option<&str>,
1675 module: &str,
1676) -> String {
1677 import
1678 .identifier
1679 .clone()
1680 .or_else(|| alias.map(str::to_string))
1681 .unwrap_or_else(|| module.to_string())
1682}
1683
1684pub fn python_namespace_binding_module(
1685 import: &ImportInfo,
1686 alias: Option<&str>,
1687 module: &str,
1688) -> String {
1689 if alias.is_some() {
1690 return module.to_string();
1691 }
1692 import
1693 .path
1694 .as_ref()
1695 .and_then(|path| path.segments.first().cloned())
1696 .unwrap_or_else(|| module.to_string())
1697}
1698
1699pub fn resolve_python_relative_module(
1700 project: &dyn brokk_bifrost_core::analyzer::project::Project,
1701 source_file: &ProjectFile,
1702 module_expr: &str,
1703) -> Option<String> {
1704 resolve_python_relative_module_from_package(
1705 &python_current_package(project, source_file),
1706 module_expr,
1707 )
1708}
1709
1710pub fn resolve_python_module_specifier(
1712 project: &dyn brokk_bifrost_core::analyzer::project::Project,
1713 importing_file: &ProjectFile,
1714 module_specifier: &str,
1715) -> Option<String> {
1716 if module_specifier.starts_with('.') {
1717 resolve_python_relative_module(project, importing_file, module_specifier)
1718 } else {
1719 Some(module_specifier.to_string())
1720 }
1721}
1722
1723pub fn resolve_python_relative_module_from_package(
1727 current_package: &str,
1728 module_expr: &str,
1729) -> Option<String> {
1730 let level = module_expr.chars().take_while(|ch| *ch == '.').count();
1731 let suffix = module_expr[level..].trim_matches('.');
1732 let mut parts: Vec<_> = current_package
1733 .split('.')
1734 .filter(|part| !part.is_empty())
1735 .map(str::to_string)
1736 .collect();
1737 if level == 0 {
1738 return Some(module_expr.to_string());
1739 }
1740 if level > 0 {
1741 if level - 1 > parts.len() {
1742 return None;
1743 }
1744 parts.truncate(parts.len() - (level - 1));
1745 }
1746 if !suffix.is_empty() {
1747 parts.extend(suffix.split('.').map(str::to_string));
1748 }
1749 Some(parts.join("."))
1750}
1751
1752fn python_current_package(
1753 project: &dyn brokk_bifrost_core::analyzer::project::Project,
1754 source_file: &ProjectFile,
1755) -> String {
1756 let module = python_module_name(project, source_file);
1757 if source_file
1758 .rel_path()
1759 .file_name()
1760 .and_then(|name| name.to_str())
1761 == Some("__init__.py")
1762 {
1763 module
1764 } else {
1765 module
1766 .rsplit_once('.')
1767 .map(|(package, _)| package.to_string())
1768 .unwrap_or_default()
1769 }
1770}