use tree_sitter::Node;
#[derive(Debug)]
pub struct DecoratorInfo {
pub name: String,
pub full_name: String, pub args: Vec<String>,
}
pub fn collect_decorators(node: &Node, source: &[u8]) -> Vec<DecoratorInfo> {
let mut decorators = Vec::new();
if node.kind() == "decorated_definition" {
for i in 0..node.named_child_count() {
let child = node.named_child(i).unwrap();
if child.kind() == "decorator"
&& let Some(info) = parse_decorator(&child, source)
{
decorators.push(info);
}
}
}
decorators
}
fn parse_decorator(node: &Node, source: &[u8]) -> Option<DecoratorInfo> {
let mut full_name = String::new();
let mut args = Vec::new();
for i in 0..node.named_child_count() {
let child = node.named_child(i).unwrap();
match child.kind() {
"identifier" => {
full_name = node_text(&child, source);
}
"attribute" => {
full_name = node_text(&child, source);
}
"call" => {
if let Some(func) = child.child_by_field_name("function") {
full_name = node_text(&func, source);
}
if let Some(arg_list) = child.child_by_field_name("arguments") {
for j in 0..arg_list.named_child_count() {
let arg = arg_list.named_child(j).unwrap();
args.push(node_text(&arg, source));
}
}
}
_ => {}
}
}
if full_name.is_empty() {
return None;
}
let name = full_name
.rsplit('.')
.next()
.unwrap_or(&full_name)
.to_string();
Some(DecoratorInfo {
name,
full_name,
args,
})
}
pub fn node_text(node: &Node, source: &[u8]) -> String {
node.utf8_text(source).unwrap_or("").to_string()
}
pub fn clean_string_literal(s: &str) -> String {
s.trim_matches('"')
.trim_matches('\'')
.trim_start_matches("f\"")
.trim_start_matches("f'")
.to_string()
}
pub fn to_kebab_case(s: &str) -> String {
let mut result = String::new();
for (i, c) in s.chars().enumerate() {
if c == '_' {
result.push('-');
} else if c.is_uppercase() && i > 0 {
result.push('-');
result.push(c.to_ascii_lowercase());
} else {
result.push(c.to_ascii_lowercase());
}
}
result
}
pub fn get_def_name(node: &Node, source: &[u8]) -> String {
node.child_by_field_name("name")
.map(|n| node_text(&n, source))
.unwrap_or_else(|| "unknown".to_string())
}
pub fn get_inner_definition<'a>(node: &'a Node<'a>) -> Option<Node<'a>> {
if node.kind() == "decorated_definition" {
node.named_children(&mut node.walk())
.find(|c| c.kind() == "function_definition" || c.kind() == "class_definition")
} else {
Some(*node)
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::context::ParsedFile;
fn parse_file(source: &str, path: &str) -> ParsedFile {
ParsedFile::parse(source.to_string(), path.to_string()).unwrap()
}
#[test]
fn test_parse_python_decorators() {
let source = r#"
from fastapi import APIRouter
router = APIRouter()
@router.get("/users/{user_id}")
async def get_user(user_id: int):
return {"user_id": user_id}
@router.post("/users")
async def create_user(user: UserCreate):
return user
"#;
let file = parse_file(source, "routes.py");
let root = file.tree.root_node();
let src = file.source.as_bytes();
let mut found = Vec::new();
for i in 0..root.named_child_count() {
let node = root.named_child(i).unwrap();
if node.kind() == "decorated_definition" {
let decorators = collect_decorators(&node, src);
for d in &decorators {
found.push((d.name.clone(), d.full_name.clone(), d.args.clone()));
}
}
}
assert_eq!(found.len(), 2);
assert_eq!(found[0].0, "get"); assert_eq!(found[0].1, "router.get"); assert!(found[0].2[0].contains("/users/"));
assert_eq!(found[1].0, "post");
assert_eq!(found[1].1, "router.post");
}
#[test]
fn test_to_kebab_case() {
assert_eq!(to_kebab_case("UserService"), "user-service");
assert_eq!(to_kebab_case("user_service"), "user-service");
assert_eq!(to_kebab_case("get_users"), "get-users");
}
#[test]
fn test_clean_string_literal() {
assert_eq!(clean_string_literal("\"hello\""), "hello");
assert_eq!(clean_string_literal("'/api/users'"), "/api/users");
}
}