use crate::graph::PythonGraphSource;
use crate::graph_support::PythonUsageSource;
use crate::imports::resolve_fqn_candidates;
use crate::syntax::{
python_deferred_annotation_identifier_ranges, python_deferred_annotation_tree,
python_node_is_in_annotation,
};
use crate::usage_index::usage_seeds;
use brokk_bifrost_core::analyzer::symbol_path::parse_symbol_path;
use brokk_bifrost_core::analyzer::usages::model::{ExportEntry, ImportBinder, ImportKind};
use brokk_bifrost_core::analyzer::usages::{ImportEdge, ImportEdgeKind};
use brokk_bifrost_core::analyzer::{
BoundedDefinitionLookup, CodeUnit, CodeUnitIndex, Language, ProjectFile, Range,
};
use std::collections::BTreeSet;
use tree_sitter::Node;
pub fn infer_export_names(python: &dyn PythonUsageSource, target: &CodeUnit) -> BTreeSet<String> {
if target_owner_code_unit(python, target).is_some() {
let owner_name = top_level_identifier(python, target);
let owner_exports =
infer_export_names_for_local(python, target, target.source(), &owner_name);
if !owner_exports.is_empty() {
return owner_exports;
}
}
infer_export_names_for_local(python, target, target.source(), target.identifier())
}
pub fn infer_usage_seeds(
python: &dyn PythonUsageSource,
target: &CodeUnit,
seed_names: BTreeSet<String>,
) -> BTreeSet<(ProjectFile, String)> {
let mut seeds = BTreeSet::new();
for seed_name in &seed_names {
seeds.extend(usage_seeds(python, target.source(), seed_name));
}
if seeds.is_empty()
&& seed_names.contains(target.identifier())
&& is_module_level_target_identifier(python, target, target.source(), target.identifier())
{
seeds.insert((target.source().clone(), target.identifier().to_string()));
}
seeds
}
fn infer_export_names_for_local(
python: &dyn PythonUsageSource,
target: &CodeUnit,
file: &ProjectFile,
local_name: &str,
) -> BTreeSet<String> {
let index = python.export_index_of(file);
let mut export_names = BTreeSet::new();
if index.exports_by_name.contains_key(local_name) {
export_names.insert(local_name.to_string());
}
for (export_name, entry) in &index.exports_by_name {
if matches!(entry, ExportEntry::Local { local_name: name } if name == local_name) {
export_names.insert(export_name.clone());
}
}
if export_names.is_empty()
&& is_module_level_target_identifier(python, target, file, local_name)
{
export_names.insert(local_name.to_string());
}
export_names
}
fn is_module_level_target_identifier(
python: &dyn PythonUsageSource,
target: &CodeUnit,
file: &ProjectFile,
local_name: &str,
) -> bool {
target.source() == file
&& target.identifier() == local_name
&& python
.parent_of(target)
.is_some_and(|parent| parent.is_module() && parent.source() == file)
}
pub fn top_level_identifier(index: &dyn CodeUnitIndex, target: &CodeUnit) -> String {
let mut current = target.clone();
while let Some(parent) = index.parent_of(¤t) {
if parent.is_module() {
break;
}
current = parent;
}
current.identifier().to_string()
}
pub fn member_name(index: &dyn CodeUnitIndex, target: &CodeUnit) -> Option<String> {
target_owner_code_unit(index, target).map(|_| target.identifier().to_string())
}
pub fn target_owner_code_unit(index: &dyn CodeUnitIndex, target: &CodeUnit) -> Option<CodeUnit> {
index
.parent_of(target)
.filter(|parent| parent.source() == target.source() && parent.is_class())
}
pub fn resolve_receiver_type(
graph: &PythonGraphSource<'_>,
python: &dyn PythonUsageSource,
file: &ProjectFile,
raw_type: &str,
target_self_file: bool,
) -> Option<CodeUnit> {
let raw_type = raw_type.trim();
if raw_type.is_empty() || raw_type.contains('.') || raw_type.contains('|') {
return None;
}
if let Some(binding) = python.import_binder_of(file).bindings.get(raw_type)
&& binding.kind == ImportKind::Named
&& let Some(imported) = binding.imported_name.as_ref()
{
let fqn = format!("{}.{}", binding.module_specifier, imported);
if let Some(class) =
resolve_fqn_candidates(python, &fqn, |name| graph.index.definitions(name).collect())
.into_iter()
.find(CodeUnit::is_class)
{
return Some(class);
}
}
if let Some(provider) = graph.imports
&& let Some(imported) = provider
.imported_code_units_of(file)
.iter()
.find(|code_unit| code_unit.identifier() == raw_type && code_unit.is_class())
{
return Some(imported.clone());
}
graph
.index
.declarations(file)
.into_iter()
.find(|code_unit| code_unit.identifier() == raw_type && code_unit.is_class())
.or_else(|| {
if !target_self_file {
return None;
}
let mut resolved = None;
(graph.definitions)(&mut |support| {
resolved = resolve_indexed_receiver_type(graph.index, support, file, raw_type);
});
resolved
})
}
fn resolve_bare_annotation_symbol(
graph: &PythonGraphSource<'_>,
python: &dyn PythonUsageSource,
file: &ProjectFile,
source: &str,
node: Node<'_>,
raw_symbol: &str,
) -> Vec<CodeUnit> {
let raw_symbol = raw_symbol.trim();
if raw_symbol.is_empty() {
return Vec::new();
}
if let Some(owner) = annotation_scope_owner_class(graph, file, source, node) {
let owner_candidates: Vec<_> = exact_owner_annotation_members(graph, &owner, raw_symbol)
.into_iter()
.filter(|candidate| !candidate.is_function())
.collect();
if !owner_candidates.is_empty() {
return owner_candidates;
}
}
let mut candidates = Vec::new();
if let Some(binding) = python.import_binder_of(file).bindings.get(raw_symbol)
&& binding.kind == ImportKind::Named
&& let Some(imported) = binding.imported_name.as_ref()
{
let fqn = format!("{}.{}", binding.module_specifier, imported);
let mut imported_candidates =
resolve_fqn_candidates(python, &fqn, |name| graph.index.definitions(name).collect());
imported_candidates.retain(|candidate| {
!candidate.is_module() || candidate.fq_name() != binding.module_specifier
});
candidates.extend(imported_candidates);
}
candidates.extend(
graph
.index
.top_level_declarations(file)
.into_iter()
.filter(|code_unit| {
!code_unit.is_module()
&& code_unit.identifier() == raw_symbol
&& graph
.index
.parent_of(code_unit)
.is_some_and(|parent| parent.is_module())
}),
);
candidates.retain(|candidate| !candidate.is_function());
candidates.sort();
candidates.dedup();
candidates
}
pub fn annotation_reference_candidates(
graph: &PythonGraphSource<'_>,
python: &dyn PythonUsageSource,
file: &ProjectFile,
source: &str,
node: Node<'_>,
target_self_file: bool,
) -> Option<Vec<CodeUnit>> {
if !is_annotation_reference_node(node) {
return None;
}
let mut candidates = match node.kind() {
"identifier" => {
let mut candidates = resolve_bare_annotation_symbol(
graph,
python,
file,
source,
node,
node_text(node, source),
);
if candidates.is_empty() {
candidates.extend(resolve_receiver_type(
graph,
python,
file,
node_text(node, source),
target_self_file,
));
}
candidates
}
"string_content" => {
let Some(string) = node.parent() else {
return Some(Vec::new());
};
let Some(ranges) = python_deferred_annotation_identifier_ranges(string, source, None)
else {
return Some(Vec::new());
};
let mut candidates = Vec::new();
for range in ranges {
let Some(symbol) = source.get(range.start_byte..range.end_byte) else {
continue;
};
let mut symbol_candidates =
resolve_bare_annotation_symbol(graph, python, file, source, node, symbol);
if symbol_candidates.is_empty() {
symbol_candidates.extend(resolve_receiver_type(
graph,
python,
file,
symbol,
target_self_file,
));
}
candidates.extend(symbol_candidates);
}
candidates
}
"attribute" => resolve_annotation_attribute_types(graph, python, file, source, node),
_ => Vec::new(),
};
candidates.sort();
candidates.dedup();
Some(candidates)
}
#[allow(clippy::too_many_arguments)]
pub fn annotation_reference_candidates_at_focus(
graph: &PythonGraphSource<'_>,
python: &dyn PythonUsageSource,
file: &ProjectFile,
source: &str,
node: Node<'_>,
focus_start: usize,
focus_end: usize,
target_self_file: bool,
) -> Option<Vec<CodeUnit>> {
if !is_annotation_reference_node(node) {
return None;
}
let deferred_tree = if node.kind() == "string_content" {
Some(python_deferred_annotation_tree(
node.parent()?,
source,
None,
)?)
} else {
None
};
let search_root = deferred_tree.as_ref().map_or(node, |tree| tree.root_node());
let focused = search_root.descendant_for_byte_range(focus_start, focus_end)?;
if focused.kind() != "identifier"
|| focused.start_byte() != focus_start
|| focused.end_byte() != focus_end
{
return Some(Vec::new());
}
let mut path = focused;
while let Some(parent) = path.parent() {
if parent.kind() != "attribute" {
break;
}
path = parent;
}
if path.kind() == "attribute" {
return Some(focused_annotation_attribute_candidates(
graph, python, file, source, node, path, focused,
));
}
let symbol = node_text(focused, source);
let mut candidates = resolve_bare_annotation_symbol(graph, python, file, source, node, symbol);
if candidates.is_empty() {
candidates.extend(resolve_receiver_type(
graph,
python,
file,
symbol,
target_self_file,
));
}
candidates.sort();
candidates.dedup();
Some(candidates)
}
fn focused_annotation_attribute_candidates(
graph: &PythonGraphSource<'_>,
python: &dyn PythonUsageSource,
file: &ProjectFile,
source: &str,
scope_node: Node<'_>,
path: Node<'_>,
focused: Node<'_>,
) -> Vec<CodeUnit> {
let Some((root, attributes)) = annotation_attribute_chain(path) else {
return Vec::new();
};
if root.id() == focused.id() {
return resolve_bare_annotation_symbol(
graph,
python,
file,
source,
scope_node,
node_text(root, source),
);
}
if attributes
.last()
.is_some_and(|node| node.id() == focused.id())
{
let candidates = namespace_qualified_declarations(graph, python, file, source, path);
if !candidates.is_empty() {
return candidates;
}
}
let owners: Vec<_> = resolve_bare_annotation_symbol(
graph,
python,
file,
source,
scope_node,
node_text(root, source),
)
.into_iter()
.filter(CodeUnit::is_class)
.collect();
let [owner] = owners.as_slice() else {
return Vec::new();
};
let mut owner = owner.clone();
for attribute in attributes {
let candidates = exact_nested_annotation_class(graph, &owner, node_text(attribute, source));
if attribute.id() == focused.id() {
return candidates;
}
let [next] = candidates.as_slice() else {
return Vec::new();
};
owner = next.clone();
}
Vec::new()
}
pub fn annotation_class_qualifier_site<'tree>(
graph: &PythonGraphSource<'_>,
python: &dyn PythonUsageSource,
file: &ProjectFile,
source: &str,
node: Node<'tree>,
target: &CodeUnit,
) -> Option<Node<'tree>> {
if node.kind() != "attribute" || !target.is_class() || !is_annotation_reference_node(node) {
return None;
}
let (root, attributes) = annotation_attribute_chain(node)?;
let owners: Vec<_> =
resolve_bare_annotation_symbol(graph, python, file, source, root, node_text(root, source))
.into_iter()
.filter(CodeUnit::is_class)
.collect();
let [owner] = owners.as_slice() else {
return None;
};
let mut owner = owner.clone();
if &owner == target {
return Some(root);
}
let qualifier_count = attributes.len().saturating_sub(1);
for attribute in attributes.into_iter().take(qualifier_count) {
let next_candidates =
exact_nested_annotation_class(graph, &owner, node_text(attribute, source));
let [next] = next_candidates.as_slice() else {
return None;
};
owner = next.clone();
if &owner == target {
return Some(attribute);
}
}
None
}
fn resolve_annotation_attribute_types(
graph: &PythonGraphSource<'_>,
python: &dyn PythonUsageSource,
file: &ProjectFile,
source: &str,
node: Node<'_>,
) -> Vec<CodeUnit> {
let mut candidates = namespace_qualified_declarations(graph, python, file, source, node);
let Some((root, attributes)) = annotation_attribute_chain(node) else {
return candidates;
};
let root_text = node_text(root, source);
let owners: Vec<_> =
resolve_bare_annotation_symbol(graph, python, file, source, root, root_text)
.into_iter()
.filter(CodeUnit::is_class)
.collect();
let [owner] = owners.as_slice() else {
return candidates;
};
let mut owner = owner.clone();
for attribute in attributes {
let segment = node_text(attribute, source);
let next_candidates = exact_nested_annotation_class(graph, &owner, segment);
let [next] = next_candidates.as_slice() else {
return candidates;
};
owner = next.clone();
}
candidates.push(owner);
candidates
}
fn annotation_attribute_chain(node: Node<'_>) -> Option<(Node<'_>, Vec<Node<'_>>)> {
let mut attributes = Vec::new();
let mut current = node;
while current.kind() == "attribute" {
attributes.push(current.child_by_field_name("attribute")?);
current = current.child_by_field_name("object")?;
}
if current.kind() != "identifier" || attributes.is_empty() {
return None;
}
attributes.reverse();
Some((current, attributes))
}
fn exact_nested_annotation_class(
graph: &PythonGraphSource<'_>,
owner: &CodeUnit,
segment: &str,
) -> Vec<CodeUnit> {
let mut candidates: Vec<_> = exact_owner_annotation_members(graph, owner, segment)
.into_iter()
.filter(CodeUnit::is_class)
.collect();
candidates.sort();
candidates.dedup();
candidates
}
fn exact_owner_annotation_members(
graph: &PythonGraphSource<'_>,
owner: &CodeUnit,
segment: &str,
) -> Vec<CodeUnit> {
let mut candidates: Vec<_> = graph
.index
.declarations(owner.source())
.into_iter()
.filter(|unit| {
unit.identifier() == segment
&& graph
.index
.parent_of(unit)
.is_some_and(|parent| parent.fq_name() == owner.fq_name())
})
.collect();
candidates.sort();
candidates.dedup();
candidates
}
fn annotation_scope_owner_class(
graph: &PythonGraphSource<'_>,
file: &ProjectFile,
source: &str,
node: Node<'_>,
) -> Option<CodeUnit> {
if !annotation_expression_is_class_scoped(node) {
return None;
}
let range = Range {
start_byte: node.start_byte(),
end_byte: node.end_byte(),
start_line: 0,
end_line: 0,
};
if let Some(enclosing) = graph.index.enclosing_code_unit(file, &range) {
if enclosing.is_class() {
return Some(enclosing);
}
if let Some(owner) = target_owner_code_unit(graph.index, &enclosing) {
return Some(owner);
}
}
structural_annotation_owner_class(graph, file, source, node)
}
fn annotation_expression_is_class_scoped(node: Node<'_>) -> bool {
let site_start = node.start_byte();
let site_end = node.end_byte();
let mut current = node;
while let Some(parent) = current.parent() {
if matches!(parent.kind(), "function_definition" | "lambda")
&& parent
.child_by_field_name("body")
.is_some_and(|body| body.start_byte() <= site_start && site_end <= body.end_byte())
{
return false;
}
if parent.kind() == "class_definition" {
return true;
}
current = parent;
}
false
}
fn structural_annotation_owner_class(
graph: &PythonGraphSource<'_>,
file: &ProjectFile,
source: &str,
node: Node<'_>,
) -> Option<CodeUnit> {
let mut current = node;
while let Some(parent) = current.parent() {
if parent.kind() == "class_definition" {
let name = node_text(parent.child_by_field_name("name")?, source).trim();
if name.is_empty() {
return None;
}
let class_range = Range {
start_byte: parent.start_byte(),
end_byte: parent.end_byte(),
start_line: 0,
end_line: 0,
};
let mut matches: Vec<_> = graph
.index
.declarations(file)
.into_iter()
.filter(|unit| unit.is_class() && unit.identifier() == name)
.filter(|unit| {
graph
.index
.ranges(unit)
.into_iter()
.any(|range| range.contains(&class_range))
})
.collect();
matches.sort();
matches.dedup();
let [owner] = matches.as_slice() else {
return None;
};
return Some(owner.clone());
}
current = parent;
}
None
}
fn is_annotation_reference_node(node: Node<'_>) -> bool {
if !matches!(node.kind(), "identifier" | "attribute" | "string_content") {
return false;
}
python_node_is_in_annotation(node)
}
pub fn resolve_constructor_types(
graph: &PythonGraphSource<'_>,
python: &dyn PythonUsageSource,
file: &ProjectFile,
source: &str,
function: Node<'_>,
) -> Vec<CodeUnit> {
let candidates = match function.kind() {
"identifier" => {
let local = node_text(function, source);
if local.is_empty() {
return Vec::new();
}
let binder = python.import_binder_of(file);
let fqn = match binder.bindings.get(local) {
Some(binding) if binding.kind == ImportKind::Named => binding
.imported_name
.as_ref()
.map(|imported| format!("{}.{}", binding.module_specifier, imported)),
_ => graph
.index
.declarations(file)
.into_iter()
.find(|unit| unit.is_class() && unit.identifier() == local)
.map(|unit| unit.fq_name()),
};
let Some(fqn) = fqn else {
return Vec::new();
};
resolve_fqn_candidates(python, &fqn, |name| graph.index.definitions(name).collect())
}
"attribute" => namespace_qualified_declarations(graph, python, file, source, function),
_ => Vec::new(),
};
let mut classes: Vec<CodeUnit> = candidates.into_iter().filter(CodeUnit::is_class).collect();
classes.sort();
classes.dedup();
classes
}
pub fn resolve_callable_parameter_default_types(
graph: &PythonGraphSource<'_>,
python: &dyn PythonUsageSource,
file: &ProjectFile,
source: &str,
reference: Node<'_>,
local_name: &str,
) -> Vec<CodeUnit> {
let site_start = reference.start_byte();
let site_end = reference.end_byte();
let mut current = reference;
while let Some(parent) = current.parent() {
current = parent;
if !matches!(current.kind(), "function_definition" | "lambda") {
continue;
}
if current
.child_by_field_name("body")
.is_none_or(|body| !(body.start_byte() <= site_start && site_end <= body.end_byte()))
{
continue;
}
let Some(parameters) = current.child_by_field_name("parameters") else {
return Vec::new();
};
let mut cursor = parameters.walk();
for parameter in parameters.named_children(&mut cursor) {
let name = if parameter.kind() == "identifier" {
Some(parameter)
} else {
parameter.child_by_field_name("name")
};
if name.is_none_or(|name| node_text(name, source) != local_name) {
continue;
}
let Some(value) = parameter.child_by_field_name("value") else {
return Vec::new();
};
return resolve_constructor_types(graph, python, file, source, value);
}
return Vec::new();
}
Vec::new()
}
fn namespace_qualified_declarations(
graph: &PythonGraphSource<'_>,
python: &dyn PythonUsageSource,
file: &ProjectFile,
source: &str,
node: Node<'_>,
) -> Vec<CodeUnit> {
let binder = python.import_binder_of(file);
let Some(fqn) = namespace_constructor_fqn(&binder, source, node) else {
return Vec::new();
};
let mut candidates =
resolve_fqn_candidates(python, &fqn, |name| graph.index.definitions(name).collect());
candidates.sort();
candidates.dedup();
candidates
}
fn namespace_constructor_fqn(
binder: &ImportBinder,
source: &str,
function: Node<'_>,
) -> Option<String> {
let mut attributes = Vec::new();
let mut current = function;
while current.kind() == "attribute" {
let attribute = current.child_by_field_name("attribute")?;
let text = node_text(attribute, source);
if text.is_empty() {
return None;
}
attributes.push(text);
current = current.child_by_field_name("object")?;
}
if current.kind() != "identifier" {
return None;
}
let root = node_text(current, source);
let binding = binder.bindings.get(root)?;
if binding.kind != ImportKind::Namespace {
return None;
}
let mut fqn = binding.module_specifier.clone();
for attribute in attributes.into_iter().rev() {
fqn.push('.');
fqn.push_str(attribute);
}
Some(fqn)
}
fn node_text<'a>(node: Node<'_>, source: &'a str) -> &'a str {
brokk_bifrost_core::analyzer::common::node_source_text(node, source)
}
fn resolve_indexed_receiver_type(
index: &dyn CodeUnitIndex,
lookup: &dyn BoundedDefinitionLookup,
file: &ProjectFile,
raw_type: &str,
) -> Option<CodeUnit> {
module_fqn_for_file(index, file)
.into_iter()
.flat_map(|module| lookup.types_in_package(&module, raw_type))
.chain(lookup.fqn(raw_type))
.chain(lookup.by_normalized_fqn(raw_type))
.find(|code_unit| code_unit.identifier() == raw_type && code_unit.is_class())
}
fn module_fqn_for_file(index: &dyn CodeUnitIndex, file: &ProjectFile) -> Option<String> {
index
.declarations(file)
.into_iter()
.find(|code_unit| code_unit.is_module())
.map(|code_unit| code_unit.fq_name())
.or_else(|| {
index
.declarations(file)
.into_iter()
.find(|code_unit| !code_unit.package_name().is_empty())
.map(|code_unit| code_unit.package_name().to_string())
})
}
pub fn normalized_receiver_type(annotation: &str) -> Option<String> {
let annotation = unwrap_python_string_annotation(annotation.trim());
let annotation = unwrap_supported_receiver_wrapper(annotation);
if annotation.is_empty()
|| annotation.contains('|')
|| annotation.contains('[')
|| annotation.contains(']')
|| annotation.contains(',')
|| annotation.contains('(')
|| annotation.contains(')')
|| annotation.contains('{')
|| annotation.contains('}')
|| annotation.contains(':')
{
return None;
}
Some(annotation.to_string())
}
fn unwrap_python_string_annotation(annotation: &str) -> &str {
if annotation.len() >= 2 {
let bytes = annotation.as_bytes();
let first = bytes[0];
let last = bytes[annotation.len() - 1];
if (first == b'\'' || first == b'"') && first == last {
return annotation[1..annotation.len() - 1].trim();
}
}
annotation
}
fn unwrap_supported_receiver_wrapper(annotation: &str) -> &str {
let mut current = annotation.trim();
loop {
let next = current
.strip_prefix("Optional[")
.or_else(|| current.strip_prefix("typing.Optional["))
.and_then(|inner| inner.strip_suffix(']'))
.map(str::trim);
let Some(unwrapped) = next else {
return current;
};
current = unwrapped;
}
}
pub fn receiver_annotation_matches_target(
annotation: &str,
edges: &[ImportEdge],
target_short: &str,
target_self_file: bool,
) -> bool {
let annotation = annotation.trim();
if annotation.is_empty() {
return false;
}
if annotation.contains('|')
|| annotation.contains('[')
|| annotation.contains(']')
|| annotation.contains(',')
|| annotation.contains('(')
|| annotation.contains(')')
{
return false;
}
if annotation == target_short {
return target_self_file || edges.iter().any(|edge| edge.local_name == target_short);
}
let segments = parse_symbol_path(Language::Python, annotation);
let Some((member, qualifier_parts)) = segments.split_last() else {
return false;
};
if qualifier_parts.is_empty() {
return false;
}
let qualifier = qualifier_parts.join(".");
let member = member.as_str();
if member != target_short {
return false;
}
edges.iter().any(|edge| {
matches!(edge.kind, ImportEdgeKind::Namespace)
&& (edge.local_name == qualifier
|| qualifier.ends_with(&format!(".{}", edge.local_name)))
})
}