1use crate::bindings::python_direct_scope_bindings_bounded;
8use crate::declarations::{parse_python_tree, py_node_text, python_module_name};
9use crate::graph_support::{
10 PythonSource, import_binder_from_imports, public_declarations_in_module,
11 resolve_module_code_unit, resolve_module_code_units_batch,
12};
13use brokk_bifrost_core::analyzer::common::node_source_text;
14use brokk_bifrost_core::analyzer::model::{
15 ImportInfo, StructuredImportPath, StructuredImportPathKind,
16};
17use brokk_bifrost_core::analyzer::usages::model::{ExportEntry, ImportBinding, ImportKind};
18use brokk_bifrost_core::analyzer::{CodeUnit, CodeUnitIndex, ProjectFile};
19use brokk_bifrost_core::hash::{HashMap, HashSet};
20use std::collections::VecDeque;
21use tree_sitter::Node;
22
23#[derive(Debug, Clone, PartialEq, Eq)]
24pub struct PythonModuleReplacement {
25 pub target_module: String,
26}
27
28fn module_replacement_from_assignment(
36 assignment: Node<'_>,
37 source: &str,
38 bindings: &HashMap<String, ImportBinding>,
39) -> Option<PythonModuleReplacement> {
40 let (left, right) = (
41 assignment.child_by_field_name("left")?,
42 assignment.child_by_field_name("right")?,
43 );
44 if left.kind() != "subscript" || right.kind() != "identifier" {
45 return None;
46 }
47 let (value, subscript) = (
48 left.child_by_field_name("value")?,
49 left.child_by_field_name("subscript")?,
50 );
51 if value.kind() != "attribute"
52 || subscript.kind() != "identifier"
53 || node_source_text(subscript, source) != "__name__"
54 {
55 return None;
56 }
57 let (sys_local, modules) = (
58 value.child_by_field_name("object")?,
59 value.child_by_field_name("attribute")?,
60 );
61 if sys_local.kind() != "identifier"
62 || modules.kind() != "identifier"
63 || node_source_text(modules, source) != "modules"
64 {
65 return None;
66 }
67 let sys_binding = bindings.get(node_source_text(sys_local, source))?;
68 let sys_module = sys_binding
69 .namespace_imported_module
70 .as_deref()
71 .unwrap_or(&sys_binding.module_specifier);
72 if sys_binding.kind != ImportKind::Namespace || sys_module != "sys" {
73 return None;
74 }
75
76 let target_binding = bindings.get(node_source_text(right, source))?;
77 if target_binding.kind != ImportKind::Namespace {
78 return None;
79 }
80 Some(PythonModuleReplacement {
81 target_module: target_binding
82 .namespace_imported_module
83 .clone()
84 .unwrap_or_else(|| target_binding.module_specifier.clone()),
85 })
86}
87
88fn remove_direct_scope_bindings(
89 statement: Node<'_>,
90 source: &str,
91 bindings: &mut HashMap<String, ImportBinding>,
92) {
93 let mut stack = vec![statement];
94 while let Some(node) = stack.pop() {
95 for binding in python_direct_scope_bindings_bounded(node, source, || true)
96 .expect("unbounded binding collection cannot be cancelled")
97 {
98 bindings.remove(node_source_text(binding.declaration, source));
99 }
100 if matches!(
101 node.kind(),
102 "function_definition" | "class_definition" | "lambda"
103 ) {
104 continue;
105 }
106 let mut cursor = node.walk();
107 stack.extend(node.named_children(&mut cursor));
108 }
109}
110
111pub fn parse_python_import_infos(source: &str) -> Vec<ImportInfo> {
115 let mut parser = tree_sitter::Parser::new();
116 parser
117 .set_language(&tree_sitter_python::LANGUAGE.into())
118 .expect("failed to load Python parser");
119 let Some(tree) = parser.parse(source, None) else {
120 return Vec::new();
121 };
122 let mut pending = vec![tree.root_node()];
123 let mut imports = Vec::new();
124 while let Some(node) = pending.pop() {
125 if matches!(node.kind(), "import_statement" | "import_from_statement") {
126 imports.extend(python_import_infos_from_node(node, source));
127 continue;
128 }
129 let mut cursor = node.walk();
130 pending.extend(node.named_children(&mut cursor));
131 }
132 imports
133}
134
135#[derive(Debug, Clone, PartialEq, Eq)]
136pub struct PythonImportBinding {
137 pub start_byte: usize,
138 pub scope_start_byte: usize,
139 pub scope_end_byte: usize,
140 function_scoped: bool,
141 pub local_name: String,
142 pub qualified_name: String,
143}
144
145impl PythonImportBinding {
146 pub fn is_function_scoped(&self) -> bool {
147 self.function_scoped
148 }
149}
150
151pub fn parse_python_import_bindings(source: &str) -> Vec<PythonImportBinding> {
155 let mut parser = tree_sitter::Parser::new();
156 parser
157 .set_language(&tree_sitter_python::LANGUAGE.into())
158 .expect("failed to load Python parser");
159 let Some(tree) = parser.parse(source, None) else {
160 return Vec::new();
161 };
162 let mut pending = vec![tree.root_node()];
163 let mut nodes = Vec::new();
164 while let Some(node) = pending.pop() {
165 if matches!(node.kind(), "import_statement" | "import_from_statement") {
166 nodes.push(node);
167 continue;
168 }
169 let mut cursor = node.walk();
170 pending.extend(node.named_children(&mut cursor));
171 }
172 nodes.sort_by_key(Node::start_byte);
173 nodes
174 .into_iter()
175 .flat_map(|node| {
176 let (scope_start_byte, scope_end_byte, function_scoped) =
177 python_import_binding_scope(node, source.len());
178 python_import_infos_from_node(node, source)
179 .into_iter()
180 .filter_map(move |import| {
181 let path = import.path.as_ref()?;
182 let details = python_import_details(&import)?;
183 match details {
184 PythonImportDetails::Import { module, alias } => {
185 Some(PythonImportBinding {
186 start_byte: node.start_byte(),
187 scope_start_byte,
188 scope_end_byte,
189 function_scoped,
190 local_name: alias.or_else(|| path.segments.first().cloned())?,
191 qualified_name: module,
192 })
193 }
194 PythonImportDetails::FromImport {
195 module,
196 name,
197 alias,
198 wildcard: false,
199 } => Some(PythonImportBinding {
200 start_byte: node.start_byte(),
201 scope_start_byte,
202 scope_end_byte,
203 function_scoped,
204 local_name: alias.unwrap_or(name.clone()),
205 qualified_name: format!("{module}.{name}"),
206 }),
207 PythonImportDetails::FromImport { wildcard: true, .. } => None,
208 }
209 })
210 })
211 .collect()
212}
213
214fn python_import_binding_scope(node: Node<'_>, source_len: usize) -> (usize, usize, bool) {
215 let mut parent = node.parent();
216 while let Some(scope) = parent {
217 if matches!(scope.kind(), "function_definition" | "lambda") {
218 return (scope.start_byte(), scope.end_byte(), true);
219 }
220 parent = scope.parent();
221 }
222 (0, source_len, false)
223}
224
225#[cfg(test)]
226mod tests {
227 use super::*;
228 use brokk_bifrost_core::analyzer::usages::model::ImportBinder;
229
230 fn replacement_for(source: &str, binder: &ImportBinder) -> Option<PythonModuleReplacement> {
231 let tree = parse_python_tree(source).expect("valid Python fixture");
232 let root = tree.root_node();
233 let mut replacement = None;
234 let mut cursor = root.walk();
235 for statement in root.named_children(&mut cursor) {
236 let statement = if statement.kind() == "expression_statement" {
237 statement.named_child(0).expect("fixture expression")
238 } else {
239 statement
240 };
241 if statement.kind() == "assignment"
242 && let Some(next) =
243 module_replacement_from_assignment(statement, source, &binder.bindings)
244 && replacement.replace(next).is_some()
245 {
246 return None;
247 }
248 }
249 replacement
250 }
251
252 #[test]
253 fn module_replacement_requires_exact_structured_import_bindings() {
254 let source = r#"import sys as _sys
255from routes.contacts import contacts_routes as _canonical
256
257_sys.modules[__name__] = _canonical
258"#;
259 let mut binder = ImportBinder::empty();
260 binder.bindings.insert(
261 "_sys".to_string(),
262 ImportBinding {
263 module_specifier: "sys".to_string(),
264 namespace_imported_module: Some("sys".to_string()),
265 kind: ImportKind::Namespace,
266 imported_name: None,
267 },
268 );
269 binder.bindings.insert(
270 "_canonical".to_string(),
271 ImportBinding {
272 module_specifier: "routes.contacts.contacts_routes".to_string(),
273 namespace_imported_module: None,
274 kind: ImportKind::Namespace,
275 imported_name: None,
276 },
277 );
278
279 assert_eq!(
280 replacement_for(source, &binder),
281 Some(PythonModuleReplacement {
282 target_module: "routes.contacts.contacts_routes".to_string(),
283 })
284 );
285
286 for near_miss in [
287 "cache.modules[__name__] = _canonical\n",
288 "_sys.modules[module_name] = _canonical\n",
289 "_sys.modules[__name__] = build()\n",
290 "def replace():\n _sys.modules[__name__] = _canonical\n",
291 ] {
292 assert_eq!(
293 replacement_for(near_miss, &binder),
294 None,
295 "near miss must not replace module identity: {near_miss:?}"
296 );
297 }
298 }
299}
300
301pub fn module_replacement_of(
302 python: &dyn PythonSource,
303 file: &ProjectFile,
304 source: &str,
305) -> Option<PythonModuleReplacement> {
306 let tree = parse_python_tree(source)?;
307 let root = tree.root_node();
308 let mut bindings: HashMap<String, ImportBinding> = HashMap::default();
309 let mut replacement = None;
310 let mut cursor = root.walk();
311 for statement in root.named_children(&mut cursor) {
312 let statement = if statement.kind() == "expression_statement" {
313 let Some(expression) = statement.named_child(0) else {
314 continue;
315 };
316 expression
317 } else {
318 statement
319 };
320 match statement.kind() {
321 "import_statement" | "import_from_statement" => {
322 let imports = python_import_infos_from_node(statement, source);
323 bindings.extend(import_binder_from_imports(python, file, &imports).bindings);
324 }
325 "assignment" => {
326 if let Some(next) = module_replacement_from_assignment(statement, source, &bindings)
327 && replacement.replace(next).is_some()
328 {
329 return None;
330 }
331 remove_direct_scope_bindings(statement, source, &mut bindings);
332 }
333 _ => remove_direct_scope_bindings(statement, source, &mut bindings),
334 }
335 }
336 replacement
337}
338
339pub fn resolve_import_bindings(
340 python: &dyn PythonSource,
341 file: &ProjectFile,
342) -> HashMap<String, CodeUnit> {
343 let imports = python.import_info_of(file);
344 let mut bindings = HashMap::default();
345 for resolved in resolve_imports_batched(python, file, &imports) {
346 for (binding, code_unit) in resolved {
347 bindings.insert(binding, code_unit);
348 }
349 }
350 bindings
351}
352
353pub fn resolve_imports_batched(
360 python: &dyn PythonSource,
361 file: &ProjectFile,
362 imports: &[ImportInfo],
363) -> Vec<Vec<(String, CodeUnit)>> {
364 let primary_fqns: Vec<Option<String>> = imports
365 .iter()
366 .map(|import| primary_module_fqn(file, import))
367 .collect();
368 let to_resolve: Vec<String> = primary_fqns.iter().flatten().cloned().collect();
369 let mut batch_results = resolve_module_code_units_batch(python, &to_resolve).into_iter();
370
371 imports
372 .iter()
373 .zip(primary_fqns.iter())
374 .map(|(import, primary_fqn)| {
375 let hint = primary_fqn.as_ref().map(|_| batch_results.next().unwrap());
376 resolve_import_with_hint(python, file, import, hint.as_ref())
377 })
378 .collect()
379}
380
381fn primary_module_fqn(file: &ProjectFile, import: &ImportInfo) -> Option<String> {
385 match python_import_details(import)? {
386 PythonImportDetails::Import { module, alias } => Some(python_namespace_binding_module(
387 import,
388 alias.as_deref(),
389 &module,
390 )),
391 PythonImportDetails::FromImport {
392 module,
393 name,
394 wildcard,
395 ..
396 } => {
397 if wildcard {
398 return None;
399 }
400 let resolved_module = if module.starts_with('.') {
401 resolve_python_relative_module(file, &module)
402 } else {
403 Some(module)
404 };
405 resolved_module.map(|resolved_module| format!("{resolved_module}.{name}"))
406 }
407 }
408}
409
410pub fn resolve_import(
411 python: &dyn PythonSource,
412 file: &ProjectFile,
413 import: &ImportInfo,
414) -> Vec<(String, CodeUnit)> {
415 resolve_import_with_hint(python, file, import, None)
416}
417
418fn resolve_import_with_hint(
421 python: &dyn PythonSource,
422 file: &ProjectFile,
423 import: &ImportInfo,
424 primary_hint: Option<&Option<CodeUnit>>,
425) -> Vec<(String, CodeUnit)> {
426 if let Some(details) = python_import_details(import) {
427 match details {
428 PythonImportDetails::Import { module, alias } => {
429 let binding = python_namespace_binding_name(import, alias.as_deref(), &module);
430 let bound_module =
431 python_namespace_binding_module(import, alias.as_deref(), &module);
432 let resolved = match primary_hint {
433 Some(hint) => hint.clone(),
434 None => resolve_module_code_unit(python, &bound_module),
435 };
436 if let Some(module_code_unit) = resolved {
437 return vec![(binding, module_code_unit)];
438 }
439 }
440 PythonImportDetails::FromImport {
441 module,
442 name,
443 alias,
444 wildcard,
445 } => {
446 let resolved_module = if module.starts_with('.') {
447 resolve_python_relative_module(file, &module)
448 } else {
449 Some(module)
450 };
451 let Some(resolved_module) = resolved_module else {
452 return Vec::new();
453 };
454 if wildcard {
455 return public_declarations_in_module(python, &resolved_module)
456 .into_iter()
457 .map(|code_unit| (code_unit.identifier().to_string(), code_unit))
458 .collect();
459 }
460
461 let binding = alias.clone().unwrap_or_else(|| name.clone());
462 let module_candidate = format!("{resolved_module}.{name}");
463 let resolved = match primary_hint {
464 Some(hint) => hint.clone(),
465 None => resolve_module_code_unit(python, &module_candidate),
466 };
467 if let Some(code_unit) = resolved {
468 return vec![(binding, code_unit)];
469 }
470 let exported = resolve_exported_name_from_module(python, &resolved_module, &name);
471 if !exported.is_empty() {
472 return exported
473 .into_iter()
474 .map(|code_unit| (binding.clone(), code_unit))
475 .collect();
476 }
477 let definitions: Vec<_> = python.definitions(&module_candidate).collect();
478 if !definitions.is_empty() {
479 return definitions
480 .into_iter()
481 .map(|code_unit| (binding.clone(), code_unit))
482 .collect();
483 }
484 let package_candidate: Vec<_> = python
485 .definitions(&format!("{resolved_module}.{name}"))
486 .collect();
487 if !package_candidate.is_empty() {
488 return package_candidate
489 .into_iter()
490 .map(|code_unit| (binding.clone(), code_unit))
491 .collect();
492 }
493 }
494 }
495 }
496 Vec::new()
497}
498
499pub fn resolve_exported_fqn(python: &dyn PythonSource, fqn: &str) -> Vec<CodeUnit> {
500 let Some((module, name)) = fqn.rsplit_once('.') else {
501 return Vec::new();
502 };
503 resolve_exported_name_from_module(python, module, name)
504}
505
506fn resolve_direct_named_exported_fqn(
511 python: &dyn PythonSource,
512 fqn: &str,
513) -> Option<Vec<CodeUnit>> {
514 let (module, name) = fqn.rsplit_once('.')?;
515 let mut results = Vec::new();
516 let mut queue = VecDeque::from([(module.to_string(), name.to_string())]);
517 let mut visited = HashSet::default();
518
519 while let Some((module, export_name)) = queue.pop_front() {
520 if !visited.insert((module.clone(), export_name.clone())) {
521 continue;
522 }
523 let module_unit = resolve_module_code_unit(python, &module)?;
524 let file = module_unit.source();
525 let local = local_export_declarations(python, file, &export_name);
526 let binder = python.import_binder_of(file);
527 let binding = binder.bindings.get(&export_name);
528 if !local.is_empty() && binding.is_some() {
529 return None;
530 }
531 if !local.is_empty() {
532 results.extend(local);
533 continue;
534 }
535 let binding = binding?;
536 if binding.kind != ImportKind::Named {
537 return None;
538 }
539 let imported_name = binding.imported_name.as_ref()?;
540 queue.push_back((binding.module_specifier.clone(), imported_name.clone()));
541 }
542
543 results.sort_by(|left, right| {
544 left.source()
545 .cmp(right.source())
546 .then_with(|| left.fq_name().cmp(&right.fq_name()))
547 });
548 results.dedup();
549 (!results.is_empty()).then_some(results)
550}
551
552pub fn resolve_fqn_candidates(
557 python: &dyn PythonSource,
558 fqn: &str,
559 exact: impl FnOnce(&str) -> Vec<CodeUnit>,
560) -> Vec<CodeUnit> {
561 if let Some(candidates) = resolve_direct_named_exported_fqn(python, fqn) {
562 return candidates;
563 }
564 let candidates = resolve_exported_fqn(python, fqn);
565 if !candidates.is_empty() {
566 return candidates;
567 }
568 exact(fqn)
569}
570
571fn resolve_exported_name_from_module(
572 python: &dyn PythonSource,
573 module: &str,
574 name: &str,
575) -> Vec<CodeUnit> {
576 let Some(module_unit) = resolve_module_code_unit(python, module) else {
577 return Vec::new();
578 };
579 resolve_exported_name(python, module_unit.source(), name)
580}
581
582fn resolve_exported_name(
583 python: &dyn PythonSource,
584 module_file: &ProjectFile,
585 name: &str,
586) -> Vec<CodeUnit> {
587 let mut results = Vec::new();
588 let mut queue = VecDeque::from([(module_file.clone(), name.to_string())]);
589 let mut visited = HashSet::default();
590
591 while let Some((file, export_name)) = queue.pop_front() {
592 if !visited.insert((file.clone(), export_name.clone())) {
593 continue;
594 }
595
596 let index = python.export_index_of(&file);
597 if let Some(entry) = index.exports_by_name.get(&export_name) {
598 match entry {
599 ExportEntry::Local { local_name } => {
600 results.extend(local_export_declarations(python, &file, local_name));
601 }
602 ExportEntry::ReexportedNamed {
603 module_specifier,
604 imported_name,
605 } => {
606 for target_file in
607 resolve_module_files_for_export(python, &file, module_specifier)
608 {
609 queue.push_back((target_file, imported_name.clone()));
610 }
611 }
612 ExportEntry::ReexportedModule { module_specifier } => {
613 results.extend(resolve_module_code_unit(python, module_specifier));
616 }
617 ExportEntry::Default { local_name } => {
618 if let Some(local_name) = local_name {
619 results.extend(local_export_declarations(python, &file, local_name));
620 }
621 }
622 }
623 continue;
624 }
625
626 if !export_name.starts_with('_') {
627 for star in &index.reexport_stars {
628 for target_file in
629 resolve_module_files_for_export(python, &file, &star.module_specifier)
630 {
631 queue.push_back((target_file, export_name.clone()));
632 }
633 }
634 }
635 }
636
637 results.sort_by(|left, right| {
638 left.source()
639 .cmp(right.source())
640 .then_with(|| left.fq_name().cmp(&right.fq_name()))
641 });
642 results.dedup();
643 results
644}
645
646fn local_export_declarations(
647 index: &dyn CodeUnitIndex,
648 file: &ProjectFile,
649 local_name: &str,
650) -> Vec<CodeUnit> {
651 index
652 .top_level_declarations(file)
653 .into_iter()
654 .filter(|unit| {
655 unit.identifier() == local_name
656 && index
657 .parent_of(unit)
658 .is_some_and(|parent| parent.is_module() && parent.source() == file)
659 })
660 .collect()
661}
662
663fn resolve_module_files_for_export(
664 python: &dyn PythonSource,
665 importing_file: &ProjectFile,
666 module_specifier: &str,
667) -> Vec<ProjectFile> {
668 let resolved_module = if module_specifier.starts_with('.') {
669 resolve_python_relative_module(importing_file, module_specifier)
670 } else {
671 Some(module_specifier.to_string())
672 };
673 let Some(resolved_module) = resolved_module else {
674 return Vec::new();
675 };
676 resolve_module_code_unit(python, &resolved_module)
680 .map(|unit| vec![unit.source().clone()])
681 .unwrap_or_default()
682}
683
684pub fn extract_package_from_python_wildcard(import: &ImportInfo) -> Option<String> {
685 let details = python_import_details(import)?;
686 match details {
687 PythonImportDetails::FromImport {
688 module, wildcard, ..
689 } if wildcard => Some(module),
690 _ => None,
691 }
692}
693
694#[derive(Debug, Clone)]
695pub enum PythonImportDetails {
696 Import {
697 module: String,
698 alias: Option<String>,
699 },
700 FromImport {
701 module: String,
702 name: String,
703 alias: Option<String>,
704 wildcard: bool,
705 },
706}
707
708pub fn python_import_infos_from_node(node: Node<'_>, source: &str) -> Vec<ImportInfo> {
709 match node.kind() {
710 "import_statement" => python_namespace_import_infos(node, source),
711 "import_from_statement" => python_from_import_infos(node, source),
712 _ => Vec::new(),
713 }
714}
715
716pub fn python_import_details(import: &ImportInfo) -> Option<PythonImportDetails> {
717 let path = import.path.as_ref()?;
718 match path.kind? {
719 StructuredImportPathKind::Namespace => Some(PythonImportDetails::Import {
720 module: join_python_import_segments(&path.segments),
721 alias: import.alias.clone(),
722 }),
723 StructuredImportPathKind::StaticMember => None,
725 StructuredImportPathKind::ImportFrom => {
726 let (name, module_segments) = if import.is_wildcard {
727 ("*".to_string(), path.segments.as_slice())
728 } else {
729 let (name, module_segments) = path.segments.split_last()?;
730 (name.clone(), module_segments)
731 };
732 Some(PythonImportDetails::FromImport {
733 module: join_python_import_segments(module_segments),
734 name,
735 alias: import.alias.clone(),
736 wildcard: import.is_wildcard,
737 })
738 }
739 }
740}
741
742fn python_namespace_import_infos(node: Node<'_>, source: &str) -> Vec<ImportInfo> {
743 let mut infos = Vec::new();
744 let mut cursor = node.walk();
745 for imported in node.children_by_field_name("name", &mut cursor) {
746 let (module_node, alias_node) = if imported.kind() == "aliased_import" {
747 let Some(name) = imported.child_by_field_name("name") else {
748 continue;
749 };
750 (name, imported.child_by_field_name("alias"))
751 } else {
752 (imported, None)
753 };
754 let alias = alias_node
755 .map(|alias| py_node_text(alias, source).trim().to_string())
756 .filter(|alias| !alias.is_empty());
757 let segments = python_path_segments(module_node, source);
758 if segments.is_empty() {
759 continue;
760 }
761 let binder_span = alias
764 .is_some()
765 .then_some(alias_node)
766 .flatten()
767 .or_else(|| python_first_segment_node(module_node))
768 .map(brokk_bifrost_core::analyzer::common::node_span);
769 let module = join_python_import_segments(&segments);
770 let identifier = alias.clone().or_else(|| segments.first().cloned());
771 infos.push(ImportInfo {
772 raw_snippet: if let Some(alias) = &alias {
773 format!("import {module} as {alias}")
774 } else {
775 format!("import {module}")
776 },
777 is_wildcard: false,
778 is_global: false,
779 identifier,
780 alias,
781 path: Some(StructuredImportPath {
782 segments,
783 kind: Some(StructuredImportPathKind::Namespace),
784 lexical_prefixes: Vec::new(),
785 lexical_scopes: Vec::new(),
786 declaration_start_byte: node.start_byte(),
787 }),
788 binder_span,
789 });
790 }
791 infos
792}
793
794fn python_from_import_infos(node: Node<'_>, source: &str) -> Vec<ImportInfo> {
795 let Some(module_node) = node.child_by_field_name("module_name") else {
796 return Vec::new();
797 };
798 let module_segments = python_module_segments(module_node, source);
799 if module_segments.is_empty() {
800 return Vec::new();
801 }
802
803 let mut infos = Vec::new();
804 let has_wildcard_import = {
805 let mut cursor = node.walk();
806 node.named_children(&mut cursor)
807 .any(|child| child.kind() == "wildcard_import")
808 };
809 let mut cursor = node.walk();
810 let imported_names: Vec<_> = node.children_by_field_name("name", &mut cursor).collect();
811 if has_wildcard_import {
812 let module = join_python_import_segments(&module_segments);
813 infos.push(ImportInfo {
814 raw_snippet: format!("from {module} import *"),
815 is_wildcard: true,
816 is_global: false,
817 identifier: None,
818 alias: None,
819 path: Some(StructuredImportPath {
820 segments: module_segments,
821 kind: Some(StructuredImportPathKind::ImportFrom),
822 lexical_prefixes: Vec::new(),
823 lexical_scopes: Vec::new(),
824 declaration_start_byte: node.start_byte(),
825 }),
826 binder_span: None,
827 });
828 return infos;
829 }
830 if imported_names.is_empty() {
831 return infos;
832 }
833
834 for imported in imported_names {
835 let (name_node, alias_node) = if imported.kind() == "aliased_import" {
836 let Some(name) = imported.child_by_field_name("name") else {
837 continue;
838 };
839 (name, imported.child_by_field_name("alias"))
840 } else {
841 (imported, None)
842 };
843 let alias = alias_node
844 .map(|alias| py_node_text(alias, source).trim().to_string())
845 .filter(|alias| !alias.is_empty());
846 let name_segments = python_path_segments(name_node, source);
847 if name_segments.is_empty() {
848 continue;
849 }
850 let binder_span = alias
853 .is_some()
854 .then_some(alias_node)
855 .flatten()
856 .or_else(|| {
857 (name_segments.len() == 1)
858 .then(|| python_first_segment_node(name_node))
859 .flatten()
860 })
861 .map(brokk_bifrost_core::analyzer::common::node_span);
862 let imported_name = join_python_import_segments(&name_segments);
863 let mut segments = module_segments.clone();
864 segments.extend(name_segments);
865 let module = join_python_import_segments(&module_segments);
866 infos.push(ImportInfo {
867 raw_snippet: if let Some(alias) = &alias {
868 format!("from {module} import {imported_name} as {alias}")
869 } else {
870 format!("from {module} import {imported_name}")
871 },
872 is_wildcard: false,
873 is_global: false,
874 identifier: Some(alias.clone().unwrap_or_else(|| imported_name.clone())),
875 alias,
876 path: Some(StructuredImportPath {
877 segments,
878 kind: Some(StructuredImportPathKind::ImportFrom),
879 lexical_prefixes: Vec::new(),
880 lexical_scopes: Vec::new(),
881 declaration_start_byte: node.start_byte(),
882 }),
883 binder_span,
884 });
885 }
886 infos
887}
888
889fn python_module_segments(module: Node<'_>, source: &str) -> Vec<String> {
890 if module.kind() == "relative_import" {
891 let mut cursor = module.walk();
892 let mut prefix = String::new();
893 let mut path_node = None;
894 for child in module.named_children(&mut cursor) {
895 match child.kind() {
896 "import_prefix" if prefix.is_empty() => {
897 prefix = py_node_text(child, source).trim().to_string();
898 }
899 "dotted_name" if path_node.is_none() => {
900 path_node = Some(child);
901 }
902 _ => {}
903 }
904 }
905 let mut segments = path_node
906 .map(|path| python_path_segments(path, source))
907 .unwrap_or_default();
908 if !prefix.is_empty() {
909 if let Some(first) = segments.first_mut() {
910 first.insert_str(0, &prefix);
911 } else {
912 segments.push(prefix);
913 }
914 }
915 return segments;
916 }
917 python_path_segments(module, source)
918}
919
920fn python_first_segment_node(node: Node<'_>) -> Option<Node<'_>> {
924 match node.kind() {
925 "identifier" => Some(node),
926 "dotted_name" => {
927 let mut cursor = node.walk();
928 node.named_children(&mut cursor)
929 .find(|child| child.kind() == "identifier")
930 }
931 _ => None,
932 }
933}
934
935fn python_path_segments(node: Node<'_>, source: &str) -> Vec<String> {
936 match node.kind() {
937 "identifier" => vec![py_node_text(node, source).trim().to_string()],
938 "dotted_name" => {
939 let mut segments = Vec::new();
940 let mut cursor = node.walk();
941 for child in node.named_children(&mut cursor) {
942 segments.extend(python_path_segments(child, source));
943 }
944 segments
945 }
946 _ => {
947 let mut segments = Vec::new();
948 let mut cursor = node.walk();
949 for child in node.named_children(&mut cursor) {
950 segments.extend(python_path_segments(child, source));
951 }
952 segments
953 }
954 }
955}
956
957fn join_python_import_segments(segments: &[String]) -> String {
958 let Some((first, rest)) = segments.split_first() else {
959 return String::new();
960 };
961 if first.starts_with('.') && !rest.is_empty() {
962 format!("{first}.{}", rest.join("."))
963 } else {
964 segments.join(".")
965 }
966}
967
968pub fn python_namespace_binding_name(
969 import: &ImportInfo,
970 alias: Option<&str>,
971 module: &str,
972) -> String {
973 import
974 .identifier
975 .clone()
976 .or_else(|| alias.map(str::to_string))
977 .unwrap_or_else(|| module.to_string())
978}
979
980pub fn python_namespace_binding_module(
981 import: &ImportInfo,
982 alias: Option<&str>,
983 module: &str,
984) -> String {
985 if alias.is_some() {
986 return module.to_string();
987 }
988 import
989 .path
990 .as_ref()
991 .and_then(|path| path.segments.first().cloned())
992 .unwrap_or_else(|| module.to_string())
993}
994
995pub fn resolve_python_relative_module(
996 source_file: &ProjectFile,
997 module_expr: &str,
998) -> Option<String> {
999 resolve_python_relative_module_from_package(&python_current_package(source_file), module_expr)
1000}
1001
1002pub fn resolve_python_relative_module_from_package(
1006 current_package: &str,
1007 module_expr: &str,
1008) -> Option<String> {
1009 let level = module_expr.chars().take_while(|ch| *ch == '.').count();
1010 let suffix = module_expr[level..].trim_matches('.');
1011 let mut parts: Vec<_> = current_package
1012 .split('.')
1013 .filter(|part| !part.is_empty())
1014 .map(str::to_string)
1015 .collect();
1016 if level == 0 {
1017 return Some(module_expr.to_string());
1018 }
1019 if level > 0 {
1020 if level - 1 > parts.len() {
1021 return None;
1022 }
1023 parts.truncate(parts.len() - (level - 1));
1024 }
1025 if !suffix.is_empty() {
1026 parts.extend(suffix.split('.').map(str::to_string));
1027 }
1028 Some(parts.join("."))
1029}
1030
1031fn python_current_package(source_file: &ProjectFile) -> String {
1032 let module = python_module_name(source_file);
1033 if source_file
1034 .rel_path()
1035 .file_name()
1036 .and_then(|name| name.to_str())
1037 == Some("__init__.py")
1038 {
1039 module
1040 } else {
1041 module
1042 .rsplit_once('.')
1043 .map(|(package, _)| package.to_string())
1044 .unwrap_or_default()
1045 }
1046}