use crate::bindings::python_direct_scope_bindings_bounded;
use crate::declarations::{parse_python_tree, py_node_text, python_module_name};
use crate::graph_support::{
PythonSource, import_binder_from_imports, public_declarations_in_module,
resolve_module_code_unit, resolve_module_code_units_batch,
};
use brokk_bifrost_core::analyzer::common::node_source_text;
use brokk_bifrost_core::analyzer::model::{
ImportInfo, StructuredImportPath, StructuredImportPathKind,
};
use brokk_bifrost_core::analyzer::usages::model::{ExportEntry, ImportBinding, ImportKind};
use brokk_bifrost_core::analyzer::{CodeUnit, CodeUnitIndex, ProjectFile};
use brokk_bifrost_core::hash::{HashMap, HashSet};
use std::collections::VecDeque;
use tree_sitter::Node;
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct PythonModuleReplacement {
pub target_module: String,
}
fn module_replacement_from_assignment(
assignment: Node<'_>,
source: &str,
bindings: &HashMap<String, ImportBinding>,
) -> Option<PythonModuleReplacement> {
let (left, right) = (
assignment.child_by_field_name("left")?,
assignment.child_by_field_name("right")?,
);
if left.kind() != "subscript" || right.kind() != "identifier" {
return None;
}
let (value, subscript) = (
left.child_by_field_name("value")?,
left.child_by_field_name("subscript")?,
);
if value.kind() != "attribute"
|| subscript.kind() != "identifier"
|| node_source_text(subscript, source) != "__name__"
{
return None;
}
let (sys_local, modules) = (
value.child_by_field_name("object")?,
value.child_by_field_name("attribute")?,
);
if sys_local.kind() != "identifier"
|| modules.kind() != "identifier"
|| node_source_text(modules, source) != "modules"
{
return None;
}
let sys_binding = bindings.get(node_source_text(sys_local, source))?;
let sys_module = sys_binding
.namespace_imported_module
.as_deref()
.unwrap_or(&sys_binding.module_specifier);
if sys_binding.kind != ImportKind::Namespace || sys_module != "sys" {
return None;
}
let target_binding = bindings.get(node_source_text(right, source))?;
if target_binding.kind != ImportKind::Namespace {
return None;
}
Some(PythonModuleReplacement {
target_module: target_binding
.namespace_imported_module
.clone()
.unwrap_or_else(|| target_binding.module_specifier.clone()),
})
}
fn remove_direct_scope_bindings(
statement: Node<'_>,
source: &str,
bindings: &mut HashMap<String, ImportBinding>,
) {
let mut stack = vec![statement];
while let Some(node) = stack.pop() {
for binding in python_direct_scope_bindings_bounded(node, source, || true)
.expect("unbounded binding collection cannot be cancelled")
{
bindings.remove(node_source_text(binding.declaration, source));
}
if matches!(
node.kind(),
"function_definition" | "class_definition" | "lambda"
) {
continue;
}
let mut cursor = node.walk();
stack.extend(node.named_children(&mut cursor));
}
}
pub fn parse_python_import_infos(source: &str) -> Vec<ImportInfo> {
let mut parser = tree_sitter::Parser::new();
parser
.set_language(&tree_sitter_python::LANGUAGE.into())
.expect("failed to load Python parser");
let Some(tree) = parser.parse(source, None) else {
return Vec::new();
};
let mut pending = vec![tree.root_node()];
let mut imports = Vec::new();
while let Some(node) = pending.pop() {
if matches!(node.kind(), "import_statement" | "import_from_statement") {
imports.extend(python_import_infos_from_node(node, source));
continue;
}
let mut cursor = node.walk();
pending.extend(node.named_children(&mut cursor));
}
imports
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct PythonImportBinding {
pub start_byte: usize,
pub scope_start_byte: usize,
pub scope_end_byte: usize,
function_scoped: bool,
pub local_name: String,
pub qualified_name: String,
pub consumed_attributes: usize,
}
impl PythonImportBinding {
pub fn is_function_scoped(&self) -> bool {
self.function_scoped
}
}
pub fn parse_python_import_bindings(source: &str) -> Vec<PythonImportBinding> {
let mut parser = tree_sitter::Parser::new();
parser
.set_language(&tree_sitter_python::LANGUAGE.into())
.expect("failed to load Python parser");
let Some(tree) = parser.parse(source, None) else {
return Vec::new();
};
let mut pending = vec![tree.root_node()];
let mut nodes = Vec::new();
while let Some(node) = pending.pop() {
if matches!(node.kind(), "import_statement" | "import_from_statement") {
nodes.push(node);
continue;
}
let mut cursor = node.walk();
pending.extend(node.named_children(&mut cursor));
}
nodes.sort_by_key(Node::start_byte);
nodes
.into_iter()
.flat_map(|node| {
let (scope_start_byte, scope_end_byte, function_scoped) =
python_import_binding_scope(node, source.len());
python_import_infos_from_node(node, source)
.into_iter()
.filter_map(move |import| {
let path = import.path.as_ref()?;
let details = python_import_details(&import)?;
match details {
PythonImportDetails::Import { module, alias } => {
let consumed_attributes = if alias.is_some() {
0
} else {
path.segments.len().saturating_sub(1)
};
Some(PythonImportBinding {
start_byte: node.start_byte(),
scope_start_byte,
scope_end_byte,
function_scoped,
local_name: alias.or_else(|| path.segments.first().cloned())?,
qualified_name: module,
consumed_attributes,
})
}
PythonImportDetails::FromImport {
module,
name,
alias,
wildcard: false,
} => Some(PythonImportBinding {
start_byte: node.start_byte(),
scope_start_byte,
scope_end_byte,
function_scoped,
local_name: alias.unwrap_or(name.clone()),
qualified_name: format!("{module}.{name}"),
consumed_attributes: 0,
}),
PythonImportDetails::FromImport { wildcard: true, .. } => None,
}
})
})
.collect()
}
fn python_import_binding_scope(node: Node<'_>, source_len: usize) -> (usize, usize, bool) {
let mut parent = node.parent();
while let Some(scope) = parent {
if matches!(scope.kind(), "function_definition" | "lambda") {
return (scope.start_byte(), scope.end_byte(), true);
}
parent = scope.parent();
}
(0, source_len, false)
}
#[cfg(test)]
mod tests {
use super::*;
use brokk_bifrost_core::analyzer::usages::model::ImportBinder;
fn replacement_for(source: &str, binder: &ImportBinder) -> Option<PythonModuleReplacement> {
let tree = parse_python_tree(source).expect("valid Python fixture");
let root = tree.root_node();
let mut replacement = None;
let mut cursor = root.walk();
for statement in root.named_children(&mut cursor) {
let statement = if statement.kind() == "expression_statement" {
statement.named_child(0).expect("fixture expression")
} else {
statement
};
if statement.kind() == "assignment"
&& let Some(next) =
module_replacement_from_assignment(statement, source, &binder.bindings)
&& replacement.replace(next).is_some()
{
return None;
}
}
replacement
}
#[test]
fn module_replacement_requires_exact_structured_import_bindings() {
let source = r#"import sys as _sys
from routes.contacts import contacts_routes as _canonical
_sys.modules[__name__] = _canonical
"#;
let mut binder = ImportBinder::empty();
binder.bindings.insert(
"_sys".to_string(),
ImportBinding {
module_specifier: "sys".to_string(),
namespace_imported_module: Some("sys".to_string()),
kind: ImportKind::Namespace,
imported_name: None,
},
);
binder.bindings.insert(
"_canonical".to_string(),
ImportBinding {
module_specifier: "routes.contacts.contacts_routes".to_string(),
namespace_imported_module: None,
kind: ImportKind::Namespace,
imported_name: None,
},
);
assert_eq!(
replacement_for(source, &binder),
Some(PythonModuleReplacement {
target_module: "routes.contacts.contacts_routes".to_string(),
})
);
for near_miss in [
"cache.modules[__name__] = _canonical\n",
"_sys.modules[module_name] = _canonical\n",
"_sys.modules[__name__] = build()\n",
"def replace():\n _sys.modules[__name__] = _canonical\n",
] {
assert_eq!(
replacement_for(near_miss, &binder),
None,
"near miss must not replace module identity: {near_miss:?}"
);
}
}
}
pub fn module_replacement_of(
python: &dyn PythonSource,
file: &ProjectFile,
source: &str,
) -> Option<PythonModuleReplacement> {
let tree = parse_python_tree(source)?;
let root = tree.root_node();
let mut bindings: HashMap<String, ImportBinding> = HashMap::default();
let mut replacement = None;
let mut cursor = root.walk();
for statement in root.named_children(&mut cursor) {
let statement = if statement.kind() == "expression_statement" {
let Some(expression) = statement.named_child(0) else {
continue;
};
expression
} else {
statement
};
match statement.kind() {
"import_statement" | "import_from_statement" => {
let imports = python_import_infos_from_node(statement, source);
bindings.extend(import_binder_from_imports(python, file, &imports).bindings);
}
"assignment" => {
if let Some(next) = module_replacement_from_assignment(statement, source, &bindings)
&& replacement.replace(next).is_some()
{
return None;
}
remove_direct_scope_bindings(statement, source, &mut bindings);
}
_ => remove_direct_scope_bindings(statement, source, &mut bindings),
}
}
replacement
}
pub fn resolve_import_bindings(
python: &dyn PythonSource,
file: &ProjectFile,
) -> HashMap<String, CodeUnit> {
let imports = python.import_info_of(file);
let mut bindings = HashMap::default();
for resolved in resolve_imports_batched(python, file, &imports) {
for (binding, code_unit) in resolved {
bindings.insert(binding, code_unit);
}
}
bindings
}
pub fn resolve_imports_batched(
python: &dyn PythonSource,
file: &ProjectFile,
imports: &[ImportInfo],
) -> Vec<Vec<(String, CodeUnit)>> {
let primary_fqns: Vec<Option<String>> = imports
.iter()
.map(|import| primary_module_fqn(file, import))
.collect();
let to_resolve: Vec<String> = primary_fqns.iter().flatten().cloned().collect();
let mut batch_results = resolve_module_code_units_batch(python, &to_resolve).into_iter();
imports
.iter()
.zip(primary_fqns.iter())
.map(|(import, primary_fqn)| {
let hint = primary_fqn.as_ref().map(|_| batch_results.next().unwrap());
resolve_import_with_hint(python, file, import, hint.as_ref())
})
.collect()
}
fn primary_module_fqn(file: &ProjectFile, import: &ImportInfo) -> Option<String> {
match python_import_details(import)? {
PythonImportDetails::Import { module, alias } => Some(python_namespace_binding_module(
import,
alias.as_deref(),
&module,
)),
PythonImportDetails::FromImport {
module,
name,
wildcard,
..
} => {
if wildcard {
return None;
}
let resolved_module = if module.starts_with('.') {
resolve_python_relative_module(file, &module)
} else {
Some(module)
};
resolved_module.map(|resolved_module| format!("{resolved_module}.{name}"))
}
}
}
pub fn resolve_import(
python: &dyn PythonSource,
file: &ProjectFile,
import: &ImportInfo,
) -> Vec<(String, CodeUnit)> {
resolve_import_with_hint(python, file, import, None)
}
fn resolve_import_with_hint(
python: &dyn PythonSource,
file: &ProjectFile,
import: &ImportInfo,
primary_hint: Option<&Option<CodeUnit>>,
) -> Vec<(String, CodeUnit)> {
if let Some(details) = python_import_details(import) {
match details {
PythonImportDetails::Import { module, alias } => {
let binding = python_namespace_binding_name(import, alias.as_deref(), &module);
let bound_module =
python_namespace_binding_module(import, alias.as_deref(), &module);
let resolved = match primary_hint {
Some(hint) => hint.clone(),
None => resolve_module_code_unit(python, &bound_module),
};
if let Some(module_code_unit) = resolved {
return vec![(binding, module_code_unit)];
}
}
PythonImportDetails::FromImport {
module,
name,
alias,
wildcard,
} => {
let resolved_module = if module.starts_with('.') {
resolve_python_relative_module(file, &module)
} else {
Some(module)
};
let Some(resolved_module) = resolved_module else {
return Vec::new();
};
if wildcard {
return public_declarations_in_module(python, &resolved_module)
.into_iter()
.map(|code_unit| (code_unit.identifier().to_string(), code_unit))
.collect();
}
let binding = alias.clone().unwrap_or_else(|| name.clone());
let module_candidate = format!("{resolved_module}.{name}");
let resolved = match primary_hint {
Some(hint) => hint.clone(),
None => resolve_module_code_unit(python, &module_candidate),
};
if let Some(code_unit) = resolved {
return vec![(binding, code_unit)];
}
let exported = resolve_exported_name_from_module(python, &resolved_module, &name);
if !exported.is_empty() {
return exported
.into_iter()
.map(|code_unit| (binding.clone(), code_unit))
.collect();
}
let definitions: Vec<_> = python.definitions(&module_candidate).collect();
if !definitions.is_empty() {
return definitions
.into_iter()
.map(|code_unit| (binding.clone(), code_unit))
.collect();
}
let package_candidate: Vec<_> = python
.definitions(&format!("{resolved_module}.{name}"))
.collect();
if !package_candidate.is_empty() {
return package_candidate
.into_iter()
.map(|code_unit| (binding.clone(), code_unit))
.collect();
}
}
}
}
Vec::new()
}
pub fn resolve_exported_fqn(python: &dyn PythonSource, fqn: &str) -> Vec<CodeUnit> {
let Some((module, name)) = fqn.rsplit_once('.') else {
return Vec::new();
};
resolve_exported_name_from_module(python, module, name)
}
fn resolve_direct_named_exported_fqn(
python: &dyn PythonSource,
fqn: &str,
) -> Option<Vec<CodeUnit>> {
let (module, name) = fqn.rsplit_once('.')?;
let mut results = Vec::new();
let mut queue = VecDeque::from([(module.to_string(), name.to_string())]);
let mut visited = HashSet::default();
while let Some((module, export_name)) = queue.pop_front() {
if !visited.insert((module.clone(), export_name.clone())) {
continue;
}
let module_unit = resolve_module_code_unit(python, &module)?;
let file = module_unit.source();
let local = local_export_declarations(python, file, &export_name);
let binder = python.import_binder_of(file);
let binding = binder.bindings.get(&export_name);
if !local.is_empty() && binding.is_some() {
return None;
}
if !local.is_empty() {
results.extend(local);
continue;
}
let binding = binding?;
if binding.kind != ImportKind::Named {
return None;
}
let imported_name = binding.imported_name.as_ref()?;
queue.push_back((binding.module_specifier.clone(), imported_name.clone()));
}
results.sort_by(|left, right| {
left.source()
.cmp(right.source())
.then_with(|| left.fq_name().cmp(&right.fq_name()))
});
results.dedup();
(!results.is_empty()).then_some(results)
}
pub fn resolve_fqn_candidates(
python: &dyn PythonSource,
fqn: &str,
exact: impl FnOnce(&str) -> Vec<CodeUnit>,
) -> Vec<CodeUnit> {
if let Some(candidates) = resolve_direct_named_exported_fqn(python, fqn) {
return candidates;
}
let candidates = resolve_exported_fqn(python, fqn);
if !candidates.is_empty() {
return candidates;
}
exact(fqn)
}
fn resolve_exported_name_from_module(
python: &dyn PythonSource,
module: &str,
name: &str,
) -> Vec<CodeUnit> {
let Some(module_unit) = resolve_module_code_unit(python, module) else {
return Vec::new();
};
resolve_exported_name(python, module_unit.source(), name)
}
fn resolve_exported_name(
python: &dyn PythonSource,
module_file: &ProjectFile,
name: &str,
) -> Vec<CodeUnit> {
let mut results = Vec::new();
let mut queue = VecDeque::from([(module_file.clone(), name.to_string())]);
let mut visited = HashSet::default();
while let Some((file, export_name)) = queue.pop_front() {
if !visited.insert((file.clone(), export_name.clone())) {
continue;
}
let index = python.export_index_of(&file);
if let Some(entry) = index.exports_by_name.get(&export_name) {
match entry {
ExportEntry::Local { local_name } => {
results.extend(local_export_declarations(python, &file, local_name));
}
ExportEntry::ReexportedNamed {
module_specifier,
imported_name,
} => {
for target_file in
resolve_module_files_for_export(python, &file, module_specifier)
{
queue.push_back((target_file, imported_name.clone()));
}
}
ExportEntry::ReexportedModule { module_specifier } => {
results.extend(resolve_module_code_unit(python, module_specifier));
}
ExportEntry::Default { local_name } => {
if let Some(local_name) = local_name {
results.extend(local_export_declarations(python, &file, local_name));
}
}
}
continue;
}
if !export_name.starts_with('_') {
for star in &index.reexport_stars {
for target_file in
resolve_module_files_for_export(python, &file, &star.module_specifier)
{
queue.push_back((target_file, export_name.clone()));
}
}
}
}
results.sort_by(|left, right| {
left.source()
.cmp(right.source())
.then_with(|| left.fq_name().cmp(&right.fq_name()))
});
results.dedup();
results
}
fn local_export_declarations(
index: &dyn CodeUnitIndex,
file: &ProjectFile,
local_name: &str,
) -> Vec<CodeUnit> {
index
.top_level_declarations(file)
.into_iter()
.filter(|unit| {
unit.identifier() == local_name
&& index
.parent_of(unit)
.is_some_and(|parent| parent.is_module() && parent.source() == file)
})
.collect()
}
fn resolve_module_files_for_export(
python: &dyn PythonSource,
importing_file: &ProjectFile,
module_specifier: &str,
) -> Vec<ProjectFile> {
let resolved_module = if module_specifier.starts_with('.') {
resolve_python_relative_module(importing_file, module_specifier)
} else {
Some(module_specifier.to_string())
};
let Some(resolved_module) = resolved_module else {
return Vec::new();
};
resolve_module_code_unit(python, &resolved_module)
.map(|unit| vec![unit.source().clone()])
.unwrap_or_default()
}
pub fn extract_package_from_python_wildcard(import: &ImportInfo) -> Option<String> {
let details = python_import_details(import)?;
match details {
PythonImportDetails::FromImport {
module, wildcard, ..
} if wildcard => Some(module),
_ => None,
}
}
#[derive(Debug, Clone)]
pub enum PythonImportDetails {
Import {
module: String,
alias: Option<String>,
},
FromImport {
module: String,
name: String,
alias: Option<String>,
wildcard: bool,
},
}
pub fn python_import_infos_from_node(node: Node<'_>, source: &str) -> Vec<ImportInfo> {
match node.kind() {
"import_statement" => python_namespace_import_infos(node, source),
"import_from_statement" => python_from_import_infos(node, source),
_ => Vec::new(),
}
}
pub fn python_import_details(import: &ImportInfo) -> Option<PythonImportDetails> {
let path = import.path.as_ref()?;
match path.kind? {
StructuredImportPathKind::Namespace => Some(PythonImportDetails::Import {
module: join_python_import_segments(&path.segments),
alias: import.alias.clone(),
}),
StructuredImportPathKind::StaticMember => None,
StructuredImportPathKind::ImportFrom => {
let (name, module_segments) = if import.is_wildcard {
("*".to_string(), path.segments.as_slice())
} else {
let (name, module_segments) = path.segments.split_last()?;
(name.clone(), module_segments)
};
Some(PythonImportDetails::FromImport {
module: join_python_import_segments(module_segments),
name,
alias: import.alias.clone(),
wildcard: import.is_wildcard,
})
}
}
}
fn python_namespace_import_infos(node: Node<'_>, source: &str) -> Vec<ImportInfo> {
let mut infos = Vec::new();
let mut cursor = node.walk();
for imported in node.children_by_field_name("name", &mut cursor) {
let (module_node, alias_node) = if imported.kind() == "aliased_import" {
let Some(name) = imported.child_by_field_name("name") else {
continue;
};
(name, imported.child_by_field_name("alias"))
} else {
(imported, None)
};
let alias = alias_node
.map(|alias| py_node_text(alias, source).trim().to_string())
.filter(|alias| !alias.is_empty());
let segments = python_path_segments(module_node, source);
if segments.is_empty() {
continue;
}
let binder_span = alias
.is_some()
.then_some(alias_node)
.flatten()
.or_else(|| python_first_segment_node(module_node))
.map(brokk_bifrost_core::analyzer::common::node_span);
let module = join_python_import_segments(&segments);
let identifier = alias.clone().or_else(|| segments.first().cloned());
infos.push(ImportInfo {
raw_snippet: if let Some(alias) = &alias {
format!("import {module} as {alias}")
} else {
format!("import {module}")
},
is_wildcard: false,
is_global: false,
identifier,
alias,
path: Some(StructuredImportPath {
segments,
kind: Some(StructuredImportPathKind::Namespace),
lexical_prefixes: Vec::new(),
lexical_scopes: Vec::new(),
declaration_start_byte: node.start_byte(),
}),
binder_span,
});
}
infos
}
fn python_from_import_infos(node: Node<'_>, source: &str) -> Vec<ImportInfo> {
let Some(module_node) = node.child_by_field_name("module_name") else {
return Vec::new();
};
let module_segments = python_module_segments(module_node, source);
if module_segments.is_empty() {
return Vec::new();
}
let mut infos = Vec::new();
let has_wildcard_import = {
let mut cursor = node.walk();
node.named_children(&mut cursor)
.any(|child| child.kind() == "wildcard_import")
};
let mut cursor = node.walk();
let imported_names: Vec<_> = node.children_by_field_name("name", &mut cursor).collect();
if has_wildcard_import {
let module = join_python_import_segments(&module_segments);
infos.push(ImportInfo {
raw_snippet: format!("from {module} import *"),
is_wildcard: true,
is_global: false,
identifier: None,
alias: None,
path: Some(StructuredImportPath {
segments: module_segments,
kind: Some(StructuredImportPathKind::ImportFrom),
lexical_prefixes: Vec::new(),
lexical_scopes: Vec::new(),
declaration_start_byte: node.start_byte(),
}),
binder_span: None,
});
return infos;
}
if imported_names.is_empty() {
return infos;
}
for imported in imported_names {
let (name_node, alias_node) = if imported.kind() == "aliased_import" {
let Some(name) = imported.child_by_field_name("name") else {
continue;
};
(name, imported.child_by_field_name("alias"))
} else {
(imported, None)
};
let alias = alias_node
.map(|alias| py_node_text(alias, source).trim().to_string())
.filter(|alias| !alias.is_empty());
let name_segments = python_path_segments(name_node, source);
if name_segments.is_empty() {
continue;
}
let binder_span = alias
.is_some()
.then_some(alias_node)
.flatten()
.or_else(|| {
(name_segments.len() == 1)
.then(|| python_first_segment_node(name_node))
.flatten()
})
.map(brokk_bifrost_core::analyzer::common::node_span);
let imported_name = join_python_import_segments(&name_segments);
let mut segments = module_segments.clone();
segments.extend(name_segments);
let module = join_python_import_segments(&module_segments);
infos.push(ImportInfo {
raw_snippet: if let Some(alias) = &alias {
format!("from {module} import {imported_name} as {alias}")
} else {
format!("from {module} import {imported_name}")
},
is_wildcard: false,
is_global: false,
identifier: Some(alias.clone().unwrap_or_else(|| imported_name.clone())),
alias,
path: Some(StructuredImportPath {
segments,
kind: Some(StructuredImportPathKind::ImportFrom),
lexical_prefixes: Vec::new(),
lexical_scopes: Vec::new(),
declaration_start_byte: node.start_byte(),
}),
binder_span,
});
}
infos
}
fn python_module_segments(module: Node<'_>, source: &str) -> Vec<String> {
if module.kind() == "relative_import" {
let mut cursor = module.walk();
let mut prefix = String::new();
let mut path_node = None;
for child in module.named_children(&mut cursor) {
match child.kind() {
"import_prefix" if prefix.is_empty() => {
prefix = py_node_text(child, source).trim().to_string();
}
"dotted_name" if path_node.is_none() => {
path_node = Some(child);
}
_ => {}
}
}
let mut segments = path_node
.map(|path| python_path_segments(path, source))
.unwrap_or_default();
if !prefix.is_empty() {
if let Some(first) = segments.first_mut() {
first.insert_str(0, &prefix);
} else {
segments.push(prefix);
}
}
return segments;
}
python_path_segments(module, source)
}
fn python_first_segment_node(node: Node<'_>) -> Option<Node<'_>> {
match node.kind() {
"identifier" => Some(node),
"dotted_name" => {
let mut cursor = node.walk();
node.named_children(&mut cursor)
.find(|child| child.kind() == "identifier")
}
_ => None,
}
}
fn python_path_segments(node: Node<'_>, source: &str) -> Vec<String> {
match node.kind() {
"identifier" => vec![py_node_text(node, source).trim().to_string()],
"dotted_name" => {
let mut segments = Vec::new();
let mut cursor = node.walk();
for child in node.named_children(&mut cursor) {
segments.extend(python_path_segments(child, source));
}
segments
}
_ => {
let mut segments = Vec::new();
let mut cursor = node.walk();
for child in node.named_children(&mut cursor) {
segments.extend(python_path_segments(child, source));
}
segments
}
}
}
fn join_python_import_segments(segments: &[String]) -> String {
let Some((first, rest)) = segments.split_first() else {
return String::new();
};
if first.starts_with('.') && !rest.is_empty() {
format!("{first}.{}", rest.join("."))
} else {
segments.join(".")
}
}
pub fn python_namespace_binding_name(
import: &ImportInfo,
alias: Option<&str>,
module: &str,
) -> String {
import
.identifier
.clone()
.or_else(|| alias.map(str::to_string))
.unwrap_or_else(|| module.to_string())
}
pub fn python_namespace_binding_module(
import: &ImportInfo,
alias: Option<&str>,
module: &str,
) -> String {
if alias.is_some() {
return module.to_string();
}
import
.path
.as_ref()
.and_then(|path| path.segments.first().cloned())
.unwrap_or_else(|| module.to_string())
}
pub fn resolve_python_relative_module(
source_file: &ProjectFile,
module_expr: &str,
) -> Option<String> {
resolve_python_relative_module_from_package(&python_current_package(source_file), module_expr)
}
pub fn resolve_python_relative_module_from_package(
current_package: &str,
module_expr: &str,
) -> Option<String> {
let level = module_expr.chars().take_while(|ch| *ch == '.').count();
let suffix = module_expr[level..].trim_matches('.');
let mut parts: Vec<_> = current_package
.split('.')
.filter(|part| !part.is_empty())
.map(str::to_string)
.collect();
if level == 0 {
return Some(module_expr.to_string());
}
if level > 0 {
if level - 1 > parts.len() {
return None;
}
parts.truncate(parts.len() - (level - 1));
}
if !suffix.is_empty() {
parts.extend(suffix.split('.').map(str::to_string));
}
Some(parts.join("."))
}
fn python_current_package(source_file: &ProjectFile) -> String {
let module = python_module_name(source_file);
if source_file
.rel_path()
.file_name()
.and_then(|name| name.to_str())
== Some("__init__.py")
{
module
} else {
module
.rsplit_once('.')
.map(|(package, _)| package.to_string())
.unwrap_or_default()
}
}