Skip to main content

brokk_bifrost_python/
imports.rs

1//! Python import syntax, binding, and FQN resolution.
2//!
3//! Everything here is a free function over [`PythonSource`]; the
4//! `ImportAnalysisProvider` impl that memoizes the per-file results stays on
5//! `PythonAnalyzer` in `analyzer/python/imports.rs`.
6
7use 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
28/// Recognize the compatibility-shim idiom that replaces the current module
29/// object with an imported workspace module:
30///
31/// `sys.modules[__name__] = imported_module`
32///
33/// Every component is resolved from tree-sitter fields and the import binder;
34/// similarly spelled attributes or locals do not create an alias edge.
35fn 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
111/// Parse the structured import facts for a Python document without executing
112/// it. LSP model binding uses the same AST-derived representation as the
113/// analyzer so external APIs are never selected by a terminal-name scan.
114pub 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
151/// Return source-ordered, parser-derived local bindings for explicit Python
152/// imports. The byte offset lets callers select the last visible binding rather
153/// than treating a whole document as one unordered import scope.
154pub 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
353/// Resolves every import in `imports` (`file`'s own imports), batching each import's primary
354/// module FQN lookup (see `primary_module_fqn`) into one store transaction instead of one per
355/// import. Shared by `resolve_import_bindings` and `resolve_import_target_files`, the two per-file
356/// "resolve everything" entry points -- both are called once per candidate file by the usages
357/// candidate walker, so unbatched resolution here means one store transaction per import times
358/// every file in the workspace.
359pub 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
381/// The module FQN `resolve_import`'s fast path checks first, if any -- must stay in sync with the
382/// two `resolve_module_code_unit` call sites in `resolve_import_with_hint` below, since it's what
383/// lets `resolve_import_target_files` batch-resolve them ahead of the serial fallback logic.
384fn 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
418/// `primary_hint`, when `Some`, is the already-resolved result of this import's primary module FQN
419/// (see `primary_module_fqn`) so the batched caller doesn't pay for a second lookup of the same FQN.
420fn 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
506/// Resolve an unambiguous chain of explicit named reexports without
507/// constructing export indexes for each intermediate module. Star exports,
508/// shadowing, and every other ambiguous shape return `None` so callers can
509/// use the complete, source-order-aware export resolver below.
510fn 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 = python
526            .top_level_declarations(file)
527            .into_iter()
528            .filter(|unit| unit.identifier() == export_name)
529            .collect::<Vec<_>>();
530        let binder = python.import_binder_of(file);
531        let binding = binder.bindings.get(&export_name);
532        if !local.is_empty() && binding.is_some() {
533            return None;
534        }
535        if !local.is_empty() {
536            results.extend(local);
537            continue;
538        }
539        let binding = binding?;
540        if binding.kind != ImportKind::Named {
541            return None;
542        }
543        let imported_name = binding.imported_name.as_ref()?;
544        queue.push_back((binding.module_specifier.clone(), imported_name.clone()));
545    }
546
547    results.sort_by(|left, right| {
548        left.source()
549            .cmp(right.source())
550            .then_with(|| left.fq_name().cmp(&right.fq_name()))
551    });
552    results.dedup();
553    (!results.is_empty()).then_some(results)
554}
555
556/// Resolve a Python FQN with the cheapest semantically complete tier that
557/// can answer it. The direct reexport walk handles only proven,
558/// collision-free chains; ambiguous shapes use the ordered export index,
559/// and the exact lookup remains the final fallback for non-export symbols.
560pub fn resolve_fqn_candidates(
561    python: &dyn PythonSource,
562    fqn: &str,
563    exact: impl FnOnce(&str) -> Vec<CodeUnit>,
564) -> Vec<CodeUnit> {
565    if let Some(candidates) = resolve_direct_named_exported_fqn(python, fqn) {
566        return candidates;
567    }
568    let candidates = resolve_exported_fqn(python, fqn);
569    if !candidates.is_empty() {
570        return candidates;
571    }
572    exact(fqn)
573}
574
575fn resolve_exported_name_from_module(
576    python: &dyn PythonSource,
577    module: &str,
578    name: &str,
579) -> Vec<CodeUnit> {
580    let Some(module_unit) = resolve_module_code_unit(python, module) else {
581        return Vec::new();
582    };
583    resolve_exported_name(python, module_unit.source(), name)
584}
585
586fn resolve_exported_name(
587    python: &dyn PythonSource,
588    module_file: &ProjectFile,
589    name: &str,
590) -> Vec<CodeUnit> {
591    let mut results = Vec::new();
592    let mut queue = VecDeque::from([(module_file.clone(), name.to_string())]);
593    let mut visited = HashSet::default();
594
595    while let Some((file, export_name)) = queue.pop_front() {
596        if !visited.insert((file.clone(), export_name.clone())) {
597            continue;
598        }
599
600        let index = python.export_index_of(&file);
601        if let Some(entry) = index.exports_by_name.get(&export_name) {
602            match entry {
603                ExportEntry::Local { local_name } => {
604                    results.extend(local_export_declarations(python, &file, local_name));
605                }
606                ExportEntry::ReexportedNamed {
607                    module_specifier,
608                    imported_name,
609                } => {
610                    for target_file in
611                        resolve_module_files_for_export(python, &file, module_specifier)
612                    {
613                        queue.push_back((target_file, imported_name.clone()));
614                    }
615                }
616                ExportEntry::ReexportedModule { module_specifier } => {
617                    // Terminal: the export *is* the module, so the walk stops
618                    // here instead of looking the name up inside it.
619                    results.extend(resolve_module_code_unit(python, module_specifier));
620                }
621                ExportEntry::Default { local_name } => {
622                    if let Some(local_name) = local_name {
623                        results.extend(local_export_declarations(python, &file, local_name));
624                    }
625                }
626            }
627            continue;
628        }
629
630        if !export_name.starts_with('_') {
631            for star in &index.reexport_stars {
632                for target_file in
633                    resolve_module_files_for_export(python, &file, &star.module_specifier)
634                {
635                    queue.push_back((target_file, export_name.clone()));
636                }
637            }
638        }
639    }
640
641    results.sort_by(|left, right| {
642        left.source()
643            .cmp(right.source())
644            .then_with(|| left.fq_name().cmp(&right.fq_name()))
645    });
646    results.dedup();
647    results
648}
649
650fn local_export_declarations(
651    index: &dyn CodeUnitIndex,
652    file: &ProjectFile,
653    local_name: &str,
654) -> Vec<CodeUnit> {
655    index
656        .top_level_declarations(file)
657        .into_iter()
658        .filter(|unit| unit.identifier() == local_name)
659        .collect()
660}
661
662fn resolve_module_files_for_export(
663    python: &dyn PythonSource,
664    importing_file: &ProjectFile,
665    module_specifier: &str,
666) -> Vec<ProjectFile> {
667    let resolved_module = if module_specifier.starts_with('.') {
668        resolve_python_relative_module(importing_file, module_specifier)
669    } else {
670        Some(module_specifier.to_string())
671    };
672    let Some(resolved_module) = resolved_module else {
673        return Vec::new();
674    };
675    // Tree-sitter tells us the import syntax, but module-to-file resolution
676    // is analyzer state. Use the prebuilt module code-unit map here instead
677    // of the usage index so interactive definition lookup stays lightweight.
678    resolve_module_code_unit(python, &resolved_module)
679        .map(|unit| vec![unit.source().clone()])
680        .unwrap_or_default()
681}
682
683pub fn extract_package_from_python_wildcard(import: &ImportInfo) -> Option<String> {
684    let details = python_import_details(import)?;
685    match details {
686        PythonImportDetails::FromImport {
687            module, wildcard, ..
688        } if wildcard => Some(module),
689        _ => None,
690    }
691}
692
693#[derive(Debug, Clone)]
694pub enum PythonImportDetails {
695    Import {
696        module: String,
697        alias: Option<String>,
698    },
699    FromImport {
700        module: String,
701        name: String,
702        alias: Option<String>,
703        wildcard: bool,
704    },
705}
706
707pub fn python_import_infos_from_node(node: Node<'_>, source: &str) -> Vec<ImportInfo> {
708    match node.kind() {
709        "import_statement" => python_namespace_import_infos(node, source),
710        "import_from_statement" => python_from_import_infos(node, source),
711        _ => Vec::new(),
712    }
713}
714
715pub fn python_import_details(import: &ImportInfo) -> Option<PythonImportDetails> {
716    let path = import.path.as_ref()?;
717    match path.kind? {
718        StructuredImportPathKind::Namespace => Some(PythonImportDetails::Import {
719            module: join_python_import_segments(&path.segments),
720            alias: import.alias.clone(),
721        }),
722        // Python has no static imports; the variant belongs to Java.
723        StructuredImportPathKind::StaticMember => None,
724        StructuredImportPathKind::ImportFrom => {
725            let (name, module_segments) = if import.is_wildcard {
726                ("*".to_string(), path.segments.as_slice())
727            } else {
728                let (name, module_segments) = path.segments.split_last()?;
729                (name.clone(), module_segments)
730            };
731            Some(PythonImportDetails::FromImport {
732                module: join_python_import_segments(module_segments),
733                name,
734                alias: import.alias.clone(),
735                wildcard: import.is_wildcard,
736            })
737        }
738    }
739}
740
741fn python_namespace_import_infos(node: Node<'_>, source: &str) -> Vec<ImportInfo> {
742    let mut infos = Vec::new();
743    let mut cursor = node.walk();
744    for imported in node.children_by_field_name("name", &mut cursor) {
745        let (module_node, alias_node) = if imported.kind() == "aliased_import" {
746            let Some(name) = imported.child_by_field_name("name") else {
747                continue;
748            };
749            (name, imported.child_by_field_name("alias"))
750        } else {
751            (imported, None)
752        };
753        let alias = alias_node
754            .map(|alias| py_node_text(alias, source).trim().to_string())
755            .filter(|alias| !alias.is_empty());
756        let segments = python_path_segments(module_node, source);
757        if segments.is_empty() {
758            continue;
759        }
760        // `import a.b` binds `a`: the first segment's own token. A renamed
761        // import binds its alias token instead.
762        let binder_span = alias
763            .is_some()
764            .then_some(alias_node)
765            .flatten()
766            .or_else(|| python_first_segment_node(module_node))
767            .map(brokk_bifrost_core::analyzer::common::node_span);
768        let module = join_python_import_segments(&segments);
769        let identifier = alias.clone().or_else(|| segments.first().cloned());
770        infos.push(ImportInfo {
771            raw_snippet: if let Some(alias) = &alias {
772                format!("import {module} as {alias}")
773            } else {
774                format!("import {module}")
775            },
776            is_wildcard: false,
777            is_global: false,
778            identifier,
779            alias,
780            path: Some(StructuredImportPath {
781                segments,
782                kind: Some(StructuredImportPathKind::Namespace),
783                lexical_prefixes: Vec::new(),
784                lexical_scopes: Vec::new(),
785                declaration_start_byte: node.start_byte(),
786            }),
787            binder_span,
788        });
789    }
790    infos
791}
792
793fn python_from_import_infos(node: Node<'_>, source: &str) -> Vec<ImportInfo> {
794    let Some(module_node) = node.child_by_field_name("module_name") else {
795        return Vec::new();
796    };
797    let module_segments = python_module_segments(module_node, source);
798    if module_segments.is_empty() {
799        return Vec::new();
800    }
801
802    let mut infos = Vec::new();
803    let has_wildcard_import = {
804        let mut cursor = node.walk();
805        node.named_children(&mut cursor)
806            .any(|child| child.kind() == "wildcard_import")
807    };
808    let mut cursor = node.walk();
809    let imported_names: Vec<_> = node.children_by_field_name("name", &mut cursor).collect();
810    if has_wildcard_import {
811        let module = join_python_import_segments(&module_segments);
812        infos.push(ImportInfo {
813            raw_snippet: format!("from {module} import *"),
814            is_wildcard: true,
815            is_global: false,
816            identifier: None,
817            alias: None,
818            path: Some(StructuredImportPath {
819                segments: module_segments,
820                kind: Some(StructuredImportPathKind::ImportFrom),
821                lexical_prefixes: Vec::new(),
822                lexical_scopes: Vec::new(),
823                declaration_start_byte: node.start_byte(),
824            }),
825            binder_span: None,
826        });
827        return infos;
828    }
829    if imported_names.is_empty() {
830        return infos;
831    }
832
833    for imported in imported_names {
834        let (name_node, alias_node) = if imported.kind() == "aliased_import" {
835            let Some(name) = imported.child_by_field_name("name") else {
836                continue;
837            };
838            (name, imported.child_by_field_name("alias"))
839        } else {
840            (imported, None)
841        };
842        let alias = alias_node
843            .map(|alias| py_node_text(alias, source).trim().to_string())
844            .filter(|alias| !alias.is_empty());
845        let name_segments = python_path_segments(name_node, source);
846        if name_segments.is_empty() {
847            continue;
848        }
849        // `from m import x` binds `x`'s own token; a rename binds the alias
850        // token. A multi-segment imported name binds no single token.
851        let binder_span = alias
852            .is_some()
853            .then_some(alias_node)
854            .flatten()
855            .or_else(|| {
856                (name_segments.len() == 1)
857                    .then(|| python_first_segment_node(name_node))
858                    .flatten()
859            })
860            .map(brokk_bifrost_core::analyzer::common::node_span);
861        let imported_name = join_python_import_segments(&name_segments);
862        let mut segments = module_segments.clone();
863        segments.extend(name_segments);
864        let module = join_python_import_segments(&module_segments);
865        infos.push(ImportInfo {
866            raw_snippet: if let Some(alias) = &alias {
867                format!("from {module} import {imported_name} as {alias}")
868            } else {
869                format!("from {module} import {imported_name}")
870            },
871            is_wildcard: false,
872            is_global: false,
873            identifier: Some(alias.clone().unwrap_or_else(|| imported_name.clone())),
874            alias,
875            path: Some(StructuredImportPath {
876                segments,
877                kind: Some(StructuredImportPathKind::ImportFrom),
878                lexical_prefixes: Vec::new(),
879                lexical_scopes: Vec::new(),
880                declaration_start_byte: node.start_byte(),
881            }),
882            binder_span,
883        });
884    }
885    infos
886}
887
888fn python_module_segments(module: Node<'_>, source: &str) -> Vec<String> {
889    if module.kind() == "relative_import" {
890        let mut cursor = module.walk();
891        let mut prefix = String::new();
892        let mut path_node = None;
893        for child in module.named_children(&mut cursor) {
894            match child.kind() {
895                "import_prefix" if prefix.is_empty() => {
896                    prefix = py_node_text(child, source).trim().to_string();
897                }
898                "dotted_name" if path_node.is_none() => {
899                    path_node = Some(child);
900                }
901                _ => {}
902            }
903        }
904        let mut segments = path_node
905            .map(|path| python_path_segments(path, source))
906            .unwrap_or_default();
907        if !prefix.is_empty() {
908            if let Some(first) = segments.first_mut() {
909                first.insert_str(0, &prefix);
910            } else {
911                segments.push(prefix);
912            }
913        }
914        return segments;
915    }
916    python_path_segments(module, source)
917}
918
919/// The token that spells a path's first segment: the identifier itself, or a
920/// dotted name's first identifier. `None` when the shape has no leading
921/// identifier token of its own (e.g. a relative-import prefix).
922fn python_first_segment_node(node: Node<'_>) -> Option<Node<'_>> {
923    match node.kind() {
924        "identifier" => Some(node),
925        "dotted_name" => {
926            let mut cursor = node.walk();
927            node.named_children(&mut cursor)
928                .find(|child| child.kind() == "identifier")
929        }
930        _ => None,
931    }
932}
933
934fn python_path_segments(node: Node<'_>, source: &str) -> Vec<String> {
935    match node.kind() {
936        "identifier" => vec![py_node_text(node, source).trim().to_string()],
937        "dotted_name" => {
938            let mut segments = Vec::new();
939            let mut cursor = node.walk();
940            for child in node.named_children(&mut cursor) {
941                segments.extend(python_path_segments(child, source));
942            }
943            segments
944        }
945        _ => {
946            let mut segments = Vec::new();
947            let mut cursor = node.walk();
948            for child in node.named_children(&mut cursor) {
949                segments.extend(python_path_segments(child, source));
950            }
951            segments
952        }
953    }
954}
955
956fn join_python_import_segments(segments: &[String]) -> String {
957    let Some((first, rest)) = segments.split_first() else {
958        return String::new();
959    };
960    if first.starts_with('.') && !rest.is_empty() {
961        format!("{first}.{}", rest.join("."))
962    } else {
963        segments.join(".")
964    }
965}
966
967pub fn python_namespace_binding_name(
968    import: &ImportInfo,
969    alias: Option<&str>,
970    module: &str,
971) -> String {
972    import
973        .identifier
974        .clone()
975        .or_else(|| alias.map(str::to_string))
976        .unwrap_or_else(|| module.to_string())
977}
978
979pub fn python_namespace_binding_module(
980    import: &ImportInfo,
981    alias: Option<&str>,
982    module: &str,
983) -> String {
984    if alias.is_some() {
985        return module.to_string();
986    }
987    import
988        .path
989        .as_ref()
990        .and_then(|path| path.segments.first().cloned())
991        .unwrap_or_else(|| module.to_string())
992}
993
994pub fn resolve_python_relative_module(
995    source_file: &ProjectFile,
996    module_expr: &str,
997) -> Option<String> {
998    let level = module_expr.chars().take_while(|ch| *ch == '.').count();
999    let suffix = module_expr[level..].trim_matches('.');
1000    let current_package = python_current_package(source_file);
1001    let mut parts: Vec<_> = current_package
1002        .split('.')
1003        .filter(|part| !part.is_empty())
1004        .map(str::to_string)
1005        .collect();
1006    if level == 0 {
1007        return Some(module_expr.to_string());
1008    }
1009    if level > 0 {
1010        if level - 1 > parts.len() {
1011            return None;
1012        }
1013        parts.truncate(parts.len() - (level - 1));
1014    }
1015    if !suffix.is_empty() {
1016        parts.extend(suffix.split('.').map(str::to_string));
1017    }
1018    Some(parts.join("."))
1019}
1020
1021fn python_current_package(source_file: &ProjectFile) -> String {
1022    let module = python_module_name(source_file);
1023    if source_file
1024        .rel_path()
1025        .file_name()
1026        .and_then(|name| name.to_str())
1027        == Some("__init__.py")
1028    {
1029        module
1030    } else {
1031        module
1032            .rsplit_once('.')
1033            .map(|(package, _)| package.to_string())
1034            .unwrap_or_default()
1035    }
1036}