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 let mut bindings = imports
576 .iter()
577 .filter_map(|import| {
578 let path = import.path.as_ref()?;
579 let details = python_import_details(import)?;
580 let mut binding_scope = None;
581 for scope in path.lexical_scopes.iter().rev() {
582 let node = root.descendant_for_byte_range(scope.start_byte, scope.end_byte)?;
583 if node.start_byte() != scope.start_byte || node.end_byte() != scope.end_byte {
584 return None;
585 }
586 if matches!(node.kind(), "function_definition" | "lambda") {
587 binding_scope = Some((scope.start_byte, scope.end_byte, true));
588 break;
589 }
590 }
591 let (scope_start_byte, scope_end_byte, function_scoped) =
592 binding_scope.unwrap_or((0, source_len, false));
593 let start_byte = path.declaration_start_byte;
594 match details {
595 PythonImportDetails::Import { module, alias } => {
596 let consumed_attributes = if alias.is_some() {
597 0
598 } else {
599 path.segments.len().saturating_sub(1)
600 };
601 Some(PythonImportBinding {
602 start_byte,
603 scope_start_byte,
604 scope_end_byte,
605 function_scoped,
606 local_name: import.identifier.clone()?,
607 qualified_name: module,
608 consumed_attributes,
609 })
610 }
611 PythonImportDetails::FromImport {
612 module,
613 name,
614 alias: _,
615 wildcard: false,
616 } => Some(PythonImportBinding {
617 start_byte,
618 scope_start_byte,
619 scope_end_byte,
620 function_scoped,
621 local_name: import.identifier.clone()?,
622 qualified_name: format!("{module}.{name}"),
623 consumed_attributes: 0,
624 }),
625 PythonImportDetails::FromImport { wildcard: true, .. } => None,
626 }
627 })
628 .collect::<Vec<_>>();
629 bindings.sort_by_key(|binding| binding.start_byte);
630 bindings
631}
632
633fn python_import_binding_scope(node: Node<'_>, source_len: usize) -> (usize, usize, bool) {
634 let mut parent = node.parent();
635 while let Some(scope) = parent {
636 if matches!(scope.kind(), "function_definition" | "lambda") {
637 return (scope.start_byte(), scope.end_byte(), true);
638 }
639 parent = scope.parent();
640 }
641 (0, source_len, false)
642}
643
644#[cfg(test)]
645mod tests {
646 use super::*;
647 use brokk_bifrost_core::analyzer::usages::model::ImportBinder;
648 use brokk_bifrost_test_macros::hardwon;
649
650 fn importlib_assignment_for(source: &str, importlib_local: &str) -> Option<String> {
651 let tree = parse_python_tree(source).expect("valid Python fixture");
652 let reference_start = source.rfind("canonical.setup").expect("reference site");
653 let mut stack = vec![tree.root_node()];
654 let mut reference = None;
655 while let Some(node) = stack.pop() {
656 if node.kind() == "identifier"
657 && node.start_byte() == reference_start
658 && node_source_text(node, source) == "canonical"
659 {
660 reference = Some(node);
661 break;
662 }
663 let mut cursor = node.walk();
664 stack.extend(node.named_children(&mut cursor));
665 }
666 imported_module_assignment_at(
667 reference.expect("canonical reference node"),
668 "canonical",
669 source,
670 |local| local == importlib_local,
671 )
672 }
673
674 #[test]
675 fn importlib_assignment_requires_a_direct_literal_dominating_write() {
676 let valid = r#"import importlib as loader
677
678def test():
679 canonical = loader.import_module("routes.contacts.contacts_routes")
680 assert canonical.setup_contacts_routes
681"#;
682 assert_eq!(
683 importlib_assignment_for(valid, "loader"),
684 Some("routes.contacts.contacts_routes".to_string())
685 );
686
687 for near_miss in [
688 r#"import importlib as loader
689def test(module_name):
690 canonical = loader.import_module(module_name)
691 assert canonical.setup_contacts_routes
692"#,
693 r#"import importlib as loader
694def test():
695 if enabled:
696 canonical = loader.import_module("routes.contacts.contacts_routes")
697 assert canonical.setup_contacts_routes
698"#,
699 r#"import importlib as loader
700def test():
701 canonical = loader.import_module("routes.contacts.contacts_routes")
702 canonical = object()
703 assert canonical.setup_contacts_routes
704"#,
705 ] {
706 assert_eq!(
707 importlib_assignment_for(near_miss, "loader"),
708 None,
709 "near miss must not infer a module receiver: {near_miss:?}"
710 );
711 }
712 assert_eq!(
713 importlib_assignment_for(valid, "shadowed_loader"),
714 None,
715 "a similarly shaped non-importlib receiver must not resolve"
716 );
717 }
718
719 fn replacement_for(source: &str, binder: &ImportBinder) -> Option<PythonModuleReplacement> {
720 let tree = parse_python_tree(source).expect("valid Python fixture");
721 let root = tree.root_node();
722 let mut replacement = None;
723 let mut cursor = root.walk();
724 for statement in root.named_children(&mut cursor) {
725 let statement = if statement.kind() == "expression_statement" {
726 statement.named_child(0).expect("fixture expression")
727 } else {
728 statement
729 };
730 if statement.kind() == "assignment"
731 && let Some(next) =
732 module_replacement_from_assignment(statement, source, &binder.bindings)
733 && replacement.replace(next).is_some()
734 {
735 return None;
736 }
737 }
738 replacement
739 }
740
741 #[test]
742 #[hardwon("59dfb80c3bb8: module replacement accepted malformed or shadowed import shapes")]
743 fn module_replacement_requires_exact_structured_import_bindings() {
744 let source = r#"import sys as _sys
745from routes.contacts import contacts_routes as _canonical
746
747_sys.modules[__name__] = _canonical
748"#;
749 let mut binder = ImportBinder::empty();
750 binder.bindings.insert(
751 "_sys".to_string(),
752 ImportBinding {
753 module_specifier: "sys".to_string(),
754 namespace_imported_module: Some("sys".to_string()),
755 kind: ImportKind::Namespace,
756 imported_name: None,
757 },
758 );
759 binder.bindings.insert(
760 "_canonical".to_string(),
761 ImportBinding {
762 module_specifier: "routes.contacts.contacts_routes".to_string(),
763 namespace_imported_module: None,
764 kind: ImportKind::Namespace,
765 imported_name: None,
766 },
767 );
768
769 assert_eq!(
770 replacement_for(source, &binder),
771 Some(PythonModuleReplacement {
772 target_module: "routes.contacts.contacts_routes".to_string(),
773 })
774 );
775
776 for near_miss in [
777 "cache.modules[__name__] = _canonical\n",
778 "_sys.modules[module_name] = _canonical\n",
779 "_sys.modules[__name__] = build()\n",
780 "def replace():\n _sys.modules[__name__] = _canonical\n",
781 ] {
782 assert_eq!(
783 replacement_for(near_miss, &binder),
784 None,
785 "near miss must not replace module identity: {near_miss:?}"
786 );
787 }
788 }
789
790 #[test]
791 #[hardwon(
792 "b58235b247fd: Python import bindings lost function scope or relative module identity"
793 )]
794 fn existing_tree_import_bindings_preserve_function_scope_and_relative_identity() {
795 let source = r#"def register():
796 from .workers.special import Worker
797 return Worker()
798"#;
799 let tree = parse_python_tree(source).expect("Python tree");
800 let bindings = python_import_bindings_from_tree(tree.root_node(), source);
801
802 assert_eq!(bindings.len(), 1, "{bindings:#?}");
803 assert!(bindings[0].is_function_scoped(), "{bindings:#?}");
804 assert_eq!(bindings[0].local_name, "Worker");
805 assert_eq!(bindings[0].qualified_name, ".workers.special.Worker");
806 assert!(bindings[0].scope_start_byte <= bindings[0].start_byte);
807 assert!(bindings[0].scope_end_byte >= source.rfind("Worker()").unwrap());
808 }
809}
810
811pub(crate) fn python_imports_by_statement(imports: &[ImportInfo]) -> HashMap<usize, &[ImportInfo]> {
814 let mut statement_imports = HashMap::default();
815 for group in imports.chunk_by(|left, right| {
816 left.path.as_ref().map(|path| path.declaration_start_byte)
817 == right.path.as_ref().map(|path| path.declaration_start_byte)
818 }) {
819 if let Some(path) = &group[0].path {
820 assert!(
821 statement_imports
822 .insert(path.declaration_start_byte, group)
823 .is_none(),
824 "canonical Python import leaves must remain in statement order"
825 );
826 }
827 }
828 statement_imports
829}
830
831pub fn module_replacement_of(
832 python: &dyn PythonSource,
833 file: &ProjectFile,
834 source: &str,
835 imports: &[ImportInfo],
836) -> Option<PythonModuleReplacement> {
837 let tree = parse_python_tree(source)?;
838 let root = tree.root_node();
839 let mut bindings: HashMap<String, ImportBinding> = HashMap::default();
840 let statement_imports = python_imports_by_statement(imports);
841 let mut replacement = None;
842 let mut cursor = root.walk();
843 for statement in root.named_children(&mut cursor) {
844 let statement = if statement.kind() == "expression_statement" {
845 let Some(expression) = statement.named_child(0) else {
846 continue;
847 };
848 expression
849 } else {
850 statement
851 };
852 match statement.kind() {
853 "import_statement" | "import_from_statement" => {
854 if let Some(imports) = statement_imports.get(&statement.start_byte()) {
855 bindings.extend(import_binder_from_imports(python, file, imports).bindings);
856 }
857 }
858 "assignment" => {
859 if let Some(next) = module_replacement_from_assignment(statement, source, &bindings)
860 && replacement.replace(next).is_some()
861 {
862 return None;
863 }
864 remove_direct_scope_bindings(statement, source, &mut bindings);
865 }
866 _ => remove_direct_scope_bindings(statement, source, &mut bindings),
867 }
868 }
869 replacement
870}
871
872pub fn resolve_import_bindings(
873 python: &dyn PythonSource,
874 token: QueryToken<'_>,
875 file: &ProjectFile,
876) -> HashMap<String, CodeUnit> {
877 let imports = python.import_info_of(token, file);
878 let mut bindings = HashMap::default();
879 for resolved in resolve_imports_batched(python, file, &imports) {
880 for (binding, code_unit) in resolved {
881 bindings.insert(binding, code_unit);
882 }
883 }
884 bindings
885}
886
887pub fn resolve_imports_batched(
894 python: &dyn PythonSource,
895 file: &ProjectFile,
896 imports: &[ImportInfo],
897) -> Vec<Vec<(String, CodeUnit)>> {
898 let primary_fqns: Vec<Option<String>> = imports
899 .iter()
900 .map(|import| primary_module_fqn(python, file, import))
901 .collect();
902 let to_resolve: Vec<String> = primary_fqns.iter().flatten().cloned().collect();
903 let mut batch_results =
904 resolve_module_code_units_batch(python, Some(file), &to_resolve).into_iter();
905
906 imports
907 .iter()
908 .zip(primary_fqns.iter())
909 .map(|(import, primary_fqn)| {
910 let hint = primary_fqn.as_ref().map(|_| batch_results.next().unwrap());
911 resolve_import_with_hint(python, file, import, hint.as_ref())
912 })
913 .collect()
914}
915
916pub fn resolve_import_files_batched(
926 python: &dyn PythonSource,
927 file: &ProjectFile,
928 imports: &[ImportInfo],
929) -> HashSet<ProjectFile> {
930 let mut module_fqns = Vec::new();
931 for import in imports {
932 match python_import_details(import) {
933 Some(PythonImportDetails::Import { module, .. }) => module_fqns.push(module),
934 Some(PythonImportDetails::FromImport {
935 module,
936 name,
937 wildcard,
938 ..
939 }) => {
940 let resolved_module = if module.starts_with('.') {
941 resolve_python_relative_module(python.project(), file, &module)
942 } else {
943 Some(module)
944 };
945 let Some(resolved_module) = resolved_module else {
946 continue;
947 };
948 if !wildcard {
949 module_fqns.push(format!("{resolved_module}.{name}"));
950 }
951 module_fqns.push(resolved_module);
952 }
953 None => {}
954 }
955 }
956 module_fqns.sort();
957 module_fqns.dedup();
958
959 python
960 .path_module_fqns_batch(&module_fqns)
961 .into_iter()
962 .flatten()
963 .flatten()
964 .filter(CodeUnit::is_module)
965 .map(|unit| unit.source().clone())
966 .collect()
967}
968
969fn primary_module_fqn(
973 python: &dyn PythonSource,
974 file: &ProjectFile,
975 import: &ImportInfo,
976) -> Option<String> {
977 match python_import_details(import)? {
978 PythonImportDetails::Import { module, alias } => Some(python_namespace_binding_module(
979 import,
980 alias.as_deref(),
981 &module,
982 )),
983 PythonImportDetails::FromImport {
984 module,
985 name,
986 wildcard,
987 ..
988 } => {
989 if wildcard {
990 return None;
991 }
992 let resolved_module = resolve_python_module_specifier(python.project(), file, &module);
993 resolved_module.map(|resolved_module| format!("{resolved_module}.{name}"))
994 }
995 }
996}
997
998pub fn resolve_import(
999 python: &dyn PythonSource,
1000 file: &ProjectFile,
1001 import: &ImportInfo,
1002) -> Vec<(String, CodeUnit)> {
1003 resolve_import_with_hint(python, file, import, None)
1004}
1005
1006fn resolve_import_with_hint(
1009 python: &dyn PythonSource,
1010 file: &ProjectFile,
1011 import: &ImportInfo,
1012 primary_hint: Option<&Option<CodeUnit>>,
1013) -> Vec<(String, CodeUnit)> {
1014 if let Some(details) = python_import_details(import) {
1015 match details {
1016 PythonImportDetails::Import { module, alias } => {
1017 let binding = python_namespace_binding_name(import, alias.as_deref(), &module);
1018 let bound_module =
1019 python_namespace_binding_module(import, alias.as_deref(), &module);
1020 let resolved = match primary_hint {
1021 Some(hint) => hint.clone(),
1022 None => resolve_module_code_unit(python, Some(file), &bound_module),
1023 };
1024 if let Some(module_code_unit) = resolved {
1025 return vec![(binding, module_code_unit)];
1026 }
1027 }
1028 PythonImportDetails::FromImport {
1029 module,
1030 name,
1031 alias,
1032 wildcard,
1033 } => {
1034 let resolved_module =
1035 resolve_python_module_specifier(python.project(), file, &module);
1036 let Some(resolved_module) = resolved_module else {
1037 return Vec::new();
1038 };
1039 if wildcard {
1040 return public_declarations_in_module(python, file, &resolved_module)
1041 .into_iter()
1042 .map(|code_unit| (code_unit.identifier().to_string(), code_unit))
1043 .collect();
1044 }
1045
1046 let binding = alias.clone().unwrap_or_else(|| name.clone());
1047 let module_candidate = format!("{resolved_module}.{name}");
1048 let exported =
1049 resolve_exported_name_from_module(python, Some(file), &resolved_module, &name);
1050 if !exported.is_empty() {
1051 return exported
1052 .into_iter()
1053 .map(|code_unit| (binding.clone(), code_unit))
1054 .collect();
1055 }
1056 let resolved = match primary_hint {
1057 Some(hint) => hint.clone(),
1058 None => resolve_module_code_unit(python, Some(file), &module_candidate),
1059 };
1060 if let Some(code_unit) = resolved {
1061 return vec![(binding, code_unit)];
1062 }
1063 let definitions: Vec<_> = python.definitions(&module_candidate).collect();
1064 if !definitions.is_empty() {
1065 return definitions
1066 .into_iter()
1067 .map(|code_unit| (binding.clone(), code_unit))
1068 .collect();
1069 }
1070 let package_candidate: Vec<_> = python
1071 .definitions(&format!("{resolved_module}.{name}"))
1072 .collect();
1073 if !package_candidate.is_empty() {
1074 return package_candidate
1075 .into_iter()
1076 .map(|code_unit| (binding.clone(), code_unit))
1077 .collect();
1078 }
1079 }
1080 }
1081 }
1082 Vec::new()
1083}
1084
1085pub fn resolve_exported_fqn(
1089 python: &dyn PythonSource,
1090 importer: Option<&ProjectFile>,
1091 fqn: &str,
1092) -> Vec<CodeUnit> {
1093 let Some((module, name)) = fqn.rsplit_once('.') else {
1094 return Vec::new();
1095 };
1096 resolve_exported_name_from_module(python, importer, module, name)
1097}
1098
1099fn resolve_direct_named_exported_fqn(
1108 python: &dyn PythonSource,
1109 importer: Option<&ProjectFile>,
1110 fqn: &str,
1111) -> Option<Vec<CodeUnit>> {
1112 let (module, name) = fqn.rsplit_once('.')?;
1113 let mut results = Vec::new();
1114 let mut queue = VecDeque::from([(importer.cloned(), module.to_string(), name.to_string())]);
1115 let mut visited = HashSet::default();
1116
1117 while let Some((importer, module, export_name)) = queue.pop_front() {
1118 if !visited.insert((module.clone(), export_name.clone())) {
1119 continue;
1120 }
1121 let module_unit = resolve_module_code_unit(python, importer.as_ref(), &module)?;
1122 let file = module_unit.source();
1123 let local = local_export_declarations(python, file, &export_name);
1124 let binder = python.import_binder_of(file);
1125 let binding = binder.bindings.get(&export_name);
1126 if !local.is_empty() && binding.is_some() {
1127 return None;
1128 }
1129 if !local.is_empty() {
1130 results.extend(local);
1131 continue;
1132 }
1133 let binding = binding?;
1134 if binding.kind != ImportKind::Named {
1135 return None;
1136 }
1137 let imported_name = binding.imported_name.as_ref()?;
1138 queue.push_back((
1139 Some(file.clone()),
1140 binding.module_specifier.clone(),
1141 imported_name.clone(),
1142 ));
1143 }
1144
1145 results.sort_by(|left, right| {
1146 left.source()
1147 .cmp(right.source())
1148 .then_with(|| left.fq_name().cmp(&right.fq_name()))
1149 });
1150 results.dedup();
1151 (!results.is_empty()).then_some(results)
1152}
1153
1154pub fn resolve_fqn_candidates(
1163 python: &dyn PythonSource,
1164 importer: Option<&ProjectFile>,
1165 fqn: &str,
1166 exact: impl FnOnce(&str) -> Vec<CodeUnit>,
1167) -> Vec<CodeUnit> {
1168 if let Some(candidates) = resolve_direct_named_exported_fqn(python, importer, fqn) {
1169 return candidates;
1170 }
1171 let candidates = resolve_exported_fqn(python, importer, fqn);
1172 if !candidates.is_empty() {
1173 return candidates;
1174 }
1175 exact(fqn)
1176}
1177
1178fn resolve_exported_name_from_module(
1179 python: &dyn PythonSource,
1180 importer: Option<&ProjectFile>,
1181 module: &str,
1182 name: &str,
1183) -> Vec<CodeUnit> {
1184 let Some(module_unit) = resolve_module_code_unit(python, importer, module) else {
1185 return Vec::new();
1186 };
1187 resolve_exported_name(python, module_unit.source(), name)
1188}
1189
1190fn resolve_exported_name(
1191 python: &dyn PythonSource,
1192 module_file: &ProjectFile,
1193 name: &str,
1194) -> Vec<CodeUnit> {
1195 let mut results = Vec::new();
1196 let mut queue = VecDeque::from([(module_file.clone(), name.to_string())]);
1197 let mut visited = HashSet::default();
1198
1199 while let Some((file, export_name)) = queue.pop_front() {
1200 if !visited.insert((file.clone(), export_name.clone())) {
1201 continue;
1202 }
1203
1204 let index = python.export_index_of(&file);
1205 if let Some(entry) = index.exports_by_name.get(&export_name) {
1206 match entry {
1207 ExportEntry::Local { local_name } => {
1208 results.extend(local_export_declarations(python, &file, local_name));
1209 }
1210 ExportEntry::ReexportedNamed {
1211 module_specifier,
1212 imported_name,
1213 } => {
1214 for target_file in
1215 resolve_module_files_for_export(python, &file, module_specifier)
1216 {
1217 queue.push_back((target_file, imported_name.clone()));
1218 }
1219 }
1220 ExportEntry::ReexportedModule { module_specifier } => {
1221 results.extend(resolve_module_code_unit(
1224 python,
1225 Some(&file),
1226 module_specifier,
1227 ));
1228 }
1229 ExportEntry::Default { local_name } => {
1230 if let Some(local_name) = local_name {
1231 results.extend(local_export_declarations(python, &file, local_name));
1232 }
1233 }
1234 }
1235 continue;
1236 }
1237
1238 if !export_name.starts_with('_') {
1239 for star in &index.reexport_stars {
1240 for target_file in
1241 resolve_module_files_for_export(python, &file, &star.module_specifier)
1242 {
1243 queue.push_back((target_file, export_name.clone()));
1244 }
1245 }
1246 }
1247 }
1248
1249 results.sort_by(|left, right| {
1250 left.source()
1251 .cmp(right.source())
1252 .then_with(|| left.fq_name().cmp(&right.fq_name()))
1253 });
1254 results.dedup();
1255 results
1256}
1257
1258fn local_export_declarations(
1259 index: &dyn CodeUnitIndex,
1260 file: &ProjectFile,
1261 local_name: &str,
1262) -> Vec<CodeUnit> {
1263 index
1264 .top_level_declarations(file)
1265 .into_iter()
1266 .filter(|unit| {
1267 unit.identifier() == local_name
1268 && index
1269 .parent_of(unit)
1270 .is_some_and(|parent| parent.is_module() && parent.source() == file)
1271 })
1272 .collect()
1273}
1274
1275fn resolve_module_files_for_export(
1276 python: &dyn PythonSource,
1277 importing_file: &ProjectFile,
1278 module_specifier: &str,
1279) -> Vec<ProjectFile> {
1280 let resolved_module =
1281 resolve_python_module_specifier(python.project(), importing_file, module_specifier);
1282 let Some(resolved_module) = resolved_module else {
1283 return Vec::new();
1284 };
1285 resolve_module_code_unit(python, Some(importing_file), &resolved_module)
1289 .map(|unit| vec![unit.source().clone()])
1290 .unwrap_or_default()
1291}
1292
1293pub fn extract_package_from_python_wildcard(import: &ImportInfo) -> Option<String> {
1294 let details = python_import_details(import)?;
1295 match details {
1296 PythonImportDetails::FromImport {
1297 module, wildcard, ..
1298 } if wildcard => Some(module),
1299 _ => None,
1300 }
1301}
1302
1303#[derive(Debug, Clone)]
1304pub enum PythonImportDetails {
1305 Import {
1306 module: String,
1307 alias: Option<String>,
1308 },
1309 FromImport {
1310 module: String,
1311 name: String,
1312 alias: Option<String>,
1313 wildcard: bool,
1314 },
1315}
1316
1317#[derive(Debug, Clone)]
1324pub struct PythonImportSyntax<'tree> {
1325 pub info: ImportInfo,
1326 pub declaration: Node<'tree>,
1327 pub target: Option<Node<'tree>>,
1328 pub alias: Option<Node<'tree>>,
1329}
1330
1331pub fn python_import_infos_from_node(node: Node<'_>, source: &str) -> Vec<ImportInfo> {
1332 python_import_syntaxes_from_node(node, source)
1333 .into_iter()
1334 .map(|syntax| syntax.info)
1335 .collect()
1336}
1337
1338pub fn python_import_syntaxes_from_node<'tree>(
1342 node: Node<'tree>,
1343 source: &str,
1344) -> Vec<PythonImportSyntax<'tree>> {
1345 match node.kind() {
1346 "import_statement" => python_namespace_import_syntaxes(node, source),
1347 "import_from_statement" => python_from_import_syntaxes(node, source),
1348 _ => Vec::new(),
1349 }
1350}
1351
1352pub fn python_import_details(import: &ImportInfo) -> Option<PythonImportDetails> {
1353 let path = import.path.as_ref()?;
1354 match path.kind? {
1355 StructuredImportPathKind::Namespace => Some(PythonImportDetails::Import {
1356 module: join_python_import_segments(&path.segments),
1357 alias: import.alias.clone(),
1358 }),
1359 StructuredImportPathKind::StaticMember | StructuredImportPathKind::ExternCrate => None,
1361 StructuredImportPathKind::FileInclude | StructuredImportPathKind::DynamicInclude => None,
1363 StructuredImportPathKind::ImportFrom => {
1364 let (name, module_segments) = if import.is_wildcard {
1365 ("*".to_string(), path.segments.as_slice())
1366 } else {
1367 let (name, module_segments) = path.segments.split_last()?;
1368 (name.clone(), module_segments)
1369 };
1370 Some(PythonImportDetails::FromImport {
1371 module: join_python_import_segments(module_segments),
1372 name,
1373 alias: import.alias.clone(),
1374 wildcard: import.is_wildcard,
1375 })
1376 }
1377 }
1378}
1379
1380fn python_namespace_import_syntaxes<'tree>(
1381 node: Node<'tree>,
1382 source: &str,
1383) -> Vec<PythonImportSyntax<'tree>> {
1384 let mut infos = Vec::new();
1385 let mut cursor = node.walk();
1386 for imported in node.children_by_field_name("name", &mut cursor) {
1387 let (module_node, alias_node) = if imported.kind() == "aliased_import" {
1388 let Some(name) = imported.child_by_field_name("name") else {
1389 continue;
1390 };
1391 (name, imported.child_by_field_name("alias"))
1392 } else {
1393 (imported, None)
1394 };
1395 let alias = alias_node
1396 .map(|alias| py_node_text(alias, source).trim().to_string())
1397 .filter(|alias| !alias.is_empty());
1398 let segments = python_path_segments(module_node, source);
1399 if segments.is_empty() {
1400 continue;
1401 }
1402 let target = alias
1405 .is_none()
1406 .then(|| python_first_segment_node(module_node))
1407 .flatten();
1408 let binder_span = alias_node
1409 .or(target)
1410 .map(brokk_bifrost_core::analyzer::common::node_span);
1411 let module = join_python_import_segments(&segments);
1412 let identifier = alias.clone().or_else(|| segments.first().cloned());
1413 let info = ImportInfo {
1414 raw_snippet: if let Some(alias) = &alias {
1415 format!("import {module} as {alias}")
1416 } else {
1417 format!("import {module}")
1418 },
1419 is_wildcard: false,
1420 is_global: false,
1421 identifier,
1422 alias,
1423 path: Some(StructuredImportPath {
1424 segments,
1425 kind: Some(StructuredImportPathKind::Namespace),
1426 lexical_prefixes: Vec::new(),
1427 lexical_scopes: Vec::new(),
1428 declaration_start_byte: node.start_byte(),
1429 }),
1430 binder_span,
1431 };
1432 infos.push(PythonImportSyntax {
1433 info,
1434 declaration: node,
1435 target,
1436 alias: alias_node,
1437 });
1438 }
1439 infos
1440}
1441
1442fn python_from_import_syntaxes<'tree>(
1443 node: Node<'tree>,
1444 source: &str,
1445) -> Vec<PythonImportSyntax<'tree>> {
1446 let Some(module_node) = node.child_by_field_name("module_name") else {
1447 return Vec::new();
1448 };
1449 let module_segments = python_module_segments(module_node, source);
1450 if module_segments.is_empty() {
1451 return Vec::new();
1452 }
1453
1454 let mut infos = Vec::new();
1455 let has_wildcard_import = {
1456 let mut cursor = node.walk();
1457 node.named_children(&mut cursor)
1458 .any(|child| child.kind() == "wildcard_import")
1459 };
1460 let mut cursor = node.walk();
1461 let imported_names: Vec<_> = node.children_by_field_name("name", &mut cursor).collect();
1462 if has_wildcard_import {
1463 let module = join_python_import_segments(&module_segments);
1464 let info = ImportInfo {
1465 raw_snippet: format!("from {module} import *"),
1466 is_wildcard: true,
1467 is_global: false,
1468 identifier: None,
1469 alias: None,
1470 path: Some(StructuredImportPath {
1471 segments: module_segments,
1472 kind: Some(StructuredImportPathKind::ImportFrom),
1473 lexical_prefixes: Vec::new(),
1474 lexical_scopes: Vec::new(),
1475 declaration_start_byte: node.start_byte(),
1476 }),
1477 binder_span: None,
1478 };
1479 infos.push(PythonImportSyntax {
1480 info,
1481 declaration: node,
1482 target: None,
1483 alias: None,
1484 });
1485 return infos;
1486 }
1487 if imported_names.is_empty() {
1488 return infos;
1489 }
1490
1491 for imported in imported_names {
1492 let (name_node, alias_node) = if imported.kind() == "aliased_import" {
1493 let Some(name) = imported.child_by_field_name("name") else {
1494 continue;
1495 };
1496 (name, imported.child_by_field_name("alias"))
1497 } else {
1498 (imported, None)
1499 };
1500 let alias = alias_node
1501 .map(|alias| py_node_text(alias, source).trim().to_string())
1502 .filter(|alias| !alias.is_empty());
1503 let name_segments = python_path_segments(name_node, source);
1504 if name_segments.is_empty() {
1505 continue;
1506 }
1507 let target = alias
1510 .is_none()
1511 .then(|| {
1512 (name_segments.len() == 1)
1513 .then(|| python_first_segment_node(name_node))
1514 .flatten()
1515 })
1516 .flatten();
1517 let binder_span = alias_node
1518 .or(target)
1519 .map(brokk_bifrost_core::analyzer::common::node_span);
1520 let imported_name = join_python_import_segments(&name_segments);
1521 let mut segments = module_segments.clone();
1522 segments.extend(name_segments);
1523 let module = join_python_import_segments(&module_segments);
1524 let info = ImportInfo {
1525 raw_snippet: if let Some(alias) = &alias {
1526 format!("from {module} import {imported_name} as {alias}")
1527 } else {
1528 format!("from {module} import {imported_name}")
1529 },
1530 is_wildcard: false,
1531 is_global: false,
1532 identifier: Some(alias.clone().unwrap_or_else(|| imported_name.clone())),
1533 alias,
1534 path: Some(StructuredImportPath {
1535 segments,
1536 kind: Some(StructuredImportPathKind::ImportFrom),
1537 lexical_prefixes: Vec::new(),
1538 lexical_scopes: Vec::new(),
1539 declaration_start_byte: node.start_byte(),
1540 }),
1541 binder_span,
1542 };
1543 infos.push(PythonImportSyntax {
1544 info,
1545 declaration: node,
1546 target,
1547 alias: alias_node,
1548 });
1549 }
1550 infos
1551}
1552
1553fn python_module_segments(module: Node<'_>, source: &str) -> Vec<String> {
1554 if module.kind() == "relative_import" {
1555 let mut cursor = module.walk();
1556 let mut prefix = String::new();
1557 let mut path_node = None;
1558 for child in module.named_children(&mut cursor) {
1559 match child.kind() {
1560 "import_prefix" if prefix.is_empty() => {
1561 prefix = py_node_text(child, source).trim().to_string();
1562 }
1563 "dotted_name" if path_node.is_none() => {
1564 path_node = Some(child);
1565 }
1566 _ => {}
1567 }
1568 }
1569 let mut segments = path_node
1570 .map(|path| python_path_segments(path, source))
1571 .unwrap_or_default();
1572 if !prefix.is_empty() {
1573 if let Some(first) = segments.first_mut() {
1574 first.insert_str(0, &prefix);
1575 } else {
1576 segments.push(prefix);
1577 }
1578 }
1579 return segments;
1580 }
1581 python_path_segments(module, source)
1582}
1583
1584fn python_first_segment_node(node: Node<'_>) -> Option<Node<'_>> {
1588 match node.kind() {
1589 "identifier" => Some(node),
1590 "dotted_name" => {
1591 let mut cursor = node.walk();
1592 node.named_children(&mut cursor)
1593 .find(|child| child.kind() == "identifier")
1594 }
1595 _ => None,
1596 }
1597}
1598
1599fn python_path_segments(node: Node<'_>, source: &str) -> Vec<String> {
1600 match node.kind() {
1601 "identifier" => vec![py_node_text(node, source).trim().to_string()],
1602 "dotted_name" => {
1603 let mut segments = Vec::new();
1604 let mut cursor = node.walk();
1605 for child in node.named_children(&mut cursor) {
1606 segments.extend(python_path_segments(child, source));
1607 }
1608 segments
1609 }
1610 _ => {
1611 let mut segments = Vec::new();
1612 let mut cursor = node.walk();
1613 for child in node.named_children(&mut cursor) {
1614 segments.extend(python_path_segments(child, source));
1615 }
1616 segments
1617 }
1618 }
1619}
1620
1621fn join_python_import_segments(segments: &[String]) -> String {
1622 let Some((first, rest)) = segments.split_first() else {
1623 return String::new();
1624 };
1625 if first.starts_with('.') && !rest.is_empty() {
1626 format!("{first}.{}", rest.join("."))
1627 } else {
1628 segments.join(".")
1629 }
1630}
1631
1632pub fn python_namespace_binding_name(
1633 import: &ImportInfo,
1634 alias: Option<&str>,
1635 module: &str,
1636) -> String {
1637 import
1638 .identifier
1639 .clone()
1640 .or_else(|| alias.map(str::to_string))
1641 .unwrap_or_else(|| module.to_string())
1642}
1643
1644pub fn python_namespace_binding_module(
1645 import: &ImportInfo,
1646 alias: Option<&str>,
1647 module: &str,
1648) -> String {
1649 if alias.is_some() {
1650 return module.to_string();
1651 }
1652 import
1653 .path
1654 .as_ref()
1655 .and_then(|path| path.segments.first().cloned())
1656 .unwrap_or_else(|| module.to_string())
1657}
1658
1659pub fn resolve_python_relative_module(
1660 project: &dyn brokk_bifrost_core::analyzer::project::Project,
1661 source_file: &ProjectFile,
1662 module_expr: &str,
1663) -> Option<String> {
1664 resolve_python_relative_module_from_package(
1665 &python_current_package(project, source_file),
1666 module_expr,
1667 )
1668}
1669
1670pub fn resolve_python_module_specifier(
1672 project: &dyn brokk_bifrost_core::analyzer::project::Project,
1673 importing_file: &ProjectFile,
1674 module_specifier: &str,
1675) -> Option<String> {
1676 if module_specifier.starts_with('.') {
1677 resolve_python_relative_module(project, importing_file, module_specifier)
1678 } else {
1679 Some(module_specifier.to_string())
1680 }
1681}
1682
1683pub fn resolve_python_relative_module_from_package(
1687 current_package: &str,
1688 module_expr: &str,
1689) -> Option<String> {
1690 let level = module_expr.chars().take_while(|ch| *ch == '.').count();
1691 let suffix = module_expr[level..].trim_matches('.');
1692 let mut parts: Vec<_> = current_package
1693 .split('.')
1694 .filter(|part| !part.is_empty())
1695 .map(str::to_string)
1696 .collect();
1697 if level == 0 {
1698 return Some(module_expr.to_string());
1699 }
1700 if level > 0 {
1701 if level - 1 > parts.len() {
1702 return None;
1703 }
1704 parts.truncate(parts.len() - (level - 1));
1705 }
1706 if !suffix.is_empty() {
1707 parts.extend(suffix.split('.').map(str::to_string));
1708 }
1709 Some(parts.join("."))
1710}
1711
1712fn python_current_package(
1713 project: &dyn brokk_bifrost_core::analyzer::project::Project,
1714 source_file: &ProjectFile,
1715) -> String {
1716 let module = python_module_name(project, source_file);
1717 if source_file
1718 .rel_path()
1719 .file_name()
1720 .and_then(|name| name.to_str())
1721 == Some("__init__.py")
1722 {
1723 module
1724 } else {
1725 module
1726 .rsplit_once('.')
1727 .map(|(package, _)| package.to_string())
1728 .unwrap_or_default()
1729 }
1730}