use brokk_bifrost_core::analyzer::capabilities::ImportAnalysisProvider;
use brokk_bifrost_core::analyzer::model::ImportInfo;
use brokk_bifrost_core::analyzer::prepared_syntax::{IndexedFileFacts, PreparedSyntaxTree};
use brokk_bifrost_core::analyzer::project::Project;
use brokk_bifrost_core::analyzer::query_token::QueryToken;
use brokk_bifrost_core::analyzer::symbol_path::parse_symbol_path;
use brokk_bifrost_core::analyzer::tree_walk::{WalkControl, walk_named_tree_preorder};
use brokk_bifrost_core::analyzer::usages::model::{
ExportEntry, ExportIndex, ImportBinder, ImportBinding, ImportKind, ReexportStar,
};
use brokk_bifrost_core::analyzer::{CodeUnit, CodeUnitIndex, Language, ProjectFile};
use brokk_bifrost_core::hash::{HashMap, HashSet};
use std::path::{Path, PathBuf};
use std::sync::Arc;
use tree_sitter::Node;
use crate::bindings::{python_direct_scope_bindings_bounded, python_executable_scope_children};
use crate::declarations::{
module_code_unit, parse_python_tree, python_base_origin_node, python_declared_import_root,
python_import_root, python_project_module_components, python_repeated_root_extension_dir,
};
use crate::imports::{
PythonImportDetails, python_import_details, python_namespace_binding_module,
python_namespace_binding_name, resolve_exported_fqn, resolve_import_bindings,
resolve_python_relative_module,
};
use crate::syntax::{python_plain_string_literal, python_static_attribute_path};
use crate::usage_index::PythonUsageIndex;
pub trait PythonSource:
CodeUnitIndex + ImportAnalysisProvider + crate::source_facts::PythonSourceFactProvider
{
fn path_module_fqn(&self, module_fq: &str) -> Option<Vec<CodeUnit>>;
fn path_module_fqns_batch(&self, module_fqs: &[String]) -> Vec<Option<Vec<CodeUnit>>>;
fn definition_fqn(&self, fqn: &str) -> Vec<CodeUnit>;
fn module_spellings(&self) -> Arc<PythonModuleSpellings>;
fn import_binder_of(&self, file: &ProjectFile) -> Arc<ImportBinder>;
fn export_index_of(&self, file: &ProjectFile) -> Arc<ExportIndex>;
fn prepared_syntax(
&self,
token: QueryToken<'_>,
file: &ProjectFile,
) -> Option<Arc<PreparedSyntaxTree>>;
fn visit_file_facts(
&self,
files: &[ProjectFile],
visit: &mut dyn FnMut(&ProjectFile, Option<&dyn IndexedFileFacts>),
);
}
pub trait PythonUsageSource: PythonSource {
fn usage_index(&self) -> Arc<PythonUsageIndex>;
}
pub fn retain_modules_for_importer<T>(
project: &dyn Project,
importer: &ProjectFile,
candidates: &mut Vec<T>,
source_of: impl Fn(&T) -> &ProjectFile,
) {
retain_nearest_owning_root(importer, candidates, |candidate| {
let source = source_of(candidate);
debug_assert_eq!(
importer.root(),
source.root(),
"one workspace's Python files share one project root"
);
python_import_root(project, source)
});
}
fn retain_nearest_owning_root<T>(
importer: &ProjectFile,
candidates: &mut Vec<T>,
root_of: impl Fn(&T) -> PathBuf,
) {
if candidates.len() < 2 {
return;
}
let owning_depth = |candidate: &T| {
let root = root_of(candidate);
importer
.rel_path()
.starts_with(&root)
.then(|| root.components().count())
};
let depths = candidates.iter().map(owning_depth).collect::<Vec<_>>();
let Some(nearest) = depths.iter().copied().flatten().max() else {
return;
};
*candidates = std::mem::take(candidates)
.into_iter()
.zip(depths)
.filter(|(_, depth)| *depth == Some(nearest))
.map(|(candidate, _)| candidate)
.collect();
}
#[derive(Clone, Debug)]
struct ProjectRootModuleName {
file: ProjectFile,
components: Vec<String>,
}
#[derive(Clone, Debug)]
pub struct PythonSourceRootModule {
pub file: ProjectFile,
pub source_root: PathBuf,
}
#[derive(Default)]
pub struct PythonModuleSpellings {
project: Option<Arc<dyn Project>>,
by_last_component: HashMap<String, Vec<ProjectRootModuleName>>,
package_dirs: HashSet<PathBuf>,
path_extensions: HashMap<PathBuf, HashSet<PathBuf>>,
}
impl std::fmt::Debug for PythonModuleSpellings {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("PythonModuleSpellings")
.field("by_last_component", &self.by_last_component)
.field("package_dirs", &self.package_dirs)
.field("path_extensions", &self.path_extensions)
.finish_non_exhaustive()
}
}
impl PythonModuleSpellings {
pub(crate) fn project(&self) -> Option<&dyn Project> {
self.project.as_deref()
}
pub fn build(files: impl IntoIterator<Item = ProjectFile>, project: Arc<dyn Project>) -> Self {
let mut by_last_component: HashMap<String, Vec<ProjectRootModuleName>> = HashMap::default();
let mut package_dirs: HashSet<PathBuf> = HashSet::default();
let mut package_markers: Vec<ProjectFile> = Vec::new();
for file in files {
if file.rel_path().file_name().and_then(|name| name.to_str()) == Some("__init__.py")
&& let Some(directory) = file.rel_path().parent()
{
package_dirs.insert(directory.to_path_buf());
package_markers.push(file.clone());
}
let components = python_project_module_components(&file);
let Some(last) = components.last() else {
continue;
};
by_last_component
.entry(last.clone())
.or_default()
.push(ProjectRootModuleName { file, components });
}
let mut path_extensions: HashMap<PathBuf, HashSet<PathBuf>> = HashMap::default();
for marker in package_markers {
let Some(directory) = marker.rel_path().parent() else {
continue;
};
if !python_package_holds_a_repeat_of_itself(directory, &package_dirs) {
continue;
}
let Ok(source) = project.read_source(&marker) else {
continue;
};
let extended = python_appended_path_dirs(&source, marker.rel_path());
if !extended.is_empty() {
path_extensions
.entry(directory.to_path_buf())
.or_default()
.extend(extended);
}
}
Self {
project: Some(project),
by_last_component,
package_dirs,
path_extensions,
}
}
pub fn suffix_candidates(&self, module_fq: &str) -> Vec<PythonSourceRootModule> {
if module_fq.starts_with('.') {
return Vec::new();
}
let components = parse_symbol_path(Language::Python, module_fq);
let Some(last) = components.last().filter(|_| components.len() > 1) else {
return Vec::new();
};
let mut modules = Vec::new();
for candidate in self
.by_last_component
.get(last.as_str())
.into_iter()
.flatten()
{
let Some(dropped) = candidate
.components
.len()
.checked_sub(components.len())
.filter(|dropped| *dropped > 0)
else {
continue;
};
if candidate.components[dropped..] != components[..] {
continue;
}
let source_root: PathBuf = candidate
.file
.rel_path()
.components()
.take(dropped)
.collect();
debug_assert!(
!source_root.as_os_str().is_empty(),
"dropping a component of {:?} leaves a directory",
candidate.file.rel_path()
);
modules.push(PythonSourceRootModule {
file: candidate.file.clone(),
source_root,
});
}
modules
}
pub fn establishes_source_root(&self, candidate: &PythonSourceRootModule) -> bool {
if let Some(declared) = python_declared_import_root(
self.project()
.expect("source-root candidates have a selected project"),
&candidate.file,
) {
if declared == candidate.source_root {
return true;
}
let Some(extension) = python_repeated_root_extension_dir(
&candidate.file,
&candidate.source_root,
&declared,
) else {
return false;
};
return self.package_chain_reaches_module(candidate)
&& self
.path_extensions
.get(&candidate.source_root)
.is_some_and(|extended| extended.contains(&extension));
}
self.package_chain_reaches_module(candidate)
}
fn package_chain_reaches_module(&self, candidate: &PythonSourceRootModule) -> bool {
debug_assert!(
candidate
.file
.rel_path()
.starts_with(&candidate.source_root),
"a candidate's implied root is a prefix of its own path"
);
let mut directory = candidate.file.rel_path().parent();
while let Some(current) = directory {
if !self.package_dirs.contains(current) {
return false;
}
if current == candidate.source_root {
return true;
}
directory = current.parent();
}
false
}
}
fn python_package_holds_a_repeat_of_itself(
directory: &Path,
package_dirs: &HashSet<PathBuf>,
) -> bool {
directory
.file_name()
.is_some_and(|name| package_dirs.contains(&directory.join(name)))
}
fn python_appended_path_dirs(source: &str, marker_rel_path: &Path) -> Vec<PathBuf> {
let Some(tree) = parse_python_tree(source) else {
return Vec::new();
};
let mut bound: HashMap<String, PathBuf> = HashMap::default();
bound.insert("__file__".to_owned(), marker_rel_path.to_path_buf());
let mut appended: Vec<PathBuf> = Vec::new();
walk_named_tree_preorder(tree.root_node(), true, |node| {
match node.kind() {
"function_definition" | "class_definition" | "lambda" => {
return WalkControl::SkipChildren;
}
"assignment" => {
if python_names_path(node.child_by_field_name("left"), source) {
appended.clear();
}
python_bind_assignment(node, source, &mut bound);
}
"call" => {
let directories = python_path_appended_dirs_in_call(node, source, &bound);
if directories.is_empty()
|| !python_extension_statement_runs(node, source, &bound, &directories)
{
return WalkControl::Continue;
}
for directory in directories {
if directory.is_relative()
&& !directory.as_os_str().is_empty()
&& !appended.contains(&directory)
{
appended.push(directory);
}
}
}
_ => {}
}
WalkControl::Continue
});
appended
}
fn python_names_path(node: Option<Node<'_>>, source: &str) -> bool {
node.is_some_and(|node| {
node.kind() == "identifier" && node.utf8_text(source.as_bytes()).ok() == Some("__path__")
})
}
fn python_extension_statement_runs(
call: Node<'_>,
source: &str,
bound: &HashMap<String, PathBuf>,
directories: &[PathBuf],
) -> bool {
let Some(statement) = call.parent() else {
return false;
};
if statement.kind() != "expression_statement" {
return false;
}
let Some(container) = statement.parent() else {
return false;
};
if container.kind() == "module" {
return true;
}
let guarded = if container.kind() == "block" {
container.parent().filter(|parent| {
parent.kind() == "if_statement"
&& parent.child_by_field_name("consequence") == Some(container)
})
} else if container.kind() == "if_statement"
&& container.child_by_field_name("consequence") == Some(statement)
{
Some(container)
} else {
None
};
let Some(guarded) = guarded else {
return false;
};
if guarded.parent().map(|parent| parent.kind()) != Some("module") {
return false;
}
let Some(condition) = guarded.child_by_field_name("condition") else {
return false;
};
python_condition_justifies_path_extension(condition, source, bound, directories)
}
fn python_condition_justifies_path_extension(
condition: Node<'_>,
source: &str,
bound: &HashMap<String, PathBuf>,
directories: &[PathBuf],
) -> bool {
if condition.kind() != "comparison_operator" {
return false;
}
let mut cursor = condition.walk();
let operators = condition
.children_by_field_name("operators", &mut cursor)
.map(|operator| operator.kind())
.collect::<Vec<_>>();
if operators.as_slice() != ["not in"] {
return false;
}
let mut cursor = condition.walk();
let operands: Vec<Node<'_>> = condition.named_children(&mut cursor).collect();
let [left, right] = operands.as_slice() else {
return false;
};
if !python_names_path(Some(*right), source) {
return false;
}
python_path_expression_dir(*left, source, bound)
.is_some_and(|directory| directories.contains(&directory))
}
fn python_bind_assignment(
assignment: Node<'_>,
source: &str,
bound: &mut HashMap<String, PathBuf>,
) {
let mut targets: Vec<Node<'_>> = Vec::new();
let mut current = assignment;
loop {
let (Some(left), Some(right)) = (
current.child_by_field_name("left"),
current.child_by_field_name("right"),
) else {
return;
};
targets.push(left);
if right.kind() != "assignment" {
let resolved = python_path_expression_dir(right, source, bound);
for target in targets {
if target.kind() != "identifier" {
continue;
}
let Ok(name) = target.utf8_text(source.as_bytes()) else {
continue;
};
match &resolved {
Some(directory) => {
bound.insert(name.to_owned(), directory.clone());
}
None => {
bound.remove(name);
}
}
}
return;
}
current = right;
}
}
fn python_path_appended_dirs_in_call(
call: Node<'_>,
source: &str,
bound: &HashMap<String, PathBuf>,
) -> Vec<PathBuf> {
let Some(function) = call.child_by_field_name("function") else {
return Vec::new();
};
let Some(path) = python_static_attribute_path(function) else {
return Vec::new();
};
let [object, method] = path.as_slice() else {
return Vec::new();
};
if object.utf8_text(source.as_bytes()).ok() != Some("__path__") {
return Vec::new();
}
let Ok(method) = method.utf8_text(source.as_bytes()) else {
return Vec::new();
};
let Some(arguments) = call.child_by_field_name("arguments") else {
return Vec::new();
};
let mut cursor = arguments.walk();
let arguments: Vec<Node<'_>> = arguments.named_children(&mut cursor).collect();
let named: Vec<Node<'_>> = match method {
"append" => arguments.first().copied().into_iter().collect(),
"insert" => arguments.get(1).copied().into_iter().collect(),
"extend" => match arguments.first() {
Some(list) if list.kind() == "list" => {
let mut cursor = list.walk();
list.named_children(&mut cursor).collect()
}
Some(iterator) => vec![*iterator],
None => Vec::new(),
},
_ => Vec::new(),
};
named
.into_iter()
.filter_map(|node| python_path_expression_dir(node, source, bound))
.collect()
}
fn python_path_expression_dir(
node: Node<'_>,
source: &str,
bound: &HashMap<String, PathBuf>,
) -> Option<PathBuf> {
if let Some(literal) = python_plain_string_literal(node, source) {
let directory = PathBuf::from(literal);
return (!directory.as_os_str().is_empty()).then_some(directory);
}
match node.kind() {
"identifier" => bound.get(node.utf8_text(source.as_bytes()).ok()?).cloned(),
"call" => python_path_call_dir(node, source, bound),
_ => None,
}
}
fn python_path_call_dir(
call: Node<'_>,
source: &str,
bound: &HashMap<String, PathBuf>,
) -> Option<PathBuf> {
let path = python_static_attribute_path(call.child_by_field_name("function")?)?;
let names: Vec<&str> = path
.iter()
.map(|node| node.utf8_text(source.as_bytes()).ok())
.collect::<Option<Vec<_>>>()?;
let arguments = call.child_by_field_name("arguments")?;
let mut cursor = arguments.walk();
let arguments: Vec<Node<'_>> = arguments.named_children(&mut cursor).collect();
match names.as_slice() {
["os", "path", "dirname"] => {
let directory = python_path_expression_dir(*arguments.first()?, source, bound)?;
Some(directory.parent()?.to_path_buf())
}
["os", "path", "join"] => {
let mut parts = arguments.iter();
let mut joined = python_path_expression_dir(*parts.next()?, source, bound)?;
for part in parts {
joined.push(python_path_expression_dir(*part, source, bound)?);
}
Some(joined)
}
_ => None,
}
}
pub fn resolve_module_code_unit(
python: &dyn PythonSource,
importer: Option<&ProjectFile>,
module_fq: &str,
) -> Option<CodeUnit> {
let indexed = match python.path_module_fqn(module_fq) {
Some(units) => indexed_module_for_importer(python.project(), importer, units),
None => indexed_module_for_importer(
python.project(),
importer,
python.definition_fqn(module_fq),
),
};
match indexed {
IndexedModuleLookup::Unique(module) => Some(module),
IndexedModuleLookup::Ambiguous => None,
IndexedModuleLookup::Missing => {
let importer = importer?;
module_under_source_root(&python.module_spellings(), importer, module_fq)
}
}
}
enum IndexedModuleLookup {
Unique(CodeUnit),
Ambiguous,
Missing,
}
fn indexed_module_for_importer(
project: &dyn Project,
importer: Option<&ProjectFile>,
units: Vec<CodeUnit>,
) -> IndexedModuleLookup {
let mut modules = units
.into_iter()
.filter(CodeUnit::is_module)
.collect::<Vec<_>>();
if let Some(importer) = importer {
retain_modules_for_importer(project, importer, &mut modules, CodeUnit::source);
}
match modules.len() {
0 => IndexedModuleLookup::Missing,
1 => IndexedModuleLookup::Unique(modules.remove(0)),
_ => IndexedModuleLookup::Ambiguous,
}
}
pub fn python_module_file_under_source_root(
spellings: &PythonModuleSpellings,
importer: &ProjectFile,
module_fq: &str,
) -> Option<ProjectFile> {
let mut candidates = spellings.suffix_candidates(module_fq);
retain_nearest_owning_root(importer, &mut candidates, |candidate| {
candidate.source_root.clone()
});
candidates.retain(|candidate| spellings.establishes_source_root(candidate));
let [candidate] = candidates.as_slice() else {
return None;
};
debug_assert_eq!(
importer.root(),
candidate.file.root(),
"one workspace's Python files share one project root"
);
Some(candidate.file.clone())
}
fn module_under_source_root(
spellings: &PythonModuleSpellings,
importer: &ProjectFile,
module_fq: &str,
) -> Option<CodeUnit> {
let file = python_module_file_under_source_root(spellings, importer, module_fq)?;
module_code_unit(spellings.project()?, &file)
}
pub fn resolve_module_code_units_batch(
python: &dyn PythonSource,
importer: Option<&ProjectFile>,
module_fqs: &[String],
) -> Vec<Option<CodeUnit>> {
let path_results = python.path_module_fqns_batch(module_fqs);
let mut lookups: Vec<IndexedModuleLookup> = Vec::with_capacity(module_fqs.len());
let mut needs_definition_fallback = Vec::new();
for (i, units) in path_results.into_iter().enumerate() {
match units {
Some(units) => lookups.push(indexed_module_for_importer(
python.project(),
importer,
units,
)),
None => {
needs_definition_fallback.push(i);
lookups.push(IndexedModuleLookup::Missing);
}
}
}
for i in needs_definition_fallback {
lookups[i] = indexed_module_for_importer(
python.project(),
importer,
python.definition_fqn(&module_fqs[i]),
);
}
let mut results: Vec<Option<CodeUnit>> = vec![None; module_fqs.len()];
let mut needs_source_root = Vec::new();
for (i, lookup) in lookups.into_iter().enumerate() {
match lookup {
IndexedModuleLookup::Unique(module) => results[i] = Some(module),
IndexedModuleLookup::Ambiguous => {}
IndexedModuleLookup::Missing => needs_source_root.push(i),
}
}
if let Some(importer) = importer
&& !needs_source_root.is_empty()
{
let spellings = python.module_spellings();
for i in needs_source_root {
results[i] = module_under_source_root(&spellings, importer, &module_fqs[i]);
}
}
results
}
pub fn compute_export_index_of(
python: &dyn PythonSource,
token: QueryToken<'_>,
file: &ProjectFile,
) -> ExportIndex {
let mut index = ExportIndex::empty();
let mut events = Vec::new();
let declarations = python.top_level_declarations(file);
collect_local_export_events(
declarations.iter(),
|code_unit| {
python
.ranges(code_unit)
.iter()
.map(|range| range.start_byte)
.max()
.unwrap_or(usize::MAX)
},
&mut events,
);
let imports = python.import_info_of(token, file);
collect_reexport_events_from_imports(python, file, &imports, &mut events, &mut index);
finish_export_index(events, index)
}
pub fn export_index_from_file_facts(
python: &dyn PythonSource,
file: &ProjectFile,
facts: &dyn IndexedFileFacts,
module_name: &str,
) -> ExportIndex {
let mut index = ExportIndex::empty();
let mut events = Vec::new();
collect_local_export_events(
facts.top_level_declarations().iter(),
|code_unit| {
facts
.declaration_ranges(code_unit)
.into_iter()
.flatten()
.map(|range| range.start_byte)
.max()
.unwrap_or(usize::MAX)
},
&mut events,
);
if !facts
.top_level_declarations()
.iter()
.any(CodeUnit::is_module)
&& let Some(identifier) = module_name.rsplit('.').next()
&& !identifier.is_empty()
&& !identifier.starts_with('_')
{
events.push((
0,
identifier.to_string(),
ExportEntry::Local {
local_name: identifier.to_string(),
},
));
}
collect_reexport_events_from_imports(python, file, facts.imports(), &mut events, &mut index);
finish_export_index(events, index)
}
fn collect_local_export_events<'a>(
declarations: impl IntoIterator<Item = &'a CodeUnit>,
mut start_byte: impl FnMut(&CodeUnit) -> usize,
events: &mut Vec<(usize, String, ExportEntry)>,
) -> HashSet<String> {
let mut local_names = HashSet::default();
for code_unit in declarations {
let identifier = code_unit.identifier().trim();
if identifier.is_empty() {
continue;
}
local_names.insert(identifier.to_string());
events.push((
start_byte(code_unit),
identifier.to_string(),
ExportEntry::Local {
local_name: identifier.to_string(),
},
));
}
local_names
}
fn finish_export_index(
mut events: Vec<(usize, String, ExportEntry)>,
mut index: ExportIndex,
) -> ExportIndex {
events.sort_by_key(|(start_byte, _, _)| *start_byte);
for (_, exported_name, entry) in events {
index.exports_by_name.insert(exported_name, entry);
}
index
}
fn collect_reexport_events_from_imports(
python: &dyn PythonSource,
file: &ProjectFile,
imports: &[ImportInfo],
events: &mut Vec<(usize, String, ExportEntry)>,
index: &mut ExportIndex,
) {
for import in imports {
if import
.path
.as_ref()
.is_none_or(|path| !path.lexical_scopes.is_empty())
{
continue;
}
record_single_reexport_event(python, file, import, events, index);
}
}
fn record_single_reexport_event(
python: &dyn PythonSource,
file: &ProjectFile,
import: &ImportInfo,
events: &mut Vec<(usize, String, ExportEntry)>,
index: &mut ExportIndex,
) {
let Some(PythonImportDetails::FromImport {
module,
name,
alias,
wildcard,
}) = python_import_details(import)
else {
return;
};
let start_byte = import
.path
.as_ref()
.map(|path| path.declaration_start_byte)
.unwrap_or(usize::MAX);
let resolved_module = if module.starts_with('.') {
resolve_python_relative_module(python.project(), file, &module)
} else {
Some(module.clone())
};
let Some(resolved_module) = resolved_module else {
return;
};
if wildcard {
index.reexport_stars.push(ReexportStar {
module_specifier: resolved_module,
});
return;
}
let exported_name = alias.unwrap_or(name.clone());
let module_candidate = format!("{resolved_module}.{name}");
if resolve_module_code_unit(python, Some(file), &module_candidate).is_some() {
events.push((
start_byte,
exported_name,
ExportEntry::ReexportedModule {
module_specifier: module_candidate,
},
));
return;
}
events.push((
start_byte,
exported_name,
ExportEntry::ReexportedNamed {
module_specifier: resolved_module,
imported_name: name,
},
));
}
pub fn import_binder_from_imports(
python: &dyn PythonSource,
file: &ProjectFile,
imports: &[ImportInfo],
) -> ImportBinder {
let mut binder = ImportBinder::empty();
for (local_name, binding) in import_bindings_from_imports(python, file, imports) {
binder.bindings.insert(local_name, binding);
}
binder
}
pub fn import_bindings_from_imports(
python: &dyn PythonSource,
file: &ProjectFile,
imports: &[ImportInfo],
) -> Vec<(String, ImportBinding)> {
let mut bindings = Vec::new();
for import in imports {
let Some(details) = python_import_details(import) else {
continue;
};
match details {
PythonImportDetails::Import { module, alias } => {
let local_name = python_namespace_binding_name(import, alias.as_deref(), &module);
let module_specifier =
python_namespace_binding_module(import, alias.as_deref(), &module);
bindings.push((
local_name,
ImportBinding {
module_specifier,
namespace_imported_module: Some(module),
kind: ImportKind::Namespace,
imported_name: None,
},
));
}
PythonImportDetails::FromImport {
module,
name,
wildcard,
..
} => {
let resolved_module = if module.starts_with('.') {
resolve_python_relative_module(python.project(), file, &module)
} else {
Some(module.clone())
};
let Some(resolved_module) = resolved_module else {
continue;
};
if wildcard {
bindings.extend(
public_declarations_in_module(python, file, &resolved_module)
.into_iter()
.map(|declaration| {
let name = declaration.identifier().to_string();
(
name.clone(),
ImportBinding {
module_specifier: resolved_module.clone(),
namespace_imported_module: None,
kind: ImportKind::Named,
imported_name: Some(name),
},
)
}),
);
continue;
}
let local_name = import
.local_name()
.map(str::to_string)
.unwrap_or_else(|| name.clone());
let module_candidate = format!("{resolved_module}.{name}");
let exported = resolve_exported_fqn(python, Some(file), &module_candidate);
let exported_module =
if !exported.is_empty() && exported.iter().all(CodeUnit::is_module) {
let mut paths = exported.iter().map(CodeUnit::fq_segment_texts);
let path = paths
.next()
.expect("non-empty exported modules have a first path");
paths
.all(|candidate| candidate == path)
.then(|| path.join("."))
} else {
None
};
let implicit_submodule = if exported.is_empty()
&& resolve_module_code_unit(python, Some(file), &module_candidate).is_some()
{
Some(module_candidate)
} else {
None
};
if let Some(module_specifier) = exported_module.or(implicit_submodule) {
bindings.push((
local_name,
ImportBinding {
module_specifier,
namespace_imported_module: None,
kind: ImportKind::Namespace,
imported_name: None,
},
));
continue;
}
bindings.push((
local_name,
ImportBinding {
module_specifier: resolved_module,
namespace_imported_module: None,
kind: ImportKind::Named,
imported_name: Some(name),
},
));
}
}
}
bindings
}
pub fn public_declarations_in_module(
python: &dyn PythonSource,
importer: &ProjectFile,
module_fq: &str,
) -> Vec<CodeUnit> {
let Some(module_code_unit) = resolve_module_code_unit(python, Some(importer), module_fq) else {
return Vec::new();
};
python
.direct_children(&module_code_unit)
.into_iter()
.filter(|code_unit| !code_unit.identifier().starts_with('_'))
.collect()
}
pub fn resolve_base_class(
python: &dyn PythonSource,
token: QueryToken<'_>,
code_unit: &CodeUnit,
raw: &str,
) -> Option<CodeUnit> {
let trimmed = raw.trim();
if trimmed.is_empty() {
return None;
}
let binder = python.import_binder_of(code_unit.source());
if let Some((head, tail)) = trimmed.split_once('.') {
if let Some(binding) = binder.bindings.get(head)
&& binding.kind == ImportKind::Namespace
{
let fq_name = format!("{}.{}", binding.module_specifier, tail);
return python.definitions(&fq_name).next();
}
return python.definitions(trimmed).next();
}
if let Some(binding) = binder.bindings.get(trimmed) {
match binding.kind {
ImportKind::Namespace => {
return resolve_module_code_unit(
python,
Some(code_unit.source()),
&binding.module_specifier,
);
}
ImportKind::Named => {
let imported_name = binding.imported_name.as_ref()?;
let fqn = format!("{}.{}", binding.module_specifier, imported_name);
return resolve_exported_fqn(python, Some(code_unit.source()), &fqn)
.into_iter()
.next()
.or_else(|| python.definitions(&fqn).next());
}
_ => {}
}
}
if python
.import_info_of(token, code_unit.source())
.iter()
.any(|import| import.is_wildcard)
&& let Some(imported) =
resolve_import_bindings(python, token, code_unit.source()).get(trimmed)
{
return Some(imported.clone());
}
if let Some(local) = resolve_preceding_local_class_base(python, token, code_unit, trimmed) {
return Some(local);
}
let local_fq_name = format!("{}.{}", code_unit.package_name(), trimmed);
python
.definitions(&local_fq_name)
.next()
.or_else(|| python.definitions(trimmed).next())
}
fn class_definition_in_statement(node: Node<'_>) -> Option<Node<'_>> {
if node.kind() == "class_definition" {
return Some(node);
}
(node.kind() == "decorated_definition")
.then(|| {
let mut cursor = node.walk();
node.named_children(&mut cursor)
.find(|child| child.kind() == "class_definition")
})
.flatten()
}
fn declaration_statement(node: Node<'_>) -> Option<Node<'_>> {
match node.kind() {
"decorated_definition" if class_definition_in_statement(node).is_some() => Some(node),
"class_definition" => Some(
node.parent()
.filter(|parent| parent.kind() == "decorated_definition")
.unwrap_or(node),
),
_ => None,
}
}
fn scope_statement_binds_name(node: Node<'_>, source: &str, name: &str) -> bool {
let mut stack = vec![node];
while let Some(node) = stack.pop() {
if node.kind() == "wildcard_import" {
return true;
}
if matches!(node.kind(), "global_statement" | "nonlocal_statement") {
let mut cursor = node.walk();
if node.named_children(&mut cursor).any(|child| {
child.kind() == "identifier"
&& child.utf8_text(source.as_bytes()).ok() == Some(name)
}) {
return true;
}
}
let Some(bindings) = python_direct_scope_bindings_bounded(node, source, || true) else {
return true;
};
if bindings
.iter()
.any(|binding| binding.declaration.utf8_text(source.as_bytes()).ok() == Some(name))
{
return true;
}
stack.extend(python_executable_scope_children(node));
}
false
}
fn resolve_preceding_local_class_base(
python: &dyn PythonSource,
token: QueryToken<'_>,
code_unit: &CodeUnit,
raw: &str,
) -> Option<CodeUnit> {
let parent = python.parent_of(code_unit)?;
if !parent.is_function() && !parent.is_class() {
return None;
}
let prepared = python.prepared_syntax(token, code_unit.source())?;
let source = prepared.source();
let current_node = prepared.declaration_node(code_unit)?;
let current_statement = declaration_statement(current_node)?;
let current_class = class_definition_in_statement(current_statement)?;
let base_list = current_class.child_by_field_name("superclasses")?;
let mut base_cursor = base_list.walk();
if !base_list.named_children(&mut base_cursor).any(|base| {
let origin = python_base_origin_node(base);
origin.kind() == "identifier" && origin.utf8_text(source.as_bytes()).ok() == Some(raw)
}) {
return None;
}
let block = current_statement.parent()?;
let mut block_cursor = block.walk();
let siblings = block.named_children(&mut block_cursor).collect::<Vec<_>>();
let current_index = siblings
.iter()
.position(|sibling| sibling.id() == current_statement.id())?;
let mut preceding: Option<(usize, CodeUnit, Node<'_>)> = None;
for candidate in prepared.direct_children(&parent) {
if !candidate.is_class()
|| candidate.source() != code_unit.source()
|| candidate.identifier() != raw
{
continue;
}
let Some(candidate_node) = prepared.declaration_node(candidate) else {
continue;
};
let Some(candidate_statement) = declaration_statement(candidate_node) else {
continue;
};
if candidate_statement.parent().map(|node| node.id()) != Some(block.id())
|| candidate_statement.start_byte() >= current_statement.start_byte()
{
continue;
}
if preceding
.as_ref()
.is_none_or(|(start, _, _)| candidate_statement.start_byte() > *start)
{
preceding = Some((
candidate_statement.start_byte(),
candidate.clone(),
candidate_statement,
));
}
}
let (_, candidate, candidate_statement) = preceding?;
let candidate_index = siblings
.iter()
.position(|sibling| sibling.id() == candidate_statement.id())?;
if candidate_index >= current_index
|| siblings[candidate_index + 1..current_index]
.iter()
.any(|statement| scope_statement_binds_name(*statement, source, raw))
{
return None;
}
Some(candidate)
}
pub fn render_skeleton_recursive(
index: &dyn CodeUnitIndex,
code_unit: &CodeUnit,
indent: &str,
header_only: bool,
out: &mut String,
) {
enum Work {
Declaration(CodeUnit, String),
Elision(String),
}
let mut stack = vec![Work::Declaration(code_unit.clone(), indent.to_string())];
while let Some(work) = stack.pop() {
let Work::Declaration(unit, indent) = work else {
let Work::Elision(indent) = work else {
unreachable!()
};
out.push_str(&indent);
out.push_str("[...]\n");
continue;
};
if !unit.is_module()
&& !unit.is_file_scope()
&& let Some(signature) = index.signatures(&unit).first()
{
let signature = if unit.is_field() {
signature.strip_suffix(" ...").unwrap_or(signature)
} else {
signature
};
for line in signature.lines() {
out.push_str(&indent);
out.push_str(line);
out.push('\n');
}
}
let children = index.direct_children(&unit);
let child_indent = format!("{indent} ");
if header_only && children.iter().any(|child| !child.is_field()) {
stack.push(Work::Elision(child_indent.clone()));
}
for child in children
.into_iter()
.rev()
.filter(|child| !header_only || child.is_field())
{
stack.push(Work::Declaration(child, child_indent.clone()));
}
}
}