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    pub local_name: String,
141    pub qualified_name: String,
142}
143
144/// Return source-ordered, parser-derived local bindings for explicit Python
145/// imports. The byte offset lets callers select the last visible binding rather
146/// than treating a whole document as one unordered import scope.
147pub 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
344/// Resolves every import in `imports` (`file`'s own imports), batching each import's primary
345/// module FQN lookup (see `primary_module_fqn`) into one store transaction instead of one per
346/// import. Shared by `resolve_import_bindings` and `resolve_import_target_files`, the two per-file
347/// "resolve everything" entry points -- both are called once per candidate file by the usages
348/// candidate walker, so unbatched resolution here means one store transaction per import times
349/// every file in the workspace.
350pub 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
372/// The module FQN `resolve_import`'s fast path checks first, if any -- must stay in sync with the
373/// two `resolve_module_code_unit` call sites in `resolve_import_with_hint` below, since it's what
374/// lets `resolve_import_target_files` batch-resolve them ahead of the serial fallback logic.
375fn 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
409/// `primary_hint`, when `Some`, is the already-resolved result of this import's primary module FQN
410/// (see `primary_module_fqn`) so the batched caller doesn't pay for a second lookup of the same FQN.
411fn 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
497/// Resolve an unambiguous chain of explicit named reexports without
498/// constructing export indexes for each intermediate module. Star exports,
499/// shadowing, and every other ambiguous shape return `None` so callers can
500/// use the complete, source-order-aware export resolver below.
501fn 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
547/// Resolve a Python FQN with the cheapest semantically complete tier that
548/// can answer it. The direct reexport walk handles only proven,
549/// collision-free chains; ambiguous shapes use the ordered export index,
550/// and the exact lookup remains the final fallback for non-export symbols.
551pub 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    // Tree-sitter tells us the import syntax, but module-to-file resolution
662    // is analyzer state. Use the prebuilt module code-unit map here instead
663    // of the usage index so interactive definition lookup stays lightweight.
664    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        // Python has no static imports; the variant belongs to Java.
709        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        // `import a.b` binds `a`: the first segment's own token. A renamed
747        // import binds its alias token instead.
748        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        // `from m import x` binds `x`'s own token; a rename binds the alias
834        // token. A multi-segment imported name binds no single token.
835        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
902/// The token that spells a path's first segment: the identifier itself, or a
903/// dotted name's first identifier. `None` when the shape has no leading
904/// identifier token of its own (e.g. a relative-import prefix).
905fn 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}