use std::path::{Path, PathBuf};
use tree_sitter::Language;
use astmap_core::model::SymbolKind;
use super::{
find_enclosing_symbol, first_line, node_text, parse_with_tree_sitter, symbol_from_node,
ExtractedCall, ExtractedImport, ExtractedSymbol, ExtractedTypeRef, LanguageParser,
LanguageResolver, LanguageSupport, ParseResult,
};
pub(super) fn lang() -> LanguageSupport {
LanguageSupport {
name: "python",
extensions: &["py"],
parser: &PythonParser,
resolver_factory: |_root| Box::new(PythonResolver),
config_files: &["pyproject.toml"],
sibling_fn: python_sibling_expansion,
}
}
fn python_sibling_expansion(rel_path: &str) -> Option<astmap_core::SiblingExpansion> {
let dir = rel_path
.strip_suffix("__init__.py")
.filter(|p| !p.is_empty())?;
Some(astmap_core::SiblingExpansion {
prefix: dir.to_string(),
extension: "py".to_string(),
root_only: false,
})
}
pub(crate) struct PythonParser;
impl LanguageParser for PythonParser {
fn parse(&self, source: &str, _file_path: &Path) -> ParseResult {
parse_python_file(source)
}
fn language_name(&self) -> &str {
"python"
}
}
fn python_language() -> Language {
tree_sitter_python::LANGUAGE.into()
}
fn parse_python_file(source: &str) -> ParseResult {
parse_with_tree_sitter(
source,
&python_language(),
|root, src, symbols, imports| extract_from_node(root, src, symbols, imports, None),
extract_calls,
extract_type_refs,
)
}
fn extract_from_node(
node: tree_sitter::Node,
source: &str,
symbols: &mut Vec<ExtractedSymbol>,
imports: &mut Vec<ExtractedImport>,
parent_index: Option<usize>,
) {
match node.kind() {
"function_definition" => {
if let Some(sym) = extract_function(&node, source, parent_index) {
symbols.push(sym);
return; }
}
"class_definition" => {
if let Some(sym) = extract_class(&node, source, parent_index) {
let idx = symbols.len();
symbols.push(sym);
if let Some(body) = node.child_by_field_name("body") {
for i in 0..body.child_count() {
if let Some(child) = body.child(i as u32) {
extract_from_node(child, source, symbols, imports, Some(idx));
}
}
}
return;
}
}
"decorated_definition" => {
extract_decorated(&node, source, symbols, imports, parent_index);
return;
}
"expression_statement" if parent_index.is_none() => {
for i in 0..node.child_count() {
if let Some(child) = node.child(i as u32) {
if child.kind() == "assignment" {
if let Some(sym) = extract_module_constant(&child, &node, source) {
symbols.push(sym);
}
}
}
}
}
"import_statement" => {
imports.extend(extract_import_plain(&node, source));
}
"import_from_statement" => {
if let Some(imp) = extract_import_from(&node, source) {
imports.push(imp);
}
}
_ => {}
}
for i in 0..node.child_count() {
if let Some(child) = node.child(i as u32) {
extract_from_node(child, source, symbols, imports, parent_index);
}
}
}
fn extract_function(
node: &tree_sitter::Node,
source: &str,
parent_index: Option<usize>,
) -> Option<ExtractedSymbol> {
let name_node = node.child_by_field_name("name")?;
let name = node_text(&name_node, source);
let kind = if parent_index.is_some() {
SymbolKind::Method
} else {
SymbolKind::Function
};
Some(symbol_from_node(name, kind, node, source, parent_index))
}
fn extract_class(
node: &tree_sitter::Node,
source: &str,
parent_index: Option<usize>,
) -> Option<ExtractedSymbol> {
let name_node = node.child_by_field_name("name")?;
let name = node_text(&name_node, source);
Some(symbol_from_node(
name,
SymbolKind::Class,
node,
source,
parent_index,
))
}
fn extract_decorated(
node: &tree_sitter::Node,
source: &str,
symbols: &mut Vec<ExtractedSymbol>,
imports: &mut Vec<ExtractedImport>,
parent_index: Option<usize>,
) {
let decorators: Vec<String> = (0..node.child_count())
.filter_map(|i| node.child(i as u32))
.filter(|c| c.kind() == "decorator")
.map(|c| node_text(&c, source).trim().to_string())
.collect();
let def_node = match node.child_by_field_name("definition") {
Some(d) => d,
None => return,
};
let decorator_prefix = if decorators.is_empty() {
String::new()
} else {
format!("{}\n", decorators.join("\n"))
};
match def_node.kind() {
"function_definition" => {
if let Some(mut sym) = extract_function(&def_node, source, parent_index) {
sym.signature = Some(format!(
"{}{}",
decorator_prefix,
first_line(&def_node, source)
));
sym.start_line = node.start_position().row + 1;
sym.start_byte = node.start_byte();
symbols.push(sym);
}
}
"class_definition" => {
if let Some(mut sym) = extract_class(&def_node, source, parent_index) {
sym.signature = Some(format!(
"{}{}",
decorator_prefix,
first_line(&def_node, source)
));
sym.start_line = node.start_position().row + 1;
sym.start_byte = node.start_byte();
let idx = symbols.len();
symbols.push(sym);
if let Some(body) = def_node.child_by_field_name("body") {
for i in 0..body.child_count() {
if let Some(child) = body.child(i as u32) {
extract_from_node(child, source, symbols, imports, Some(idx));
}
}
}
}
}
_ => {}
}
}
fn extract_module_constant(
assignment: &tree_sitter::Node,
stmt: &tree_sitter::Node,
source: &str,
) -> Option<ExtractedSymbol> {
let left = assignment.child_by_field_name("left")?;
if left.kind() != "identifier" {
return None;
}
let name = node_text(&left, source);
let is_all_caps = name.chars().any(|c| c.is_ascii_uppercase())
&& name.chars().all(|c| c.is_ascii_uppercase() || c == '_');
let has_type_annotation = assignment.child_by_field_name("type").is_some();
if !is_all_caps && !has_type_annotation {
return None;
}
Some(symbol_from_node(
name,
SymbolKind::Variable,
stmt,
source,
None,
))
}
fn extract_import_plain(node: &tree_sitter::Node, source: &str) -> Vec<ExtractedImport> {
let mut imports = Vec::new();
for i in 0..node.child_count() {
let child = match node.child(i as u32) {
Some(c) => c,
None => continue,
};
match child.kind() {
"dotted_name" => {
let full_path = node_text(&child, source);
let last_segment = full_path
.rsplit('.')
.next()
.unwrap_or(&full_path)
.to_string();
imports.push(ExtractedImport {
path: full_path,
imported_symbols: vec![last_segment],
});
}
"aliased_import" => {
let name_node = child.child_by_field_name("name");
let alias_node = child.child_by_field_name("alias");
if let Some(name) = name_node {
let path = node_text(&name, source);
let symbol = alias_node
.map(|a| node_text(&a, source))
.unwrap_or_else(|| path.rsplit('.').next().unwrap_or(&path).to_string());
imports.push(ExtractedImport {
path,
imported_symbols: vec![symbol],
});
}
}
_ => {}
}
}
imports
}
fn extract_import_from(node: &tree_sitter::Node, source: &str) -> Option<ExtractedImport> {
let module_node = node.child_by_field_name("module_name")?;
let path = node_text(&module_node, source);
let has_wildcard = (0..node.child_count())
.filter_map(|i| node.child(i as u32))
.any(|c| c.kind() == "wildcard_import");
if has_wildcard {
return Some(ExtractedImport {
path,
imported_symbols: vec![],
});
}
let module_node_id = module_node.id();
let mut imported_symbols = Vec::new();
for i in 0..node.child_count() {
let child = match node.child(i as u32) {
Some(c) => c,
None => continue,
};
match child.kind() {
"dotted_name" if child.id() != module_node_id => {
imported_symbols.push(node_text(&child, source));
}
"aliased_import" => {
if let Some(alias) = child.child_by_field_name("alias") {
imported_symbols.push(node_text(&alias, source));
} else if let Some(name) = child.child_by_field_name("name") {
let text = node_text(&name, source);
let last = text.rsplit('.').next().unwrap_or(&text).to_string();
imported_symbols.push(last);
}
}
_ => {}
}
}
Some(ExtractedImport {
path,
imported_symbols,
})
}
fn extract_calls(
node: tree_sitter::Node,
source: &str,
symbols: &[ExtractedSymbol],
) -> Vec<ExtractedCall> {
let mut calls = Vec::new();
collect_calls(node, source, symbols, &mut calls);
calls
}
fn collect_calls(
node: tree_sitter::Node,
source: &str,
symbols: &[ExtractedSymbol],
calls: &mut Vec<ExtractedCall>,
) {
if node.kind() == "call" {
if let Some(func_node) = node.child_by_field_name("function") {
let callee_name = match func_node.kind() {
"attribute" => func_node
.child_by_field_name("attribute")
.map(|a| node_text(&a, source))
.unwrap_or_else(|| node_text(&func_node, source)),
"identifier" => node_text(&func_node, source),
_ => node_text(&func_node, source),
};
if !is_python_builtin(&callee_name) {
let call_line = node.start_position().row + 1;
if let Some(caller) = find_enclosing_symbol(symbols, call_line) {
calls.push(ExtractedCall {
caller_name: caller.to_string(),
callee_name,
line: call_line,
});
}
}
}
}
for i in 0..node.child_count() {
if let Some(child) = node.child(i as u32) {
collect_calls(child, source, symbols, calls);
}
}
}
fn is_python_builtin(name: &str) -> bool {
matches!(
name,
"print"
| "len"
| "range"
| "enumerate"
| "zip"
| "map"
| "filter"
| "sorted"
| "reversed"
| "list"
| "dict"
| "set"
| "tuple"
| "str"
| "int"
| "float"
| "bool"
| "bytes"
| "bytearray"
| "type"
| "isinstance"
| "issubclass"
| "hasattr"
| "getattr"
| "setattr"
| "delattr"
| "callable"
| "super"
| "property"
| "staticmethod"
| "classmethod"
| "abs"
| "all"
| "any"
| "ascii"
| "bin"
| "chr"
| "complex"
| "dir"
| "divmod"
| "format"
| "globals"
| "hash"
| "hex"
| "id"
| "input"
| "iter"
| "max"
| "min"
| "next"
| "oct"
| "open"
| "ord"
| "pow"
| "repr"
| "round"
| "slice"
| "sum"
| "vars"
| "object"
| "Exception"
| "ValueError"
| "TypeError"
| "KeyError"
| "IndexError"
| "AttributeError"
| "RuntimeError"
| "StopIteration"
| "NotImplementedError"
| "OSError"
| "IOError"
)
}
fn extract_type_refs(
node: tree_sitter::Node,
source: &str,
symbols: &[ExtractedSymbol],
) -> Vec<ExtractedTypeRef> {
let mut refs = Vec::new();
collect_type_refs(node, source, symbols, &mut refs);
refs
}
fn collect_type_refs(
node: tree_sitter::Node,
source: &str,
symbols: &[ExtractedSymbol],
refs: &mut Vec<ExtractedTypeRef>,
) {
if node.kind() == "type" {
collect_type_identifiers(node, source, symbols, refs);
return; }
for i in 0..node.child_count() {
if let Some(child) = node.child(i as u32) {
collect_type_refs(child, source, symbols, refs);
}
}
}
fn collect_type_identifiers(
node: tree_sitter::Node,
source: &str,
symbols: &[ExtractedSymbol],
refs: &mut Vec<ExtractedTypeRef>,
) {
if node.kind() == "identifier" {
let name = node_text(&node, source);
if !is_python_builtin_type(&name) {
let line = node.start_position().row + 1;
if let Some(from_sym) = find_enclosing_symbol(symbols, line) {
refs.push(ExtractedTypeRef {
from_symbol: from_sym.to_string(),
to_type: name,
line,
});
}
}
}
for i in 0..node.child_count() {
if let Some(child) = node.child(i as u32) {
collect_type_identifiers(child, source, symbols, refs);
}
}
}
fn is_python_builtin_type(name: &str) -> bool {
matches!(
name,
"int"
| "float"
| "str"
| "bool"
| "bytes"
| "bytearray"
| "None"
| "list"
| "dict"
| "set"
| "tuple"
| "frozenset"
| "object"
| "type"
| "Any"
| "Union"
| "Optional"
| "List"
| "Dict"
| "Set"
| "Tuple"
| "FrozenSet"
| "Type"
| "Callable"
| "Iterator"
| "Generator"
| "Coroutine"
| "Awaitable"
| "Sequence"
| "Mapping"
| "MutableMapping"
| "Iterable"
| "ClassVar"
| "Final"
| "Literal"
| "Protocol"
| "Self"
| "Never"
| "NoReturn"
| "TypeVar"
)
}
pub(crate) struct PythonResolver;
impl LanguageResolver for PythonResolver {
fn resolve_import(
&self,
import_path: &str,
source_file: &str,
project_root: &Path,
) -> Option<PathBuf> {
resolve_python_import(import_path, source_file, project_root)
}
}
fn resolve_python_import(
import_path: &str,
source_file: &str,
project_root: &Path,
) -> Option<PathBuf> {
if import_path.is_empty() {
return None;
}
if import_path.starts_with('.') {
resolve_relative_import(import_path, source_file, project_root)
} else {
resolve_absolute_import(import_path, project_root)
}
}
fn resolve_absolute_import(import_path: &str, project_root: &Path) -> Option<PathBuf> {
let parts: Vec<&str> = import_path.split('.').collect();
if let Some(resolved) = resolve_module_path(&parts, project_root) {
return Some(resolved);
}
for src_dir in &["src", "lib"] {
let base = project_root.join(src_dir);
if base.is_dir() {
if let Some(resolved) = resolve_module_path(&parts, &base) {
return Some(resolved);
}
}
}
None
}
fn resolve_relative_import(
import_path: &str,
source_file: &str,
project_root: &Path,
) -> Option<PathBuf> {
let dot_count = import_path.chars().take_while(|&c| c == '.').count();
let remainder = &import_path[dot_count..];
let source_abs = project_root.join(source_file);
let source_dir = source_abs.parent()?;
let mut base = source_dir.to_path_buf();
for _ in 1..dot_count {
base = base.parent()?.to_path_buf();
}
if remainder.is_empty() {
let init = base.join("__init__.py");
if init.exists() {
return init.canonicalize().ok();
}
return None;
}
let parts: Vec<&str> = remainder.split('.').collect();
resolve_module_path(&parts, &base)
}
fn resolve_module_path(parts: &[&str], base: &Path) -> Option<PathBuf> {
if parts.is_empty() {
let init = base.join("__init__.py");
if init.exists() {
return init.canonicalize().ok();
}
return None;
}
let mut current = base.to_path_buf();
for (i, part) in parts.iter().enumerate() {
if i == parts.len() - 1 {
let as_file = current.join(format!("{}.py", part));
if as_file.exists() {
return as_file.canonicalize().ok();
}
let as_package = current.join(part).join("__init__.py");
if as_package.exists() {
return as_package.canonicalize().ok();
}
return None;
} else {
let next = current.join(part);
if next.is_dir() {
current = next;
} else {
return None;
}
}
}
None
}
#[cfg(test)]
mod tests {
use super::*;
fn parse(source: &str) -> ParseResult {
PythonParser.parse(source, &PathBuf::from("test.py"))
}
#[test]
fn test_parse_function() {
let result = parse("def hello():\n pass\n");
assert!(!result.has_errors);
assert_eq!(result.symbols.len(), 1);
assert_eq!(result.symbols[0].name, "hello");
assert_eq!(result.symbols[0].kind, SymbolKind::Function);
assert!(result.symbols[0].parent_index.is_none());
}
#[test]
fn test_parse_async_function() {
let result = parse("async def fetch():\n pass\n");
assert!(!result.has_errors);
assert_eq!(result.symbols.len(), 1);
assert_eq!(result.symbols[0].name, "fetch");
assert_eq!(result.symbols[0].kind, SymbolKind::Function);
assert!(result.symbols[0]
.signature
.as_ref()
.unwrap()
.contains("async"));
}
#[test]
fn test_parse_class_with_methods() {
let result = parse(
"class Foo:\n def bar(self):\n pass\n def baz(self):\n pass\n",
);
assert!(!result.has_errors);
assert_eq!(result.symbols.len(), 3);
assert_eq!(result.symbols[0].name, "Foo");
assert_eq!(result.symbols[0].kind, SymbolKind::Class);
assert_eq!(result.symbols[1].name, "bar");
assert_eq!(result.symbols[1].kind, SymbolKind::Method);
assert_eq!(result.symbols[1].parent_index, Some(0));
assert_eq!(result.symbols[2].name, "baz");
assert_eq!(result.symbols[2].kind, SymbolKind::Method);
assert_eq!(result.symbols[2].parent_index, Some(0));
}
#[test]
fn test_parse_class_basic() {
let result = parse("class Empty:\n pass\n");
assert!(!result.has_errors);
assert_eq!(result.symbols.len(), 1);
assert_eq!(result.symbols[0].name, "Empty");
assert_eq!(result.symbols[0].kind, SymbolKind::Class);
}
#[test]
fn test_parse_nested_class() {
let result = parse("class Outer:\n class Inner:\n pass\n");
assert!(!result.has_errors);
assert_eq!(result.symbols.len(), 2);
assert_eq!(result.symbols[0].name, "Outer");
assert_eq!(result.symbols[1].name, "Inner");
assert_eq!(result.symbols[1].kind, SymbolKind::Class);
assert_eq!(result.symbols[1].parent_index, Some(0));
}
#[test]
fn test_decorated_function() {
let result = parse("@app.route(\"/\")\ndef index():\n pass\n");
assert!(!result.has_errors);
assert_eq!(result.symbols.len(), 1);
assert_eq!(result.symbols[0].name, "index");
assert_eq!(result.symbols[0].kind, SymbolKind::Function);
let sig = result.symbols[0].signature.as_ref().unwrap();
assert!(sig.contains("@app.route"), "sig = {}", sig);
assert!(sig.contains("def index"), "sig = {}", sig);
}
#[test]
fn test_multiple_decorators() {
let result = parse("@login_required\n@admin_only\ndef admin_view():\n pass\n");
assert!(!result.has_errors);
assert_eq!(result.symbols.len(), 1);
let sig = result.symbols[0].signature.as_ref().unwrap();
assert!(sig.contains("@login_required"), "sig = {}", sig);
assert!(sig.contains("@admin_only"), "sig = {}", sig);
}
#[test]
fn test_decorated_method_in_class() {
let result = parse("class Foo:\n @staticmethod\n def helper():\n pass\n");
assert!(!result.has_errors);
assert_eq!(result.symbols.len(), 2);
assert_eq!(result.symbols[1].name, "helper");
assert_eq!(result.symbols[1].kind, SymbolKind::Method);
assert_eq!(result.symbols[1].parent_index, Some(0));
let sig = result.symbols[1].signature.as_ref().unwrap();
assert!(sig.contains("@staticmethod"), "sig = {}", sig);
}
#[test]
fn test_module_constant_allcaps() {
let result = parse("MAX_SIZE = 100\n");
assert!(!result.has_errors);
assert_eq!(result.symbols.len(), 1);
assert_eq!(result.symbols[0].name, "MAX_SIZE");
assert_eq!(result.symbols[0].kind, SymbolKind::Variable);
}
#[test]
fn test_module_constant_typed() {
let result = parse("config: Config = get_config()\n");
assert!(!result.has_errors);
assert_eq!(result.symbols.len(), 1);
assert_eq!(result.symbols[0].name, "config");
assert_eq!(result.symbols[0].kind, SymbolKind::Variable);
}
#[test]
fn test_non_constant_skipped() {
let result = parse("result = compute()\n");
assert!(!result.has_errors);
assert_eq!(
result.symbols.len(),
0,
"lowercase non-typed assignment should not be extracted"
);
}
#[test]
fn test_class_level_assignment_skipped() {
let result = parse("class Foo:\n MAX = 100\n");
assert!(!result.has_errors);
let names: Vec<&str> = result.symbols.iter().map(|s| s.name.as_str()).collect();
assert!(names.contains(&"Foo"), "class should be extracted");
assert!(
!names.contains(&"MAX"),
"class-level constant should NOT be extracted, got: {:?}",
names
);
}
#[test]
fn test_method_vs_function() {
let result =
parse("def top():\n pass\n\nclass C:\n def inner(self):\n pass\n");
assert!(!result.has_errors);
let top = result.symbols.iter().find(|s| s.name == "top").unwrap();
let inner = result.symbols.iter().find(|s| s.name == "inner").unwrap();
assert_eq!(top.kind, SymbolKind::Function);
assert_eq!(inner.kind, SymbolKind::Method);
}
#[test]
fn test_empty_file() {
let result = parse("");
assert!(!result.has_errors);
assert_eq!(result.symbols.len(), 0);
}
#[test]
fn test_malformed_file() {
let result = parse("def broken(\n {}\n");
assert!(result.has_errors);
}
#[test]
fn test_import_simple() {
let result = parse("import os\n");
assert_eq!(result.imports.len(), 1);
assert_eq!(result.imports[0].path, "os");
assert_eq!(result.imports[0].imported_symbols, vec!["os"]);
}
#[test]
fn test_import_dotted() {
let result = parse("import os.path\n");
assert_eq!(result.imports.len(), 1);
assert_eq!(result.imports[0].path, "os.path");
assert_eq!(result.imports[0].imported_symbols, vec!["path"]);
}
#[test]
fn test_import_aliased() {
let result = parse("import numpy as np\n");
assert_eq!(result.imports.len(), 1);
assert_eq!(result.imports[0].path, "numpy");
assert_eq!(result.imports[0].imported_symbols, vec!["np"]);
}
#[test]
fn test_from_import() {
let result = parse("from os import path\n");
assert_eq!(result.imports.len(), 1);
assert_eq!(result.imports[0].path, "os");
assert!(
result.imports[0]
.imported_symbols
.contains(&"path".to_string()),
"imports: {:?}",
result.imports[0]
);
}
#[test]
fn test_from_import_multiple() {
let result = parse("from os.path import join, exists\n");
assert_eq!(result.imports.len(), 1);
assert_eq!(result.imports[0].path, "os.path");
assert!(result.imports[0]
.imported_symbols
.contains(&"join".to_string()));
assert!(result.imports[0]
.imported_symbols
.contains(&"exists".to_string()));
}
#[test]
fn test_from_import_relative() {
let result = parse("from . import utils\n");
assert_eq!(result.imports.len(), 1);
assert!(
result.imports[0]
.imported_symbols
.contains(&"utils".to_string()),
"imports: {:?}",
result.imports[0]
);
}
#[test]
fn test_from_import_wildcard() {
let result = parse("from foo import *\n");
assert_eq!(result.imports.len(), 1);
assert_eq!(result.imports[0].path, "foo");
assert!(
result.imports[0].imported_symbols.is_empty(),
"wildcard should have empty imported_symbols"
);
}
#[test]
fn test_function_call() {
let result = parse("def main():\n helper()\n\ndef helper():\n pass\n");
let call = result
.calls
.iter()
.find(|c| c.callee_name == "helper")
.expect("should find call to helper");
assert_eq!(call.caller_name, "main");
}
#[test]
fn test_method_call() {
let result = parse("def main():\n obj.process()\n");
let call = result
.calls
.iter()
.find(|c| c.callee_name == "process")
.expect("should find call to process");
assert_eq!(call.caller_name, "main");
}
#[test]
fn test_chained_method_call() {
let result = parse("def main():\n obj.foo().bar()\n");
let names: Vec<&str> = result
.calls
.iter()
.map(|c| c.callee_name.as_str())
.collect();
assert!(names.contains(&"foo"), "calls: {:?}", names);
assert!(names.contains(&"bar"), "calls: {:?}", names);
}
#[test]
fn test_builtin_call_filtered() {
let result = parse("def main():\n print(len(x))\n");
assert!(
result.calls.is_empty(),
"builtin calls should be filtered, got: {:?}",
result.calls
);
}
#[test]
fn test_type_ref_parameter() {
let result = parse("def foo(x: Config):\n pass\n");
let refs: Vec<&str> = result
.type_refs
.iter()
.map(|r| r.to_type.as_str())
.collect();
assert!(
refs.contains(&"Config"),
"should find Config type ref, got: {:?}",
refs
);
}
#[test]
fn test_type_ref_return() {
let result = parse("def foo() -> Result:\n pass\n");
let refs: Vec<&str> = result
.type_refs
.iter()
.map(|r| r.to_type.as_str())
.collect();
assert!(
refs.contains(&"Result"),
"should find Result type ref, got: {:?}",
refs
);
}
#[test]
fn test_type_ref_assignment() {
let result = parse("def foo():\n x: Config = get()\n");
let refs: Vec<&str> = result
.type_refs
.iter()
.map(|r| r.to_type.as_str())
.collect();
assert!(
refs.contains(&"Config"),
"should find Config type ref, got: {:?}",
refs
);
}
#[test]
fn test_builtin_type_filtered() {
let result = parse("def foo(x: int, y: str) -> bool:\n pass\n");
assert!(
result.type_refs.is_empty(),
"builtin types should be filtered, got: {:?}",
result.type_refs
);
}
fn setup_python_project() -> tempfile::TempDir {
let tmp = tempfile::tempdir().unwrap();
let foo = tmp.path().join("foo");
let utils = foo.join("utils");
let src = tmp.path().join("src");
std::fs::create_dir_all(&utils).unwrap();
std::fs::create_dir_all(&src).unwrap();
std::fs::write(foo.join("__init__.py"), "").unwrap();
std::fs::write(foo.join("bar.py"), "def greet(): pass").unwrap();
std::fs::write(utils.join("__init__.py"), "").unwrap();
std::fs::write(utils.join("helpers.py"), "def help(): pass").unwrap();
std::fs::write(src.join("app.py"), "def run(): pass").unwrap();
std::fs::write(tmp.path().join("main.py"), "import foo").unwrap();
tmp
}
#[test]
fn test_resolve_absolute_simple() {
let tmp = setup_python_project();
let result = resolve_python_import("foo.bar", "main.py", tmp.path());
assert!(result.is_some(), "should resolve foo.bar");
assert!(result.unwrap().ends_with("bar.py"));
}
#[test]
fn test_resolve_absolute_dotted_package() {
let tmp = setup_python_project();
let result = resolve_python_import("foo.utils", "main.py", tmp.path());
assert!(result.is_some(), "should resolve foo.utils to __init__.py");
let path = result.unwrap();
assert!(path.ends_with("__init__.py"));
assert!(path.to_string_lossy().contains("utils"));
}
#[test]
fn test_resolve_absolute_from_src() {
let tmp = setup_python_project();
let result = resolve_python_import("app", "main.py", tmp.path());
assert!(result.is_some(), "should resolve app from src/ fallback");
assert!(result.unwrap().ends_with("app.py"));
}
#[test]
fn test_resolve_relative_single_dot() {
let tmp = setup_python_project();
let result = resolve_python_import(".bar", "foo/__init__.py", tmp.path());
assert!(result.is_some(), "should resolve .bar from foo/");
assert!(result.unwrap().ends_with("bar.py"));
}
#[test]
fn test_resolve_relative_double_dot() {
let tmp = setup_python_project();
let result = resolve_python_import("..bar", "foo/utils/helpers.py", tmp.path());
assert!(result.is_some(), "should resolve ..bar from foo/utils/");
assert!(result.unwrap().ends_with("bar.py"));
}
#[test]
fn test_resolve_package_init() {
let tmp = setup_python_project();
let result = resolve_python_import("foo", "main.py", tmp.path());
assert!(result.is_some(), "should resolve foo to __init__.py");
assert!(result.unwrap().ends_with("__init__.py"));
}
#[test]
fn test_resolve_external_returns_none() {
let tmp = setup_python_project();
let result = resolve_python_import("numpy", "main.py", tmp.path());
assert!(result.is_none(), "external package should return None");
}
#[test]
fn test_resolve_empty_path() {
let tmp = setup_python_project();
let result = resolve_python_import("", "main.py", tmp.path());
assert!(result.is_none());
}
#[test]
fn test_resolve_relative_dot_init() {
let tmp = setup_python_project();
std::fs::write(tmp.path().join("foo/__init__.py"), "from . import utils").unwrap();
let result = resolve_python_import(".utils", "foo/__init__.py", tmp.path());
assert!(result.is_some(), "should resolve .utils from foo/");
}
#[test]
fn test_resolve_deep_dotted() {
let tmp = setup_python_project();
let result = resolve_python_import("foo.utils.helpers", "main.py", tmp.path());
assert!(result.is_some(), "should resolve foo.utils.helpers");
assert!(result.unwrap().ends_with("helpers.py"));
}
#[test]
fn test_decorated_class() {
let result = parse(
"@dataclass\nclass Config:\n name: str\n def validate(self):\n pass\n",
);
assert!(!result.has_errors);
let config = result.symbols.iter().find(|s| s.name == "Config");
assert!(config.is_some(), "decorated class should be extracted");
assert_eq!(config.unwrap().kind, SymbolKind::Class);
let sig = config.unwrap().signature.as_ref().unwrap();
assert!(
sig.contains("@dataclass"),
"sig should contain decorator: {}",
sig
);
assert!(
sig.contains("class Config"),
"sig should contain class def: {}",
sig
);
let validate = result.symbols.iter().find(|s| s.name == "validate");
assert!(
validate.is_some(),
"method in decorated class should be extracted"
);
assert_eq!(validate.unwrap().kind, SymbolKind::Method);
}
#[test]
fn test_decorated_class_with_methods_parent_index() {
let result = parse("@decorator\nclass Svc:\n def run(self):\n pass\n def stop(self):\n pass\n");
assert!(!result.has_errors);
let svc = result.symbols.iter().find(|s| s.name == "Svc").unwrap();
assert_eq!(svc.kind, SymbolKind::Class);
let svc_idx = result.symbols.iter().position(|s| s.name == "Svc").unwrap();
for method in result
.symbols
.iter()
.filter(|s| s.kind == SymbolKind::Method)
{
assert_eq!(
method.parent_index,
Some(svc_idx),
"method '{}' should have Svc as parent",
method.name
);
}
}
#[test]
fn test_resolve_relative_dot_only_to_init() {
let tmp = setup_python_project();
let result = resolve_python_import(".", "foo/bar.py", tmp.path());
assert!(
result.is_some(),
"bare dot should resolve to __init__.py in foo/"
);
assert!(result.unwrap().ends_with("__init__.py"));
}
#[test]
fn test_resolve_relative_dot_only_no_init() {
let tmp = tempfile::tempdir().unwrap();
let pkg = tmp.path().join("pkg");
std::fs::create_dir_all(&pkg).unwrap();
std::fs::write(pkg.join("mod.py"), "").unwrap();
let result = resolve_python_import(".", "pkg/mod.py", tmp.path());
assert!(
result.is_none(),
"bare dot without __init__.py should return None"
);
}
#[test]
fn test_resolve_module_path_empty_parts() {
let tmp = tempfile::tempdir().unwrap();
let pkg = tmp.path().join("pkg");
std::fs::create_dir_all(&pkg).unwrap();
std::fs::write(pkg.join("__init__.py"), "").unwrap();
let result = resolve_module_path(&[], &pkg);
assert!(
result.is_some(),
"empty parts with __init__.py should resolve"
);
assert!(result.unwrap().ends_with("__init__.py"));
}
#[test]
fn test_resolve_module_path_empty_parts_no_init() {
let tmp = tempfile::tempdir().unwrap();
let pkg = tmp.path().join("pkg");
std::fs::create_dir_all(&pkg).unwrap();
let result = resolve_module_path(&[], &pkg);
assert!(
result.is_none(),
"empty parts without __init__.py should return None"
);
}
#[test]
fn test_resolve_intermediate_not_dir() {
let tmp = setup_python_project();
let result = resolve_python_import("foo.bar.nonexistent", "main.py", tmp.path());
assert!(result.is_none(), "intermediate non-directory should fail");
}
#[test]
fn test_import_aliased_from() {
let result = parse("from os import path as p\n");
assert_eq!(result.imports.len(), 1);
assert_eq!(result.imports[0].path, "os");
assert!(
result.imports[0]
.imported_symbols
.contains(&"p".to_string()),
"aliased import should use alias name: {:?}",
result.imports[0].imported_symbols
);
}
#[test]
fn test_import_from_aliased_without_alias() {
let result = parse("from os.path import join, exists\n");
assert_eq!(result.imports.len(), 1);
assert!(result.imports[0]
.imported_symbols
.contains(&"join".to_string()));
assert!(result.imports[0]
.imported_symbols
.contains(&"exists".to_string()));
}
#[test]
fn test_multiple_plain_imports() {
let result = parse("import os\nimport sys\nimport json\n");
assert_eq!(result.imports.len(), 3);
let paths: Vec<&str> = result.imports.iter().map(|i| i.path.as_str()).collect();
assert!(paths.contains(&"os"));
assert!(paths.contains(&"sys"));
assert!(paths.contains(&"json"));
}
#[test]
fn test_triple_dot_relative_import() {
let tmp = tempfile::tempdir().unwrap();
let deep = tmp.path().join("a").join("b").join("c");
std::fs::create_dir_all(&deep).unwrap();
std::fs::write(deep.join("mod.py"), "").unwrap();
std::fs::write(tmp.path().join("a").join("target.py"), "").unwrap();
let result = resolve_python_import("...target", "a/b/c/mod.py", tmp.path());
assert!(result.is_some(), "triple dot should resolve two levels up");
assert!(result.unwrap().ends_with("target.py"));
}
#[test]
fn test_resolve_nonexistent_module() {
let tmp = setup_python_project();
let result = resolve_python_import("foo.nonexistent", "main.py", tmp.path());
assert!(result.is_none(), "nonexistent module should return None");
}
#[test]
fn test_call_with_attribute_access() {
let result = parse("def main():\n config.get_value()\n");
let calls: Vec<&str> = result
.calls
.iter()
.map(|c| c.callee_name.as_str())
.collect();
assert!(
calls.contains(&"get_value"),
"should extract attribute call: {:?}",
calls
);
}
#[test]
fn test_type_ref_builtin_type_not_extracted() {
let result = parse("def foo(x: int, y: str, z: list) -> bool:\n pass\n");
assert!(
result.type_refs.is_empty(),
"builtin types should be filtered: {:?}",
result.type_refs
);
}
#[test]
fn test_resolve_from_lib_directory() {
let tmp = tempfile::tempdir().unwrap();
let lib = tmp.path().join("lib");
std::fs::create_dir_all(&lib).unwrap();
std::fs::write(lib.join("mymod.py"), "").unwrap();
let result = resolve_python_import("mymod", "main.py", tmp.path());
assert!(result.is_some(), "should resolve from lib/ fallback");
assert!(result.unwrap().ends_with("mymod.py"));
}
#[test]
fn test_expression_statement_not_assignment() {
let result = parse("print('hello')\n");
assert!(!result.has_errors);
assert_eq!(
result.symbols.len(),
0,
"bare expression should not produce symbols"
);
}
#[test]
fn test_assignment_left_not_identifier() {
let result = parse("x, y = 1, 2\n");
assert!(!result.has_errors);
}
#[test]
fn test_is_python_builtin_positive() {
assert!(is_python_builtin("print"));
assert!(is_python_builtin("len"));
assert!(is_python_builtin("isinstance"));
assert!(is_python_builtin("ValueError"));
assert!(is_python_builtin("super"));
assert!(is_python_builtin("object"));
}
#[test]
fn test_is_python_builtin_negative() {
assert!(!is_python_builtin("my_function"));
assert!(!is_python_builtin("CustomError"));
assert!(!is_python_builtin(""));
}
#[test]
fn test_is_python_builtin_type_positive() {
assert!(is_python_builtin_type("int"));
assert!(is_python_builtin_type("str"));
assert!(is_python_builtin_type("Optional"));
assert!(is_python_builtin_type("List"));
assert!(is_python_builtin_type("TypeVar"));
assert!(is_python_builtin_type("Self"));
assert!(is_python_builtin_type("Never"));
assert!(is_python_builtin_type("NoReturn"));
}
#[test]
fn test_is_python_builtin_type_negative() {
assert!(!is_python_builtin_type("MyClass"));
assert!(!is_python_builtin_type("AppConfig"));
assert!(!is_python_builtin_type(""));
}
#[test]
fn test_type_ref_inside_class_method() {
let result =
parse("class Svc:\n def process(self, cfg: AppConfig) -> Response:\n pass\n");
assert!(!result.has_errors);
let ref_types: Vec<&str> = result
.type_refs
.iter()
.map(|r| r.to_type.as_str())
.collect();
assert!(
ref_types.contains(&"AppConfig"),
"should find AppConfig type ref: {:?}",
ref_types
);
assert!(
ref_types.contains(&"Response"),
"should find Response type ref: {:?}",
ref_types
);
}
#[test]
fn test_module_constant_underscore_allcaps() {
let result = parse("MAX_RETRY_COUNT = 3\n");
assert!(!result.has_errors);
assert_eq!(result.symbols.len(), 1);
assert_eq!(result.symbols[0].name, "MAX_RETRY_COUNT");
assert_eq!(result.symbols[0].kind, SymbolKind::Variable);
}
#[test]
fn test_resolve_absolute_nonexistent_package() {
let tmp = setup_python_project();
let result = resolve_python_import("totally_missing.submod", "main.py", tmp.path());
assert!(result.is_none());
}
#[test]
fn test_resolve_relative_too_many_dots() {
let tmp = tempfile::tempdir().unwrap();
let pkg = tmp.path().join("a");
std::fs::create_dir_all(&pkg).unwrap();
std::fs::write(pkg.join("mod.py"), "").unwrap();
let result = resolve_python_import("....target", "a/mod.py", tmp.path());
assert!(result.is_none(), "too many dots should return None");
}
#[test]
fn test_import_plain_multiple_dotted_names() {
let result = parse("import os, sys, json\n");
assert!(!result.has_errors);
let paths: Vec<&str> = result.imports.iter().map(|i| i.path.as_str()).collect();
assert!(paths.contains(&"os"), "imports: {:?}", paths);
assert!(paths.contains(&"sys"), "imports: {:?}", paths);
assert!(paths.contains(&"json"), "imports: {:?}", paths);
}
#[test]
fn test_from_import_with_alias() {
let result = parse("from os import path as p\n");
assert_eq!(result.imports.len(), 1);
assert_eq!(result.imports[0].path, "os");
assert!(
result.imports[0]
.imported_symbols
.contains(&"p".to_string()),
"aliased import should use alias name: {:?}",
result.imports[0].imported_symbols
);
}
#[test]
fn test_call_identifier_branch() {
let result = parse("def main():\n do_work()\n");
assert!(!result.has_errors);
let calls: Vec<&str> = result
.calls
.iter()
.map(|c| c.callee_name.as_str())
.collect();
assert!(
calls.contains(&"do_work"),
"should capture identifier call: {:?}",
calls
);
}
#[test]
fn test_decorated_definition_no_definition_node() {
let result = parse("@property\ndef x(self):\n return 1\n");
assert!(!result.has_errors);
assert_eq!(result.symbols.len(), 1);
assert_eq!(result.symbols[0].name, "x");
}
#[test]
fn test_class_with_inheritance() {
let result = parse("class Child(Parent):\n def method(self):\n pass\n");
assert!(!result.has_errors);
let child = result.symbols.iter().find(|s| s.name == "Child").unwrap();
assert_eq!(child.kind, SymbolKind::Class);
let sig = child.signature.as_ref().unwrap();
assert!(
sig.contains("Parent"),
"sig should contain base class: {}",
sig
);
}
#[test]
fn test_class_multiple_inheritance() {
let result = parse("class Multi(Base1, Base2, Mixin):\n pass\n");
assert!(!result.has_errors);
assert_eq!(result.symbols.len(), 1);
assert_eq!(result.symbols[0].name, "Multi");
let sig = result.symbols[0].signature.as_ref().unwrap();
assert!(sig.contains("Base1"), "sig = {}", sig);
assert!(sig.contains("Mixin"), "sig = {}", sig);
}
#[test]
fn test_property_decorated_method() {
let result =
parse("class Foo:\n @property\n def name(self):\n return self._name\n");
assert!(!result.has_errors);
let name_method = result.symbols.iter().find(|s| s.name == "name").unwrap();
assert_eq!(name_method.kind, SymbolKind::Method);
let sig = name_method.signature.as_ref().unwrap();
assert!(sig.contains("@property"), "sig = {}", sig);
}
#[test]
fn test_classmethod_decorated() {
let result = parse("class Foo:\n @classmethod\n def create(cls):\n pass\n");
assert!(!result.has_errors);
let create = result.symbols.iter().find(|s| s.name == "create").unwrap();
assert_eq!(create.kind, SymbolKind::Method);
let sig = create.signature.as_ref().unwrap();
assert!(sig.contains("@classmethod"), "sig = {}", sig);
}
#[test]
fn test_function_with_complex_signature() {
let result = parse(
"def process(items: list[str], *, timeout: int = 30) -> dict[str, Any]:\n pass\n",
);
assert!(!result.has_errors);
assert_eq!(result.symbols.len(), 1);
assert_eq!(result.symbols[0].name, "process");
let sig = result.symbols[0].signature.as_ref().unwrap();
assert!(sig.contains("def process"), "sig = {}", sig);
}
#[test]
fn test_call_inside_method() {
let result =
parse("class Svc:\n def run(self):\n self.setup()\n process()\n");
let calls: Vec<&str> = result
.calls
.iter()
.map(|c| c.callee_name.as_str())
.collect();
assert!(
calls.contains(&"setup"),
"should find self.setup(): {:?}",
calls
);
assert!(
calls.contains(&"process"),
"should find process(): {:?}",
calls
);
}
#[test]
fn test_type_ref_generic_annotation() {
let result = parse("def foo(items: List[Config]) -> Optional[Response]:\n pass\n");
let refs: Vec<&str> = result
.type_refs
.iter()
.map(|r| r.to_type.as_str())
.collect();
assert!(refs.contains(&"Config"), "type_refs: {:?}", refs);
assert!(refs.contains(&"Response"), "type_refs: {:?}", refs);
}
#[test]
fn test_multiple_functions_and_classes() {
let result = parse(
"def a():\n pass\ndef b():\n pass\nclass C:\n def d(self):\n pass\n",
);
assert!(!result.has_errors);
let names: Vec<&str> = result.symbols.iter().map(|s| s.name.as_str()).collect();
assert_eq!(names, vec!["a", "b", "C", "d"]);
}
#[test]
fn test_module_constant_with_numbers_not_extracted() {
let result = parse("MAX_SIZE_3 = 1024\n");
assert!(!result.has_errors);
assert_eq!(
result.symbols.len(),
0,
"allcaps with digits should NOT be extracted (digits fail the check)"
);
}
#[test]
fn test_module_constant_pure_allcaps() {
let result = parse("MAX_SIZE = 1024\n");
assert!(!result.has_errors);
assert_eq!(result.symbols.len(), 1);
assert_eq!(result.symbols[0].name, "MAX_SIZE");
}
#[test]
fn test_from_import_relative_double_dot() {
let result = parse("from .. import utils\n");
assert!(!result.has_errors);
assert_eq!(result.imports.len(), 1);
assert!(result.imports[0]
.imported_symbols
.contains(&"utils".to_string()));
}
#[test]
fn test_language_parser_trait_name() {
let parser = PythonParser;
assert_eq!(parser.language_name(), "python");
}
#[test]
fn test_language_parser_parse_via_trait() {
let parser = PythonParser;
let result = parser.parse("def foo():\n pass\n", &PathBuf::from("test.py"));
assert!(!result.has_errors);
assert_eq!(result.symbols.len(), 1);
assert_eq!(result.symbols[0].name, "foo");
}
#[test]
fn test_resolve_from_init_subpackage() {
let tmp = setup_python_project();
let result = resolve_python_import("foo.utils.helpers", "foo/__init__.py", tmp.path());
assert!(result.is_some(), "should resolve deep submodule from init");
assert!(result.unwrap().ends_with("helpers.py"));
}
#[test]
fn test_call_nested_in_class_method() {
let result =
parse("class A:\n def m(self):\n x = helper()\n y = other()\n");
let callers: Vec<&str> = result
.calls
.iter()
.map(|c| c.caller_name.as_str())
.collect();
for caller in &callers {
assert_eq!(*caller, "m");
}
}
}