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::query_token::QueryToken;
18use brokk_bifrost_core::analyzer::usages::model::{ExportEntry, ImportBinding, ImportKind};
19use brokk_bifrost_core::analyzer::{CodeUnit, CodeUnitIndex, ProjectFile};
20use brokk_bifrost_core::hash::{HashMap, HashSet};
21use std::collections::VecDeque;
22use tree_sitter::Node;
23
24#[derive(Debug, Clone, PartialEq, Eq)]
25pub struct PythonModuleReplacement {
26    pub target_module: String,
27}
28
29/// Return the absolute module named by the last direct assignment to `local`
30/// that dominates `reference` when that assignment has the exact structured
31/// shape `local = importlib_binding.import_module("module.name")`.
32///
33/// Assignments hidden in a branch do not dominate the reference. Any later
34/// direct binding of `local` invalidates an earlier result. The caller proves
35/// that the receiver of `import_module` is an in-scope namespace binding for
36/// Python's `importlib` module; this helper only interprets the syntax and
37/// statement order.
38pub(crate) fn imported_module_assignment_at(
39    reference: Node<'_>,
40    local: &str,
41    source: &str,
42    mut is_importlib_binding: impl FnMut(&str) -> bool,
43) -> Option<String> {
44    let suite = enclosing_execution_suite(reference)?;
45    let mut resolved = None;
46    let mut cursor = suite.walk();
47    for statement in suite.named_children(&mut cursor) {
48        if statement.start_byte() >= reference.start_byte()
49            || (statement.start_byte() <= reference.start_byte()
50                && reference.end_byte() <= statement.end_byte())
51        {
52            break;
53        }
54
55        let expression = if statement.kind() == "expression_statement" {
56            statement.named_child(0).unwrap_or(statement)
57        } else {
58            statement
59        };
60        if expression.kind() == "assignment"
61            && expression.child_by_field_name("left").is_some_and(|left| {
62                left.kind() == "identifier" && node_source_text(left, source) == local
63            })
64        {
65            resolved = expression
66                .child_by_field_name("right")
67                .and_then(|right| importlib_call_module(right, source, &mut is_importlib_binding));
68            continue;
69        }
70
71        let binds_local = python_direct_scope_bindings_bounded(statement, source, || true)
72            .expect("unbounded binding collection cannot be cancelled")
73            .into_iter()
74            .any(|binding| node_source_text(binding.declaration, source) == local);
75        if binds_local {
76            resolved = None;
77        }
78    }
79    resolved
80}
81
82fn enclosing_execution_suite(mut node: Node<'_>) -> Option<Node<'_>> {
83    while let Some(parent) = node.parent() {
84        if matches!(parent.kind(), "function_definition" | "lambda") {
85            return parent.child_by_field_name("body");
86        }
87        if parent.kind() == "module" {
88            return Some(parent);
89        }
90        node = parent;
91    }
92    None
93}
94
95fn importlib_call_module(
96    call: Node<'_>,
97    source: &str,
98    is_importlib_binding: &mut impl FnMut(&str) -> bool,
99) -> Option<String> {
100    if call.kind() != "call" {
101        return None;
102    }
103    let function = call.child_by_field_name("function")?;
104    if function.kind() != "attribute" {
105        return None;
106    }
107    let (object, attribute) = (
108        function.child_by_field_name("object")?,
109        function.child_by_field_name("attribute")?,
110    );
111    if object.kind() != "identifier"
112        || attribute.kind() != "identifier"
113        || node_source_text(attribute, source) != "import_module"
114        || !is_importlib_binding(node_source_text(object, source))
115    {
116        return None;
117    }
118
119    let arguments = call.child_by_field_name("arguments")?;
120    if arguments.named_child_count() != 1 {
121        return None;
122    }
123    let literal = arguments.named_child(0)?;
124    if literal.kind() != "string"
125        || literal
126            .parent()
127            .is_some_and(|parent| parent.kind() == "concatenated_string")
128    {
129        return None;
130    }
131    let mut content = None;
132    for index in 0..literal.named_child_count() {
133        let child = literal.named_child(index)?;
134        match child.kind() {
135            "string_start" | "string_end" => {}
136            "string_content" if content.is_none() && child.named_child_count() == 0 => {
137                content = Some(child);
138            }
139            _ => return None,
140        }
141    }
142    let content = content?;
143    let module = node_source_text(content, source);
144    (!module.is_empty() && !module.starts_with('.')).then(|| module.to_string())
145}
146
147/// Collect absolute module names from structurally exact
148/// `local = importlib.import_module("...")` assignments. This is a conservative
149/// candidate-file index: dominance and local shadowing are checked later at the
150/// reference site by [`imported_module_assignment_at`].
151pub(crate) fn literal_importlib_modules(
152    source: &str,
153    bindings: &HashMap<String, ImportBinding>,
154) -> HashSet<String> {
155    let Some(tree) = parse_python_tree(source) else {
156        return HashSet::default();
157    };
158    let mut modules = HashSet::default();
159    let mut stack = vec![tree.root_node()];
160    while let Some(node) = stack.pop() {
161        if node.kind() == "assignment"
162            && node
163                .child_by_field_name("left")
164                .is_some_and(|left| left.kind() == "identifier")
165            && let Some(module) = node.child_by_field_name("right").and_then(|right| {
166                importlib_call_module(right, source, &mut |local| {
167                    bindings.get(local).is_some_and(|binding| {
168                        binding.kind == ImportKind::Namespace
169                            && binding
170                                .namespace_imported_module
171                                .as_deref()
172                                .unwrap_or(&binding.module_specifier)
173                                == "importlib"
174                    })
175                })
176            })
177        {
178            modules.insert(module);
179        }
180        let mut cursor = node.walk();
181        stack.extend(node.named_children(&mut cursor));
182    }
183    modules
184}
185
186/// Recognize the compatibility-shim idiom that replaces the current module
187/// object with an imported workspace module:
188///
189/// `sys.modules[__name__] = imported_module`
190///
191/// Every component is resolved from tree-sitter fields and the import binder;
192/// similarly spelled attributes or locals do not create an alias edge.
193fn module_replacement_from_assignment(
194    assignment: Node<'_>,
195    source: &str,
196    bindings: &HashMap<String, ImportBinding>,
197) -> Option<PythonModuleReplacement> {
198    let (left, right) = (
199        assignment.child_by_field_name("left")?,
200        assignment.child_by_field_name("right")?,
201    );
202    if left.kind() != "subscript" || right.kind() != "identifier" {
203        return None;
204    }
205    let (value, subscript) = (
206        left.child_by_field_name("value")?,
207        left.child_by_field_name("subscript")?,
208    );
209    if value.kind() != "attribute"
210        || subscript.kind() != "identifier"
211        || node_source_text(subscript, source) != "__name__"
212    {
213        return None;
214    }
215    let (sys_local, modules) = (
216        value.child_by_field_name("object")?,
217        value.child_by_field_name("attribute")?,
218    );
219    if sys_local.kind() != "identifier"
220        || modules.kind() != "identifier"
221        || node_source_text(modules, source) != "modules"
222    {
223        return None;
224    }
225    let sys_binding = bindings.get(node_source_text(sys_local, source))?;
226    let sys_module = sys_binding
227        .namespace_imported_module
228        .as_deref()
229        .unwrap_or(&sys_binding.module_specifier);
230    if sys_binding.kind != ImportKind::Namespace || sys_module != "sys" {
231        return None;
232    }
233
234    let target_binding = bindings.get(node_source_text(right, source))?;
235    if target_binding.kind != ImportKind::Namespace {
236        return None;
237    }
238    Some(PythonModuleReplacement {
239        target_module: target_binding
240            .namespace_imported_module
241            .clone()
242            .unwrap_or_else(|| target_binding.module_specifier.clone()),
243    })
244}
245
246fn remove_direct_scope_bindings(
247    statement: Node<'_>,
248    source: &str,
249    bindings: &mut HashMap<String, ImportBinding>,
250) {
251    let mut stack = vec![statement];
252    while let Some(node) = stack.pop() {
253        for binding in python_direct_scope_bindings_bounded(node, source, || true)
254            .expect("unbounded binding collection cannot be cancelled")
255        {
256            bindings.remove(node_source_text(binding.declaration, source));
257        }
258        if matches!(
259            node.kind(),
260            "function_definition" | "class_definition" | "lambda"
261        ) {
262            continue;
263        }
264        let mut cursor = node.walk();
265        stack.extend(node.named_children(&mut cursor));
266    }
267}
268
269/// Parse the structured import facts for a Python document without executing
270/// it. LSP model binding uses the same AST-derived representation as the
271/// analyzer so external APIs are never selected by a terminal-name scan.
272pub fn parse_python_import_infos(source: &str) -> Vec<ImportInfo> {
273    let mut parser = tree_sitter::Parser::new();
274    parser
275        .set_language(&tree_sitter_python::LANGUAGE.into())
276        .expect("failed to load Python parser");
277    let Some(tree) = parser.parse(source, None) else {
278        return Vec::new();
279    };
280    let mut pending = vec![tree.root_node()];
281    let mut imports = Vec::new();
282    while let Some(node) = pending.pop() {
283        if matches!(node.kind(), "import_statement" | "import_from_statement") {
284            imports.extend(python_import_infos_from_node(node, source));
285            continue;
286        }
287        let mut cursor = node.walk();
288        pending.extend(node.named_children(&mut cursor));
289    }
290    imports
291}
292
293#[derive(Debug, Clone, PartialEq, Eq)]
294pub struct PythonImportBinding {
295    pub start_byte: usize,
296    pub scope_start_byte: usize,
297    pub scope_end_byte: usize,
298    function_scoped: bool,
299    pub local_name: String,
300    pub qualified_name: String,
301    pub consumed_attributes: usize,
302}
303
304impl PythonImportBinding {
305    pub fn is_function_scoped(&self) -> bool {
306        self.function_scoped
307    }
308}
309
310/// Return source-ordered, parser-derived local bindings for explicit Python
311/// imports. The byte offset lets callers select the last visible binding rather
312/// than treating a whole document as one unordered import scope.
313pub fn parse_python_import_bindings(source: &str) -> Vec<PythonImportBinding> {
314    let mut parser = tree_sitter::Parser::new();
315    parser
316        .set_language(&tree_sitter_python::LANGUAGE.into())
317        .expect("failed to load Python parser");
318    let Some(tree) = parser.parse(source, None) else {
319        return Vec::new();
320    };
321    python_import_bindings_from_tree(tree.root_node(), source)
322}
323
324/// Return structured import bindings from a tree the caller already owns.
325/// Query paths use this form so one request never reparses a document merely
326/// to recover function-local import scope.
327pub fn python_import_bindings_from_tree(root: Node<'_>, source: &str) -> Vec<PythonImportBinding> {
328    let mut pending = vec![root];
329    let mut nodes = Vec::new();
330    while let Some(node) = pending.pop() {
331        if matches!(node.kind(), "import_statement" | "import_from_statement") {
332            nodes.push(node);
333            continue;
334        }
335        let mut cursor = node.walk();
336        pending.extend(node.named_children(&mut cursor));
337    }
338    nodes.sort_by_key(Node::start_byte);
339    nodes
340        .into_iter()
341        .flat_map(|node| {
342            let (scope_start_byte, scope_end_byte, function_scoped) =
343                python_import_binding_scope(node, source.len());
344            python_import_infos_from_node(node, source)
345                .into_iter()
346                .filter_map(move |import| {
347                    let path = import.path.as_ref()?;
348                    let details = python_import_details(&import)?;
349                    match details {
350                        PythonImportDetails::Import { module, alias } => {
351                            let consumed_attributes = if alias.is_some() {
352                                0
353                            } else {
354                                path.segments.len().saturating_sub(1)
355                            };
356                            Some(PythonImportBinding {
357                                start_byte: node.start_byte(),
358                                scope_start_byte,
359                                scope_end_byte,
360                                function_scoped,
361                                local_name: alias.or_else(|| path.segments.first().cloned())?,
362                                qualified_name: module,
363                                consumed_attributes,
364                            })
365                        }
366                        PythonImportDetails::FromImport {
367                            module,
368                            name,
369                            alias,
370                            wildcard: false,
371                        } => Some(PythonImportBinding {
372                            start_byte: node.start_byte(),
373                            scope_start_byte,
374                            scope_end_byte,
375                            function_scoped,
376                            local_name: alias.unwrap_or(name.clone()),
377                            qualified_name: format!("{module}.{name}"),
378                            consumed_attributes: 0,
379                        }),
380                        PythonImportDetails::FromImport { wildcard: true, .. } => None,
381                    }
382                })
383        })
384        .collect()
385}
386
387fn python_import_binding_scope(node: Node<'_>, source_len: usize) -> (usize, usize, bool) {
388    let mut parent = node.parent();
389    while let Some(scope) = parent {
390        if matches!(scope.kind(), "function_definition" | "lambda") {
391            return (scope.start_byte(), scope.end_byte(), true);
392        }
393        parent = scope.parent();
394    }
395    (0, source_len, false)
396}
397
398#[cfg(test)]
399mod tests {
400    use super::*;
401    use brokk_bifrost_core::analyzer::usages::model::ImportBinder;
402
403    fn importlib_assignment_for(source: &str, importlib_local: &str) -> Option<String> {
404        let tree = parse_python_tree(source).expect("valid Python fixture");
405        let reference_start = source.rfind("canonical.setup").expect("reference site");
406        let mut stack = vec![tree.root_node()];
407        let mut reference = None;
408        while let Some(node) = stack.pop() {
409            if node.kind() == "identifier"
410                && node.start_byte() == reference_start
411                && node_source_text(node, source) == "canonical"
412            {
413                reference = Some(node);
414                break;
415            }
416            let mut cursor = node.walk();
417            stack.extend(node.named_children(&mut cursor));
418        }
419        imported_module_assignment_at(
420            reference.expect("canonical reference node"),
421            "canonical",
422            source,
423            |local| local == importlib_local,
424        )
425    }
426
427    #[test]
428    fn existing_tree_import_bindings_preserve_function_scope_and_relative_identity() {
429        let source = r#"def register():
430    from .workers.special import Worker
431    return Worker()
432"#;
433        let tree = parse_python_tree(source).expect("Python tree");
434        let bindings = python_import_bindings_from_tree(tree.root_node(), source);
435
436        assert_eq!(bindings.len(), 1, "{bindings:#?}");
437        assert!(bindings[0].is_function_scoped(), "{bindings:#?}");
438        assert_eq!(bindings[0].local_name, "Worker");
439        assert_eq!(bindings[0].qualified_name, ".workers.special.Worker");
440        assert!(bindings[0].scope_start_byte <= bindings[0].start_byte);
441        assert!(bindings[0].scope_end_byte >= source.rfind("Worker()").unwrap());
442    }
443
444    #[test]
445    fn importlib_assignment_requires_a_direct_literal_dominating_write() {
446        let valid = r#"import importlib as loader
447
448def test():
449    canonical = loader.import_module("routes.contacts.contacts_routes")
450    assert canonical.setup_contacts_routes
451"#;
452        assert_eq!(
453            importlib_assignment_for(valid, "loader"),
454            Some("routes.contacts.contacts_routes".to_string())
455        );
456
457        for near_miss in [
458            r#"import importlib as loader
459def test(module_name):
460    canonical = loader.import_module(module_name)
461    assert canonical.setup_contacts_routes
462"#,
463            r#"import importlib as loader
464def test():
465    if enabled:
466        canonical = loader.import_module("routes.contacts.contacts_routes")
467    assert canonical.setup_contacts_routes
468"#,
469            r#"import importlib as loader
470def test():
471    canonical = loader.import_module("routes.contacts.contacts_routes")
472    canonical = object()
473    assert canonical.setup_contacts_routes
474"#,
475        ] {
476            assert_eq!(
477                importlib_assignment_for(near_miss, "loader"),
478                None,
479                "near miss must not infer a module receiver: {near_miss:?}"
480            );
481        }
482        assert_eq!(
483            importlib_assignment_for(valid, "shadowed_loader"),
484            None,
485            "a similarly shaped non-importlib receiver must not resolve"
486        );
487    }
488
489    fn replacement_for(source: &str, binder: &ImportBinder) -> Option<PythonModuleReplacement> {
490        let tree = parse_python_tree(source).expect("valid Python fixture");
491        let root = tree.root_node();
492        let mut replacement = None;
493        let mut cursor = root.walk();
494        for statement in root.named_children(&mut cursor) {
495            let statement = if statement.kind() == "expression_statement" {
496                statement.named_child(0).expect("fixture expression")
497            } else {
498                statement
499            };
500            if statement.kind() == "assignment"
501                && let Some(next) =
502                    module_replacement_from_assignment(statement, source, &binder.bindings)
503                && replacement.replace(next).is_some()
504            {
505                return None;
506            }
507        }
508        replacement
509    }
510
511    #[test]
512    fn module_replacement_requires_exact_structured_import_bindings() {
513        let source = r#"import sys as _sys
514from routes.contacts import contacts_routes as _canonical
515
516_sys.modules[__name__] = _canonical
517"#;
518        let mut binder = ImportBinder::empty();
519        binder.bindings.insert(
520            "_sys".to_string(),
521            ImportBinding {
522                module_specifier: "sys".to_string(),
523                namespace_imported_module: Some("sys".to_string()),
524                kind: ImportKind::Namespace,
525                imported_name: None,
526            },
527        );
528        binder.bindings.insert(
529            "_canonical".to_string(),
530            ImportBinding {
531                module_specifier: "routes.contacts.contacts_routes".to_string(),
532                namespace_imported_module: None,
533                kind: ImportKind::Namespace,
534                imported_name: None,
535            },
536        );
537
538        assert_eq!(
539            replacement_for(source, &binder),
540            Some(PythonModuleReplacement {
541                target_module: "routes.contacts.contacts_routes".to_string(),
542            })
543        );
544
545        for near_miss in [
546            "cache.modules[__name__] = _canonical\n",
547            "_sys.modules[module_name] = _canonical\n",
548            "_sys.modules[__name__] = build()\n",
549            "def replace():\n    _sys.modules[__name__] = _canonical\n",
550        ] {
551            assert_eq!(
552                replacement_for(near_miss, &binder),
553                None,
554                "near miss must not replace module identity: {near_miss:?}"
555            );
556        }
557    }
558}
559
560pub fn module_replacement_of(
561    python: &dyn PythonSource,
562    file: &ProjectFile,
563    source: &str,
564) -> Option<PythonModuleReplacement> {
565    let tree = parse_python_tree(source)?;
566    let root = tree.root_node();
567    let mut bindings: HashMap<String, ImportBinding> = HashMap::default();
568    let mut replacement = None;
569    let mut cursor = root.walk();
570    for statement in root.named_children(&mut cursor) {
571        let statement = if statement.kind() == "expression_statement" {
572            let Some(expression) = statement.named_child(0) else {
573                continue;
574            };
575            expression
576        } else {
577            statement
578        };
579        match statement.kind() {
580            "import_statement" | "import_from_statement" => {
581                let imports = python_import_infos_from_node(statement, source);
582                bindings.extend(import_binder_from_imports(python, file, &imports).bindings);
583            }
584            "assignment" => {
585                if let Some(next) = module_replacement_from_assignment(statement, source, &bindings)
586                    && replacement.replace(next).is_some()
587                {
588                    return None;
589                }
590                remove_direct_scope_bindings(statement, source, &mut bindings);
591            }
592            _ => remove_direct_scope_bindings(statement, source, &mut bindings),
593        }
594    }
595    replacement
596}
597
598pub fn resolve_import_bindings(
599    python: &dyn PythonSource,
600    token: QueryToken<'_>,
601    file: &ProjectFile,
602) -> HashMap<String, CodeUnit> {
603    let imports = python.import_info_of(token, file);
604    let mut bindings = HashMap::default();
605    for resolved in resolve_imports_batched(python, file, &imports) {
606        for (binding, code_unit) in resolved {
607            bindings.insert(binding, code_unit);
608        }
609    }
610    bindings
611}
612
613/// Resolves every import in `imports` (`file`'s own imports), batching each import's primary
614/// module FQN lookup (see `primary_module_fqn`) into one store transaction instead of one per
615/// import. Shared by `resolve_import_bindings` and `resolve_import_target_files`, the two per-file
616/// "resolve everything" entry points -- both are called once per candidate file by the usages
617/// candidate walker, so unbatched resolution here means one store transaction per import times
618/// every file in the workspace.
619pub fn resolve_imports_batched(
620    python: &dyn PythonSource,
621    file: &ProjectFile,
622    imports: &[ImportInfo],
623) -> Vec<Vec<(String, CodeUnit)>> {
624    let primary_fqns: Vec<Option<String>> = imports
625        .iter()
626        .map(|import| primary_module_fqn(file, import))
627        .collect();
628    let to_resolve: Vec<String> = primary_fqns.iter().flatten().cloned().collect();
629    let mut batch_results = resolve_module_code_units_batch(python, &to_resolve).into_iter();
630
631    imports
632        .iter()
633        .zip(primary_fqns.iter())
634        .map(|(import, primary_fqn)| {
635            let hint = primary_fqn.as_ref().map(|_| batch_results.next().unwrap());
636            resolve_import_with_hint(python, file, import, hint.as_ref())
637        })
638        .collect()
639}
640
641/// The module FQN `resolve_import`'s fast path checks first, if any -- must stay in sync with the
642/// two `resolve_module_code_unit` call sites in `resolve_import_with_hint` below, since it's what
643/// lets `resolve_import_target_files` batch-resolve them ahead of the serial fallback logic.
644fn primary_module_fqn(file: &ProjectFile, import: &ImportInfo) -> Option<String> {
645    match python_import_details(import)? {
646        PythonImportDetails::Import { module, alias } => Some(python_namespace_binding_module(
647            import,
648            alias.as_deref(),
649            &module,
650        )),
651        PythonImportDetails::FromImport {
652            module,
653            name,
654            wildcard,
655            ..
656        } => {
657            if wildcard {
658                return None;
659            }
660            let resolved_module = if module.starts_with('.') {
661                resolve_python_relative_module(file, &module)
662            } else {
663                Some(module)
664            };
665            resolved_module.map(|resolved_module| format!("{resolved_module}.{name}"))
666        }
667    }
668}
669
670pub fn resolve_import(
671    python: &dyn PythonSource,
672    file: &ProjectFile,
673    import: &ImportInfo,
674) -> Vec<(String, CodeUnit)> {
675    resolve_import_with_hint(python, file, import, None)
676}
677
678/// `primary_hint`, when `Some`, is the already-resolved result of this import's primary module FQN
679/// (see `primary_module_fqn`) so the batched caller doesn't pay for a second lookup of the same FQN.
680fn resolve_import_with_hint(
681    python: &dyn PythonSource,
682    file: &ProjectFile,
683    import: &ImportInfo,
684    primary_hint: Option<&Option<CodeUnit>>,
685) -> Vec<(String, CodeUnit)> {
686    if let Some(details) = python_import_details(import) {
687        match details {
688            PythonImportDetails::Import { module, alias } => {
689                let binding = python_namespace_binding_name(import, alias.as_deref(), &module);
690                let bound_module =
691                    python_namespace_binding_module(import, alias.as_deref(), &module);
692                let resolved = match primary_hint {
693                    Some(hint) => hint.clone(),
694                    None => resolve_module_code_unit(python, &bound_module),
695                };
696                if let Some(module_code_unit) = resolved {
697                    return vec![(binding, module_code_unit)];
698                }
699            }
700            PythonImportDetails::FromImport {
701                module,
702                name,
703                alias,
704                wildcard,
705            } => {
706                let resolved_module = if module.starts_with('.') {
707                    resolve_python_relative_module(file, &module)
708                } else {
709                    Some(module)
710                };
711                let Some(resolved_module) = resolved_module else {
712                    return Vec::new();
713                };
714                if wildcard {
715                    return public_declarations_in_module(python, &resolved_module)
716                        .into_iter()
717                        .map(|code_unit| (code_unit.identifier().to_string(), code_unit))
718                        .collect();
719                }
720
721                let binding = alias.clone().unwrap_or_else(|| name.clone());
722                let module_candidate = format!("{resolved_module}.{name}");
723                let resolved = match primary_hint {
724                    Some(hint) => hint.clone(),
725                    None => resolve_module_code_unit(python, &module_candidate),
726                };
727                if let Some(code_unit) = resolved {
728                    return vec![(binding, code_unit)];
729                }
730                let exported = resolve_exported_name_from_module(python, &resolved_module, &name);
731                if !exported.is_empty() {
732                    return exported
733                        .into_iter()
734                        .map(|code_unit| (binding.clone(), code_unit))
735                        .collect();
736                }
737                let definitions: Vec<_> = python.definitions(&module_candidate).collect();
738                if !definitions.is_empty() {
739                    return definitions
740                        .into_iter()
741                        .map(|code_unit| (binding.clone(), code_unit))
742                        .collect();
743                }
744                let package_candidate: Vec<_> = python
745                    .definitions(&format!("{resolved_module}.{name}"))
746                    .collect();
747                if !package_candidate.is_empty() {
748                    return package_candidate
749                        .into_iter()
750                        .map(|code_unit| (binding.clone(), code_unit))
751                        .collect();
752                }
753            }
754        }
755    }
756    Vec::new()
757}
758
759pub fn resolve_exported_fqn(python: &dyn PythonSource, fqn: &str) -> Vec<CodeUnit> {
760    let Some((module, name)) = fqn.rsplit_once('.') else {
761        return Vec::new();
762    };
763    resolve_exported_name_from_module(python, module, name)
764}
765
766/// Resolve an unambiguous chain of explicit named reexports without
767/// constructing export indexes for each intermediate module. Star exports,
768/// shadowing, and every other ambiguous shape return `None` so callers can
769/// use the complete, source-order-aware export resolver below.
770fn resolve_direct_named_exported_fqn(
771    python: &dyn PythonSource,
772    fqn: &str,
773) -> Option<Vec<CodeUnit>> {
774    let (module, name) = fqn.rsplit_once('.')?;
775    let mut results = Vec::new();
776    let mut queue = VecDeque::from([(module.to_string(), name.to_string())]);
777    let mut visited = HashSet::default();
778
779    while let Some((module, export_name)) = queue.pop_front() {
780        if !visited.insert((module.clone(), export_name.clone())) {
781            continue;
782        }
783        let module_unit = resolve_module_code_unit(python, &module)?;
784        let file = module_unit.source();
785        let local = local_export_declarations(python, file, &export_name);
786        let binder = python.import_binder_of(file);
787        let binding = binder.bindings.get(&export_name);
788        if !local.is_empty() && binding.is_some() {
789            return None;
790        }
791        if !local.is_empty() {
792            results.extend(local);
793            continue;
794        }
795        let binding = binding?;
796        if binding.kind != ImportKind::Named {
797            return None;
798        }
799        let imported_name = binding.imported_name.as_ref()?;
800        queue.push_back((binding.module_specifier.clone(), imported_name.clone()));
801    }
802
803    results.sort_by(|left, right| {
804        left.source()
805            .cmp(right.source())
806            .then_with(|| left.fq_name().cmp(&right.fq_name()))
807    });
808    results.dedup();
809    (!results.is_empty()).then_some(results)
810}
811
812/// Resolve a Python FQN with the cheapest semantically complete tier that
813/// can answer it. The direct reexport walk handles only proven,
814/// collision-free chains; ambiguous shapes use the ordered export index,
815/// and the exact lookup remains the final fallback for non-export symbols.
816pub fn resolve_fqn_candidates(
817    python: &dyn PythonSource,
818    fqn: &str,
819    exact: impl FnOnce(&str) -> Vec<CodeUnit>,
820) -> Vec<CodeUnit> {
821    if let Some(candidates) = resolve_direct_named_exported_fqn(python, fqn) {
822        return candidates;
823    }
824    let candidates = resolve_exported_fqn(python, fqn);
825    if !candidates.is_empty() {
826        return candidates;
827    }
828    exact(fqn)
829}
830
831fn resolve_exported_name_from_module(
832    python: &dyn PythonSource,
833    module: &str,
834    name: &str,
835) -> Vec<CodeUnit> {
836    let Some(module_unit) = resolve_module_code_unit(python, module) else {
837        return Vec::new();
838    };
839    resolve_exported_name(python, module_unit.source(), name)
840}
841
842fn resolve_exported_name(
843    python: &dyn PythonSource,
844    module_file: &ProjectFile,
845    name: &str,
846) -> Vec<CodeUnit> {
847    let mut results = Vec::new();
848    let mut queue = VecDeque::from([(module_file.clone(), name.to_string())]);
849    let mut visited = HashSet::default();
850
851    while let Some((file, export_name)) = queue.pop_front() {
852        if !visited.insert((file.clone(), export_name.clone())) {
853            continue;
854        }
855
856        let index = python.export_index_of(&file);
857        if let Some(entry) = index.exports_by_name.get(&export_name) {
858            match entry {
859                ExportEntry::Local { local_name } => {
860                    results.extend(local_export_declarations(python, &file, local_name));
861                }
862                ExportEntry::ReexportedNamed {
863                    module_specifier,
864                    imported_name,
865                } => {
866                    for target_file in
867                        resolve_module_files_for_export(python, &file, module_specifier)
868                    {
869                        queue.push_back((target_file, imported_name.clone()));
870                    }
871                }
872                ExportEntry::ReexportedModule { module_specifier } => {
873                    // Terminal: the export *is* the module, so the walk stops
874                    // here instead of looking the name up inside it.
875                    results.extend(resolve_module_code_unit(python, module_specifier));
876                }
877                ExportEntry::Default { local_name } => {
878                    if let Some(local_name) = local_name {
879                        results.extend(local_export_declarations(python, &file, local_name));
880                    }
881                }
882            }
883            continue;
884        }
885
886        if !export_name.starts_with('_') {
887            for star in &index.reexport_stars {
888                for target_file in
889                    resolve_module_files_for_export(python, &file, &star.module_specifier)
890                {
891                    queue.push_back((target_file, export_name.clone()));
892                }
893            }
894        }
895    }
896
897    results.sort_by(|left, right| {
898        left.source()
899            .cmp(right.source())
900            .then_with(|| left.fq_name().cmp(&right.fq_name()))
901    });
902    results.dedup();
903    results
904}
905
906fn local_export_declarations(
907    index: &dyn CodeUnitIndex,
908    file: &ProjectFile,
909    local_name: &str,
910) -> Vec<CodeUnit> {
911    index
912        .top_level_declarations(file)
913        .into_iter()
914        .filter(|unit| {
915            unit.identifier() == local_name
916                && index
917                    .parent_of(unit)
918                    .is_some_and(|parent| parent.is_module() && parent.source() == file)
919        })
920        .collect()
921}
922
923fn resolve_module_files_for_export(
924    python: &dyn PythonSource,
925    importing_file: &ProjectFile,
926    module_specifier: &str,
927) -> Vec<ProjectFile> {
928    let resolved_module = if module_specifier.starts_with('.') {
929        resolve_python_relative_module(importing_file, module_specifier)
930    } else {
931        Some(module_specifier.to_string())
932    };
933    let Some(resolved_module) = resolved_module else {
934        return Vec::new();
935    };
936    // Tree-sitter tells us the import syntax, but module-to-file resolution
937    // is analyzer state. Use the prebuilt module code-unit map here instead
938    // of the usage index so interactive definition lookup stays lightweight.
939    resolve_module_code_unit(python, &resolved_module)
940        .map(|unit| vec![unit.source().clone()])
941        .unwrap_or_default()
942}
943
944pub fn extract_package_from_python_wildcard(import: &ImportInfo) -> Option<String> {
945    let details = python_import_details(import)?;
946    match details {
947        PythonImportDetails::FromImport {
948            module, wildcard, ..
949        } if wildcard => Some(module),
950        _ => None,
951    }
952}
953
954#[derive(Debug, Clone)]
955pub enum PythonImportDetails {
956    Import {
957        module: String,
958        alias: Option<String>,
959    },
960    FromImport {
961        module: String,
962        name: String,
963        alias: Option<String>,
964        wildcard: bool,
965    },
966}
967
968pub fn python_import_infos_from_node(node: Node<'_>, source: &str) -> Vec<ImportInfo> {
969    match node.kind() {
970        "import_statement" => python_namespace_import_infos(node, source),
971        "import_from_statement" => python_from_import_infos(node, source),
972        _ => Vec::new(),
973    }
974}
975
976pub fn python_import_details(import: &ImportInfo) -> Option<PythonImportDetails> {
977    let path = import.path.as_ref()?;
978    match path.kind? {
979        StructuredImportPathKind::Namespace => Some(PythonImportDetails::Import {
980            module: join_python_import_segments(&path.segments),
981            alias: import.alias.clone(),
982        }),
983        // Python has no static imports; the variant belongs to Java.
984        StructuredImportPathKind::StaticMember => None,
985        StructuredImportPathKind::ImportFrom => {
986            let (name, module_segments) = if import.is_wildcard {
987                ("*".to_string(), path.segments.as_slice())
988            } else {
989                let (name, module_segments) = path.segments.split_last()?;
990                (name.clone(), module_segments)
991            };
992            Some(PythonImportDetails::FromImport {
993                module: join_python_import_segments(module_segments),
994                name,
995                alias: import.alias.clone(),
996                wildcard: import.is_wildcard,
997            })
998        }
999    }
1000}
1001
1002fn python_namespace_import_infos(node: Node<'_>, source: &str) -> Vec<ImportInfo> {
1003    let mut infos = Vec::new();
1004    let mut cursor = node.walk();
1005    for imported in node.children_by_field_name("name", &mut cursor) {
1006        let (module_node, alias_node) = if imported.kind() == "aliased_import" {
1007            let Some(name) = imported.child_by_field_name("name") else {
1008                continue;
1009            };
1010            (name, imported.child_by_field_name("alias"))
1011        } else {
1012            (imported, None)
1013        };
1014        let alias = alias_node
1015            .map(|alias| py_node_text(alias, source).trim().to_string())
1016            .filter(|alias| !alias.is_empty());
1017        let segments = python_path_segments(module_node, source);
1018        if segments.is_empty() {
1019            continue;
1020        }
1021        // `import a.b` binds `a`: the first segment's own token. A renamed
1022        // import binds its alias token instead.
1023        let binder_span = alias
1024            .is_some()
1025            .then_some(alias_node)
1026            .flatten()
1027            .or_else(|| python_first_segment_node(module_node))
1028            .map(brokk_bifrost_core::analyzer::common::node_span);
1029        let module = join_python_import_segments(&segments);
1030        let identifier = alias.clone().or_else(|| segments.first().cloned());
1031        infos.push(ImportInfo {
1032            raw_snippet: if let Some(alias) = &alias {
1033                format!("import {module} as {alias}")
1034            } else {
1035                format!("import {module}")
1036            },
1037            is_wildcard: false,
1038            is_global: false,
1039            identifier,
1040            alias,
1041            path: Some(StructuredImportPath {
1042                segments,
1043                kind: Some(StructuredImportPathKind::Namespace),
1044                lexical_prefixes: Vec::new(),
1045                lexical_scopes: Vec::new(),
1046                declaration_start_byte: node.start_byte(),
1047            }),
1048            binder_span,
1049        });
1050    }
1051    infos
1052}
1053
1054fn python_from_import_infos(node: Node<'_>, source: &str) -> Vec<ImportInfo> {
1055    let Some(module_node) = node.child_by_field_name("module_name") else {
1056        return Vec::new();
1057    };
1058    let module_segments = python_module_segments(module_node, source);
1059    if module_segments.is_empty() {
1060        return Vec::new();
1061    }
1062
1063    let mut infos = Vec::new();
1064    let has_wildcard_import = {
1065        let mut cursor = node.walk();
1066        node.named_children(&mut cursor)
1067            .any(|child| child.kind() == "wildcard_import")
1068    };
1069    let mut cursor = node.walk();
1070    let imported_names: Vec<_> = node.children_by_field_name("name", &mut cursor).collect();
1071    if has_wildcard_import {
1072        let module = join_python_import_segments(&module_segments);
1073        infos.push(ImportInfo {
1074            raw_snippet: format!("from {module} import *"),
1075            is_wildcard: true,
1076            is_global: false,
1077            identifier: None,
1078            alias: None,
1079            path: Some(StructuredImportPath {
1080                segments: module_segments,
1081                kind: Some(StructuredImportPathKind::ImportFrom),
1082                lexical_prefixes: Vec::new(),
1083                lexical_scopes: Vec::new(),
1084                declaration_start_byte: node.start_byte(),
1085            }),
1086            binder_span: None,
1087        });
1088        return infos;
1089    }
1090    if imported_names.is_empty() {
1091        return infos;
1092    }
1093
1094    for imported in imported_names {
1095        let (name_node, alias_node) = if imported.kind() == "aliased_import" {
1096            let Some(name) = imported.child_by_field_name("name") else {
1097                continue;
1098            };
1099            (name, imported.child_by_field_name("alias"))
1100        } else {
1101            (imported, None)
1102        };
1103        let alias = alias_node
1104            .map(|alias| py_node_text(alias, source).trim().to_string())
1105            .filter(|alias| !alias.is_empty());
1106        let name_segments = python_path_segments(name_node, source);
1107        if name_segments.is_empty() {
1108            continue;
1109        }
1110        // `from m import x` binds `x`'s own token; a rename binds the alias
1111        // token. A multi-segment imported name binds no single token.
1112        let binder_span = alias
1113            .is_some()
1114            .then_some(alias_node)
1115            .flatten()
1116            .or_else(|| {
1117                (name_segments.len() == 1)
1118                    .then(|| python_first_segment_node(name_node))
1119                    .flatten()
1120            })
1121            .map(brokk_bifrost_core::analyzer::common::node_span);
1122        let imported_name = join_python_import_segments(&name_segments);
1123        let mut segments = module_segments.clone();
1124        segments.extend(name_segments);
1125        let module = join_python_import_segments(&module_segments);
1126        infos.push(ImportInfo {
1127            raw_snippet: if let Some(alias) = &alias {
1128                format!("from {module} import {imported_name} as {alias}")
1129            } else {
1130                format!("from {module} import {imported_name}")
1131            },
1132            is_wildcard: false,
1133            is_global: false,
1134            identifier: Some(alias.clone().unwrap_or_else(|| imported_name.clone())),
1135            alias,
1136            path: Some(StructuredImportPath {
1137                segments,
1138                kind: Some(StructuredImportPathKind::ImportFrom),
1139                lexical_prefixes: Vec::new(),
1140                lexical_scopes: Vec::new(),
1141                declaration_start_byte: node.start_byte(),
1142            }),
1143            binder_span,
1144        });
1145    }
1146    infos
1147}
1148
1149fn python_module_segments(module: Node<'_>, source: &str) -> Vec<String> {
1150    if module.kind() == "relative_import" {
1151        let mut cursor = module.walk();
1152        let mut prefix = String::new();
1153        let mut path_node = None;
1154        for child in module.named_children(&mut cursor) {
1155            match child.kind() {
1156                "import_prefix" if prefix.is_empty() => {
1157                    prefix = py_node_text(child, source).trim().to_string();
1158                }
1159                "dotted_name" if path_node.is_none() => {
1160                    path_node = Some(child);
1161                }
1162                _ => {}
1163            }
1164        }
1165        let mut segments = path_node
1166            .map(|path| python_path_segments(path, source))
1167            .unwrap_or_default();
1168        if !prefix.is_empty() {
1169            if let Some(first) = segments.first_mut() {
1170                first.insert_str(0, &prefix);
1171            } else {
1172                segments.push(prefix);
1173            }
1174        }
1175        return segments;
1176    }
1177    python_path_segments(module, source)
1178}
1179
1180/// The token that spells a path's first segment: the identifier itself, or a
1181/// dotted name's first identifier. `None` when the shape has no leading
1182/// identifier token of its own (e.g. a relative-import prefix).
1183fn python_first_segment_node(node: Node<'_>) -> Option<Node<'_>> {
1184    match node.kind() {
1185        "identifier" => Some(node),
1186        "dotted_name" => {
1187            let mut cursor = node.walk();
1188            node.named_children(&mut cursor)
1189                .find(|child| child.kind() == "identifier")
1190        }
1191        _ => None,
1192    }
1193}
1194
1195fn python_path_segments(node: Node<'_>, source: &str) -> Vec<String> {
1196    match node.kind() {
1197        "identifier" => vec![py_node_text(node, source).trim().to_string()],
1198        "dotted_name" => {
1199            let mut segments = Vec::new();
1200            let mut cursor = node.walk();
1201            for child in node.named_children(&mut cursor) {
1202                segments.extend(python_path_segments(child, source));
1203            }
1204            segments
1205        }
1206        _ => {
1207            let mut segments = Vec::new();
1208            let mut cursor = node.walk();
1209            for child in node.named_children(&mut cursor) {
1210                segments.extend(python_path_segments(child, source));
1211            }
1212            segments
1213        }
1214    }
1215}
1216
1217fn join_python_import_segments(segments: &[String]) -> String {
1218    let Some((first, rest)) = segments.split_first() else {
1219        return String::new();
1220    };
1221    if first.starts_with('.') && !rest.is_empty() {
1222        format!("{first}.{}", rest.join("."))
1223    } else {
1224        segments.join(".")
1225    }
1226}
1227
1228pub fn python_namespace_binding_name(
1229    import: &ImportInfo,
1230    alias: Option<&str>,
1231    module: &str,
1232) -> String {
1233    import
1234        .identifier
1235        .clone()
1236        .or_else(|| alias.map(str::to_string))
1237        .unwrap_or_else(|| module.to_string())
1238}
1239
1240pub fn python_namespace_binding_module(
1241    import: &ImportInfo,
1242    alias: Option<&str>,
1243    module: &str,
1244) -> String {
1245    if alias.is_some() {
1246        return module.to_string();
1247    }
1248    import
1249        .path
1250        .as_ref()
1251        .and_then(|path| path.segments.first().cloned())
1252        .unwrap_or_else(|| module.to_string())
1253}
1254
1255pub fn resolve_python_relative_module(
1256    source_file: &ProjectFile,
1257    module_expr: &str,
1258) -> Option<String> {
1259    resolve_python_relative_module_from_package(&python_current_package(source_file), module_expr)
1260}
1261
1262/// Resolve a structured Python module expression against an already-known
1263/// package identity. Dependency-pack producers know module identities without
1264/// owning a workspace [`ProjectFile`], so they use this entry point.
1265pub fn resolve_python_relative_module_from_package(
1266    current_package: &str,
1267    module_expr: &str,
1268) -> Option<String> {
1269    let level = module_expr.chars().take_while(|ch| *ch == '.').count();
1270    let suffix = module_expr[level..].trim_matches('.');
1271    let mut parts: Vec<_> = current_package
1272        .split('.')
1273        .filter(|part| !part.is_empty())
1274        .map(str::to_string)
1275        .collect();
1276    if level == 0 {
1277        return Some(module_expr.to_string());
1278    }
1279    if level > 0 {
1280        if level - 1 > parts.len() {
1281            return None;
1282        }
1283        parts.truncate(parts.len() - (level - 1));
1284    }
1285    if !suffix.is_empty() {
1286        parts.extend(suffix.split('.').map(str::to_string));
1287    }
1288    Some(parts.join("."))
1289}
1290
1291fn python_current_package(source_file: &ProjectFile) -> String {
1292    let module = python_module_name(source_file);
1293    if source_file
1294        .rel_path()
1295        .file_name()
1296        .and_then(|name| name.to_str())
1297        == Some("__init__.py")
1298    {
1299        module
1300    } else {
1301        module
1302            .rsplit_once('.')
1303            .map(|(package, _)| package.to_string())
1304            .unwrap_or_default()
1305    }
1306}