use std::collections::HashSet;
use sqry_core::graph::local_scopes::{self, ResolutionOutcome, ScopeId, ScopeKindTrait, ScopeTree};
use sqry_core::graph::unified::build::helper::GraphBuildHelper;
use sqry_core::graph::{GraphResult, Span};
use tree_sitter::Node;
#[derive(Debug, Clone, Copy, Eq, PartialEq)]
pub(crate) enum ScopeKind {
Module,
Function,
Lambda,
Class,
Comprehension,
}
impl ScopeKindTrait for ScopeKind {
fn is_class_scope(&self) -> bool {
*self == ScopeKind::Class
}
fn is_overlap_boundary(&self) -> bool {
*self == ScopeKind::Class
}
fn allows_nested_shadowing(&self) -> bool {
true }
}
pub(crate) type PythonScopeTree = ScopeTree<ScopeKind>;
pub(crate) fn build(root: Node, content: &[u8]) -> GraphResult<PythonScopeTree> {
let content_len = content.len();
let mut tree = PythonScopeTree::new(content_len);
let mut guard = local_scopes::load_recursion_guard();
let module_scope = tree.add_scope(ScopeKind::Module, 0, content_len, None);
build_scopes_recursive(&mut tree, root, content, module_scope, &mut guard)?;
tree.rebuild_index();
bind_declarations_recursive(&mut tree, root, content, &mut guard)?;
tree.rebuild_index();
Ok(tree)
}
#[allow(clippy::only_used_in_recursion)]
fn build_scopes_recursive(
tree: &mut PythonScopeTree,
node: Node,
content: &[u8],
current_scope: Option<ScopeId>,
guard: &mut sqry_core::query::security::RecursionGuard,
) -> GraphResult<()> {
guard
.enter()
.map_err(|e| local_scopes::recursion_error_to_graph_error(&e, node))?;
let new_scope = match node.kind() {
"function_definition" => {
tree.add_scope(
ScopeKind::Function,
node.start_byte(),
node.end_byte(),
current_scope,
)
}
"lambda" => {
tree.add_scope(
ScopeKind::Lambda,
node.start_byte(),
node.end_byte(),
current_scope,
)
}
"class_definition" => {
tree.add_scope(
ScopeKind::Class,
node.start_byte(),
node.end_byte(),
current_scope,
)
}
"list_comprehension"
| "set_comprehension"
| "dictionary_comprehension"
| "generator_expression" => tree.add_scope(
ScopeKind::Comprehension,
node.start_byte(),
node.end_byte(),
current_scope,
),
_ => None,
};
let scope_for_children = new_scope.or(current_scope);
let mut cursor = node.walk();
for child in node.children(&mut cursor) {
build_scopes_recursive(tree, child, content, scope_for_children, guard)?;
}
guard.exit();
Ok(())
}
fn bind_declarations_recursive(
tree: &mut PythonScopeTree,
node: Node,
content: &[u8],
guard: &mut sqry_core::query::security::RecursionGuard,
) -> GraphResult<()> {
guard
.enter()
.map_err(|e| local_scopes::recursion_error_to_graph_error(&e, node))?;
match node.kind() {
"function_definition" => {
bind_parameters(tree, node, content);
}
"lambda" => {
bind_lambda_parameters(tree, node, content);
}
"for_statement" => {
bind_for_variable(tree, node, content);
}
"expression_statement" => {
let mut cursor = node.walk();
for child in node.children(&mut cursor) {
if child.kind() == "assignment" {
bind_assignment(tree, child, content);
}
}
}
"assignment" => {
bind_assignment(tree, node, content);
}
"named_expression" => {
bind_walrus(tree, node, content);
}
"except_clause" => {
bind_except_variable(tree, node, content);
}
"with_statement" => {
bind_with_variable(tree, node, content);
}
"for_in_clause" => {
bind_comprehension_variable(tree, node, content);
}
_ => {}
}
let mut cursor = node.walk();
for child in node.children(&mut cursor) {
bind_declarations_recursive(tree, child, content, guard)?;
}
guard.exit();
Ok(())
}
fn bind_parameters(tree: &mut PythonScopeTree, func_node: Node, content: &[u8]) {
let Some(params_node) = func_node.child_by_field_name("parameters") else {
return;
};
let Some(scope_id) = tree.innermost_scope_at(func_node.start_byte()) else {
return;
};
for param in params_node.children(&mut params_node.walk()) {
match param.kind() {
"identifier" => {
bind_identifier_param(tree, scope_id, param, content);
}
"typed_parameter" => {
if let Some(name_node) = param.child_by_field_name("name") {
bind_identifier_param(tree, scope_id, name_node, content);
} else {
let mut cursor = param.walk();
for child in param.children(&mut cursor) {
if child.kind() == "identifier" {
bind_identifier_param(tree, scope_id, child, content);
break;
}
if child.kind() == "list_splat_pattern"
|| child.kind() == "dictionary_splat_pattern"
{
if let Some(id) = local_scopes::first_child_of_kind(child, "identifier")
{
bind_identifier_param(tree, scope_id, id, content);
}
break;
}
}
}
}
"default_parameter" => {
if let Some(name_node) = param.child_by_field_name("name") {
bind_identifier_param(tree, scope_id, name_node, content);
}
}
"typed_default_parameter" => {
if let Some(name_node) = param.child_by_field_name("name") {
bind_identifier_param(tree, scope_id, name_node, content);
}
}
"list_splat_pattern" => {
if let Some(id) = local_scopes::first_child_of_kind(param, "identifier") {
bind_identifier_param(tree, scope_id, id, content);
}
}
"dictionary_splat_pattern" => {
if let Some(id) = local_scopes::first_child_of_kind(param, "identifier") {
bind_identifier_param(tree, scope_id, id, content);
}
}
_ => {}
}
}
}
fn bind_lambda_parameters(tree: &mut PythonScopeTree, lambda_node: Node, content: &[u8]) {
let Some(params_node) = lambda_node.child_by_field_name("parameters") else {
return;
};
let Some(scope_id) = tree.innermost_scope_at(lambda_node.start_byte()) else {
return;
};
for param in params_node.children(&mut params_node.walk()) {
match param.kind() {
"identifier" => {
bind_identifier_param(tree, scope_id, param, content);
}
"default_parameter" => {
if let Some(name_node) = param.child_by_field_name("name") {
bind_identifier_param(tree, scope_id, name_node, content);
}
}
"list_splat_pattern" | "dictionary_splat_pattern" => {
if let Some(id) = local_scopes::first_child_of_kind(param, "identifier") {
bind_identifier_param(tree, scope_id, id, content);
}
}
_ => {}
}
}
}
fn bind_identifier_param(
tree: &mut PythonScopeTree,
scope_id: ScopeId,
name_node: Node,
content: &[u8],
) {
let Ok(name) = name_node.utf8_text(content) else {
return;
};
let name = name.trim();
if name.is_empty() || name == "self" || name == "cls" {
return;
}
tree.add_binding(
scope_id,
name,
name_node.start_byte(),
name_node.end_byte(),
name_node.end_byte(), None,
);
}
fn bind_assignment(tree: &mut PythonScopeTree, assignment_node: Node, content: &[u8]) {
let Some(left) = assignment_node.child_by_field_name("left") else {
return;
};
let Some(scope_id) = find_binding_scope(tree, assignment_node.start_byte()) else {
return;
};
let init_start = assignment_node
.child_by_field_name("right")
.map(|r| r.start_byte());
bind_pattern(tree, scope_id, left, content, init_start, assignment_node);
}
fn bind_for_variable(tree: &mut PythonScopeTree, for_node: Node, content: &[u8]) {
let Some(left) = for_node.child_by_field_name("left") else {
return;
};
let Some(scope_id) = find_binding_scope(tree, for_node.start_byte()) else {
return;
};
bind_pattern(tree, scope_id, left, content, None, for_node);
}
fn bind_walrus(tree: &mut PythonScopeTree, walrus_node: Node, content: &[u8]) {
let Some(name_node) = walrus_node.child_by_field_name("name") else {
return;
};
let Some(scope_id) = find_function_scope(tree, walrus_node.start_byte()) else {
return;
};
let Ok(name) = name_node.utf8_text(content) else {
return;
};
let name = name.trim();
if name.is_empty() {
return;
}
let init_start = walrus_node
.child_by_field_name("value")
.map(|v| v.start_byte());
tree.add_binding(
scope_id,
name,
name_node.start_byte(),
name_node.end_byte(),
walrus_node.end_byte(),
init_start,
);
}
fn bind_except_variable(tree: &mut PythonScopeTree, except_node: Node, content: &[u8]) {
let mut cursor = except_node.walk();
for child in except_node.children(&mut cursor) {
if child.kind() == "as_pattern" {
let mut inner_cursor = child.walk();
for inner_child in child.children(&mut inner_cursor) {
if inner_child.kind() == "as_pattern_target" {
if let Some(id) = local_scopes::first_child_of_kind(inner_child, "identifier") {
let Some(scope_id) = find_binding_scope(tree, except_node.start_byte())
else {
return;
};
let Ok(name) = id.utf8_text(content) else {
return;
};
let name = name.trim();
if !name.is_empty() {
tree.add_binding(
scope_id,
name,
id.start_byte(),
id.end_byte(),
id.end_byte(),
None,
);
}
}
return;
}
}
}
}
}
fn bind_with_variable(tree: &mut PythonScopeTree, with_node: Node, content: &[u8]) {
let mut cursor = with_node.walk();
for child in with_node.children(&mut cursor) {
if child.kind() == "with_clause" {
let mut clause_cursor = child.walk();
for item in child.children(&mut clause_cursor) {
if item.kind() == "with_item" {
bind_with_item(tree, item, content, with_node.start_byte());
}
}
}
}
}
fn bind_with_item(tree: &mut PythonScopeTree, item: Node, content: &[u8], context_byte: usize) {
let mut cursor = item.walk();
for child in item.children(&mut cursor) {
if child.kind() == "as_pattern" {
let mut inner_cursor = child.walk();
for inner_child in child.children(&mut inner_cursor) {
if inner_child.kind() == "as_pattern_target" {
let Some(scope_id) = find_binding_scope(tree, context_byte) else {
return;
};
if let Some(id) = local_scopes::first_child_of_kind(inner_child, "identifier") {
let Ok(name) = id.utf8_text(content) else {
return;
};
let name = name.trim();
if !name.is_empty() {
tree.add_binding(
scope_id,
name,
id.start_byte(),
id.end_byte(),
id.end_byte(),
None,
);
}
} else {
bind_pattern(tree, scope_id, inner_child, content, None, inner_child);
}
return;
}
}
}
}
}
fn bind_comprehension_variable(tree: &mut PythonScopeTree, for_in_node: Node, content: &[u8]) {
let Some(left) = for_in_node.child_by_field_name("left") else {
return;
};
let Some(scope_id) = tree.innermost_scope_at(for_in_node.start_byte()) else {
return;
};
bind_pattern(tree, scope_id, left, content, None, for_in_node);
}
fn bind_pattern(
tree: &mut PythonScopeTree,
scope_id: ScopeId,
pattern: Node,
content: &[u8],
init_start: Option<usize>,
declarator_node: Node,
) {
match pattern.kind() {
"identifier" => {
let Ok(name) = pattern.utf8_text(content) else {
return;
};
let name = name.trim();
if name.is_empty() || name == "_" {
return;
}
tree.add_binding(
scope_id,
name,
pattern.start_byte(),
pattern.end_byte(),
declarator_node.end_byte(),
init_start,
);
}
"pattern_list" | "tuple_pattern" | "list_pattern" => {
let mut cursor = pattern.walk();
for child in pattern.children(&mut cursor) {
bind_pattern(tree, scope_id, child, content, init_start, declarator_node);
}
}
"list_splat_pattern" => {
if let Some(id) = local_scopes::first_child_of_kind(pattern, "identifier") {
let Ok(name) = id.utf8_text(content) else {
return;
};
let name = name.trim();
if !name.is_empty() && name != "_" {
tree.add_binding(
scope_id,
name,
id.start_byte(),
id.end_byte(),
declarator_node.end_byte(),
init_start,
);
}
}
}
_ => {
}
}
}
fn find_binding_scope(tree: &PythonScopeTree, byte: usize) -> Option<ScopeId> {
let innermost = tree.innermost_scope_at(byte)?;
let chain = tree.scope_chain(innermost);
for scope_id in &chain {
let kind = tree.scopes[*scope_id].kind;
match kind {
ScopeKind::Function | ScopeKind::Lambda | ScopeKind::Module => {
return Some(*scope_id);
}
ScopeKind::Class | ScopeKind::Comprehension => {
}
}
}
None
}
fn find_function_scope(tree: &PythonScopeTree, byte: usize) -> Option<ScopeId> {
let innermost = tree.innermost_scope_at(byte)?;
let chain = tree.scope_chain(innermost);
for scope_id in &chain {
let kind = tree.scopes[*scope_id].kind;
match kind {
ScopeKind::Function | ScopeKind::Lambda | ScopeKind::Module => {
return Some(*scope_id);
}
ScopeKind::Class | ScopeKind::Comprehension => {}
}
}
None
}
pub(crate) fn handle_identifier_for_reference(
node: Node,
content: &[u8],
scope_tree: &mut PythonScopeTree,
helper: &mut GraphBuildHelper,
) {
let Ok(name) = node.utf8_text(content) else {
return;
};
let name = name.trim();
if name.is_empty() || name == "_" || name == "self" || name == "cls" {
return;
}
if is_declaration_context(node) {
return;
}
if is_decorator_or_import_context(node) {
return;
}
if is_call_context(node) {
return;
}
if is_attribute_access(node) {
return;
}
if is_type_annotation_context(node) {
return;
}
if !is_inside_scope(node) {
return;
}
let usage_byte = node.start_byte();
match scope_tree.resolve_identifier(usage_byte, name) {
ResolutionOutcome::Local(binding) => {
let span = Span::from_bytes(node.start_byte(), node.end_byte());
let qualified_name = format!("{}@{}", name, binding.decl_start_byte);
let var_id = helper.add_variable(&qualified_name, Some(span));
if let Some(decl_id) = binding.node_id {
helper.add_reference_edge(var_id, decl_id);
} else {
let decl_span = Span::from_bytes(binding.decl_start_byte, binding.decl_end_byte);
let decl_id = helper.add_variable(&qualified_name, Some(decl_span));
scope_tree.attach_node_id(name, binding.decl_start_byte, decl_id);
helper.add_reference_edge(var_id, decl_id);
}
}
ResolutionOutcome::Member { .. }
| ResolutionOutcome::Ambiguous
| ResolutionOutcome::NoMatch => {}
}
}
#[allow(clippy::match_same_arms)] fn is_declaration_context(node: Node) -> bool {
let Some(parent) = node.parent() else {
return false;
};
match parent.kind() {
"function_definition" | "class_definition" => parent
.child_by_field_name("name")
.is_some_and(|n| n.id() == node.id()),
"assignment" | "augmented_assignment" => parent
.child_by_field_name("left")
.is_some_and(|left| contains_node(left, node)),
"for_statement" => parent
.child_by_field_name("left")
.is_some_and(|left| contains_node(left, node)),
"named_expression" => parent
.child_by_field_name("name")
.is_some_and(|n| n.id() == node.id()),
"typed_parameter" | "typed_default_parameter" | "default_parameter" => parent
.child_by_field_name("name")
.is_some_and(|n| n.id() == node.id()),
#[allow(clippy::match_same_arms)]
"parameters" | "lambda_parameters" => true,
"pattern_list" | "tuple_pattern" | "list_pattern" => {
if let Some(grandparent) = parent.parent() {
match grandparent.kind() {
"assignment" | "augmented_assignment" => grandparent
.child_by_field_name("left")
.is_some_and(|left| contains_node(left, node)),
"for_statement" => grandparent
.child_by_field_name("left")
.is_some_and(|left| contains_node(left, node)),
_ => false,
}
} else {
false
}
}
"list_splat_pattern" | "dictionary_splat_pattern" => true,
"for_in_clause" => parent
.child_by_field_name("left")
.is_some_and(|left| contains_node(left, node)),
"global_statement" | "nonlocal_statement" => true,
"as_pattern_target" => true,
_ => false,
}
}
fn is_decorator_or_import_context(node: Node) -> bool {
let mut current = node.parent();
while let Some(parent) = current {
match parent.kind() {
"decorator" | "import_statement" | "import_from_statement" | "aliased_import" => {
return true;
}
"function_definition" | "class_definition" | "module" => {
return false;
}
_ => current = parent.parent(),
}
}
false
}
fn is_call_context(node: Node) -> bool {
let Some(parent) = node.parent() else {
return false;
};
if parent.kind() == "call" {
return parent
.child_by_field_name("function")
.is_some_and(|f| f.id() == node.id());
}
false
}
fn is_attribute_access(node: Node) -> bool {
let Some(parent) = node.parent() else {
return false;
};
if parent.kind() == "attribute" {
return parent
.child_by_field_name("attribute")
.is_some_and(|a| a.id() == node.id());
}
false
}
fn is_type_annotation_context(node: Node) -> bool {
let mut current = node.parent();
while let Some(parent) = current {
match parent.kind() {
"type" => return true,
"assignment" | "typed_parameter" | "typed_default_parameter" => {
if parent
.child_by_field_name("type")
.is_some_and(|t| contains_node(t, node))
{
return true;
}
return false;
}
"function_definition" => {
if parent
.child_by_field_name("return_type")
.is_some_and(|r| contains_node(r, node))
{
return true;
}
return false;
}
_ => current = parent.parent(),
}
}
false
}
fn is_inside_scope(node: Node) -> bool {
let mut current = node.parent();
while let Some(parent) = current {
match parent.kind() {
"function_definition" | "class_definition" | "lambda" | "module" => return true,
_ => current = parent.parent(),
}
}
false
}
fn contains_node(haystack: Node, needle: Node) -> bool {
if haystack.id() == needle.id() {
return true;
}
let mut cursor = haystack.walk();
for child in haystack.children(&mut cursor) {
if contains_node(child, needle) {
return true;
}
}
false
}
#[allow(dead_code)] fn collect_global_nonlocal_names(func_body: Node, content: &[u8]) -> HashSet<String> {
let mut names = HashSet::new();
collect_global_nonlocal_recursive(func_body, content, &mut names);
names
}
fn collect_global_nonlocal_recursive(node: Node, content: &[u8], names: &mut HashSet<String>) {
match node.kind() {
"global_statement" | "nonlocal_statement" => {
let mut cursor = node.walk();
for child in node.children(&mut cursor) {
if child.kind() == "identifier"
&& let Ok(name) = child.utf8_text(content)
{
names.insert(name.to_string());
}
}
}
"function_definition" | "class_definition" | "lambda" => {}
_ => {
let mut cursor = node.walk();
for child in node.children(&mut cursor) {
collect_global_nonlocal_recursive(child, content, names);
}
}
}
}