use crate::indexer::languages::Language;
use tree_sitter::Node;
pub struct Python {}
impl Language for Python {
fn name(&self) -> &'static str {
"python"
}
fn get_ts_language(&self) -> tree_sitter::Language {
tree_sitter_python::LANGUAGE.into()
}
fn get_meaningful_kinds(&self) -> Vec<&'static str> {
vec![
"function_definition",
"class_definition",
"import_statement",
"import_from_statement",
]
}
fn extract_symbols(&self, node: Node, contents: &str) -> Vec<String> {
let mut symbols = Vec::new();
match node.kind() {
"function_definition" => {
for child in node.children(&mut node.walk()) {
if child.kind() == "identifier" {
if let Ok(name) = child.utf8_text(contents.as_bytes()) {
symbols.push(name.to_string());
}
break;
}
}
for child in node.children(&mut node.walk()) {
if child.kind() == "block" {
self.extract_python_variables(child, contents, &mut symbols);
break;
}
}
}
_ => self.extract_identifiers(node, contents, &mut symbols),
}
symbols.sort();
symbols.dedup();
symbols
}
fn extract_identifiers(&self, node: Node, contents: &str, symbols: &mut Vec<String>) {
let kind = node.kind();
if kind == "identifier" {
if let Ok(text) = node.utf8_text(contents.as_bytes()) {
let t = text.trim();
if !t.is_empty() && !symbols.contains(&t.to_string()) && !t.starts_with("_") {
symbols.push(t.to_string());
}
}
}
let mut cursor = node.walk();
if cursor.goto_first_child() {
loop {
self.extract_identifiers(cursor.node(), contents, symbols);
if !cursor.goto_next_sibling() {
break;
}
}
}
}
fn are_node_types_equivalent(&self, type1: &str, type2: &str) -> bool {
if type1 == type2 {
return true;
}
let semantic_groups = [
&["function_definition"] as &[&str],
&["class_definition"],
&["import_statement", "import_from_statement"],
];
for group in &semantic_groups {
let contains_type1 = group.contains(&type1);
let contains_type2 = group.contains(&type2);
if contains_type1 && contains_type2 {
return true;
}
}
false
}
fn get_node_type_description(&self, node_type: &str) -> &'static str {
match node_type {
"function_definition" => "function declarations",
"class_definition" => "class declarations",
"import_statement" | "import_from_statement" => "import statements",
_ => "declarations",
}
}
fn extract_imports_exports(&self, node: Node, contents: &str) -> (Vec<String>, Vec<String>) {
let mut imports = Vec::new();
let mut exports = Vec::new();
match node.kind() {
"import_statement" => {
if let Ok(import_text) = node.utf8_text(contents.as_bytes()) {
if let Some(imported_paths) =
parse_python_import_statement_full_path(import_text)
{
imports.extend(imported_paths);
}
}
}
"import_from_statement" => {
if let Ok(import_text) = node.utf8_text(contents.as_bytes()) {
if let Some(imported_paths) =
parse_python_from_import_statement_full_path(import_text)
{
imports.extend(imported_paths);
}
}
}
"function_definition" | "class_definition" => {
if let Some(parent) = node.parent() {
if parent.kind() == "module" {
for child in node.children(&mut node.walk()) {
if child.kind() == "identifier" {
if let Ok(name) = child.utf8_text(contents.as_bytes()) {
exports.push(name.to_string());
break;
}
}
}
}
}
}
_ => {}
}
(imports, exports)
}
fn resolve_import(
&self,
import_path: &str,
source_file: &str,
all_files: &[String],
) -> Option<String> {
use super::resolution_utils::FileRegistry;
let registry = FileRegistry::new(all_files);
if import_path.starts_with('.') {
self.resolve_relative_import(import_path, source_file, ®istry)
} else {
self.resolve_absolute_import(import_path, source_file, ®istry)
}
}
fn get_file_extensions(&self) -> Vec<&'static str> {
vec!["py"]
}
}
impl Python {
#[allow(clippy::only_used_in_recursion)]
fn extract_python_variables(&self, node: Node, contents: &str, symbols: &mut Vec<String>) {
let mut cursor = node.walk();
if cursor.goto_first_child() {
loop {
let child = cursor.node();
match child.kind() {
"assignment" => {
let mut child_cursor = child.walk();
if child_cursor.goto_first_child() {
let target = child_cursor.node();
if target.kind() == "identifier" {
if let Ok(name) = target.utf8_text(contents.as_bytes()) {
if !name.starts_with("_")
&& !symbols.contains(&name.to_string())
{
symbols.push(name.to_string());
}
}
}
}
}
"expression_statement" => {
for expr_child in child.children(&mut child.walk()) {
if expr_child.kind() == "augmented_assignment" {
let mut aug_cursor = expr_child.walk();
if aug_cursor.goto_first_child() {
let target = aug_cursor.node();
if target.kind() == "identifier" {
if let Ok(name) = target.utf8_text(contents.as_bytes()) {
if !name.starts_with("_")
&& !symbols.contains(&name.to_string())
{
symbols.push(name.to_string());
}
}
}
}
}
}
}
"for_statement" | "while_statement" | "if_statement" | "try_statement"
| "with_statement" => {
for stmt_child in child.children(&mut child.walk()) {
if stmt_child.kind() == "block" {
self.extract_python_variables(stmt_child, contents, symbols);
}
}
}
_ => {}
}
if !cursor.goto_next_sibling() {
break;
}
}
}
}
}
fn parse_python_import_statement_full_path(import_text: &str) -> Option<Vec<String>> {
let mut imports = Vec::new();
let cleaned = import_text.trim();
if let Some(rest) = cleaned.strip_prefix("import ") {
for item in rest.split(',') {
let item = item.trim();
let module_name = if let Some(as_pos) = item.find(" as ") {
&item[..as_pos]
} else {
item
};
if !module_name.is_empty() {
imports.push(module_name.to_string());
}
}
return Some(imports);
}
None
}
fn parse_python_from_import_statement_full_path(import_text: &str) -> Option<Vec<String>> {
let mut imports = Vec::new();
let cleaned = import_text.trim();
if cleaned.starts_with("from ") && cleaned.contains(" import ") {
if let Some(import_pos) = cleaned.find(" import ") {
let module_part = &cleaned[5..import_pos]; if !module_part.is_empty() {
imports.push(module_part.to_string());
}
}
return Some(imports);
}
None
}
impl Python {
fn resolve_relative_import(
&self,
import_path: &str,
source_file: &str,
registry: &super::resolution_utils::FileRegistry,
) -> Option<String> {
let dots = import_path.chars().take_while(|&c| c == '.').count();
let module_name = &import_path[dots..];
let source_path = std::path::Path::new(source_file);
let mut target_dir = source_path.parent()?;
for _ in 1..dots {
target_dir = target_dir.parent()?;
}
self.find_python_module(module_name, target_dir, registry)
}
fn resolve_absolute_import(
&self,
import_path: &str,
source_file: &str,
registry: &super::resolution_utils::FileRegistry,
) -> Option<String> {
let source_path = std::path::Path::new(source_file);
let source_dir = source_path.parent()?;
if let Some(result) = self.find_python_module(import_path, source_dir, registry) {
return Some(result);
}
if let Some(project_root) = super::resolution_utils::find_project_root(source_file) {
let root_path = std::path::Path::new(&project_root);
self.find_python_module(import_path, root_path, registry)
} else {
None
}
}
fn find_python_module(
&self,
module_name: &str,
base_dir: &std::path::Path,
registry: &super::resolution_utils::FileRegistry,
) -> Option<String> {
let module_parts: Vec<&str> = module_name.split('.').collect();
let mut current_path = base_dir.to_path_buf();
for (i, part) in module_parts.iter().enumerate() {
if i == module_parts.len() - 1 {
let file_path = current_path.join(format!("{}.py", part));
if let Some(found) = registry.find_exact_file(&file_path.to_string_lossy()) {
return Some(found);
}
let init_path = current_path.join(part).join("__init__.py");
if let Some(found) = registry.find_exact_file(&init_path.to_string_lossy()) {
return Some(found);
}
} else {
current_path = current_path.join(part);
}
}
None
}
}