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