use super::{extract_doc_comment, extract_signature, field_text, node_text};
use crate::types::*;
pub fn extract(tree: &tree_sitter::Tree, source: &[u8]) -> (Vec<Symbol>, Vec<Import>) {
let root = tree.root_node();
let mut symbols = Vec::new();
let mut imports = Vec::new();
let mut cursor = root.walk();
for child in root.children(&mut cursor) {
extract_top_level(&child, source, &mut symbols, &mut imports);
}
(symbols, imports)
}
fn extract_top_level(
child: &tree_sitter::Node,
source: &[u8],
symbols: &mut Vec<Symbol>,
imports: &mut Vec<Import>,
) {
match child.kind() {
"function_definition" => {
if let Some(sym) = extract_function(child, source, None) {
symbols.push(sym);
}
}
"class_definition" => {
if let Some(sym) = extract_class(child, source) {
symbols.push(sym);
}
}
"decorated_definition" => {
let mut inner_cursor = child.walk();
for inner in child.children(&mut inner_cursor) {
match inner.kind() {
"class_definition" => {
if let Some(sym) = extract_class(&inner, source) {
symbols.push(sym);
}
}
"function_definition" => {
if let Some(sym) = extract_function(&inner, source, None) {
symbols.push(sym);
}
}
_ => {}
}
}
}
"import_statement" | "import_from_statement" => {
imports.push(Import {
path: node_text(child, source).to_string(),
alias: None,
span: Span::from_node(child),
});
}
"expression_statement" => {
let mut inner_cursor = child.walk();
for inner in child.children(&mut inner_cursor) {
if inner.kind() == "assignment" {
if let Some(left) = inner.child_by_field_name("left") {
let name = node_text(&left, source);
if name.chars().all(|c| c.is_uppercase() || c == '_') && !name.is_empty() {
symbols.push(Symbol {
name: name.to_string(),
kind: SymbolKind::Const,
span: Span::from_node(&inner),
signature: node_text(&inner, source).to_string(),
doc_comment: None,
parent: None,
children: Vec::new(),
});
}
}
}
}
}
_ => {}
}
}
fn extract_function(
node: &tree_sitter::Node,
source: &[u8],
parent: Option<&str>,
) -> Option<Symbol> {
let name = field_text(node, "name", source)?;
let is_method = parent.is_some();
let mut sig = extract_signature(node, "body", source);
if let Some(prev) = node.prev_sibling() {
if prev.kind() == "decorator" {
let dec_text = node_text(&prev, source);
sig = format!("{}\n{}", dec_text, sig);
}
}
Some(Symbol {
name: name.to_string(),
kind: if is_method {
SymbolKind::Method
} else {
SymbolKind::Function
},
span: Span::from_node(node),
signature: sig,
doc_comment: extract_docstring(node, source),
parent: parent.map(|s| s.to_string()),
children: Vec::new(),
})
}
fn extract_class(node: &tree_sitter::Node, source: &[u8]) -> Option<Symbol> {
let name = field_text(node, "name", source)?;
let mut methods = Vec::new();
if let Some(body) = node.child_by_field_name("body") {
let mut cursor = body.walk();
for child in body.children(&mut cursor) {
match child.kind() {
"function_definition" => {
if let Some(m) = extract_function(&child, source, Some(name)) {
methods.push(m);
}
}
"decorated_definition" => {
let mut inner = child.walk();
for grandchild in child.children(&mut inner) {
if grandchild.kind() == "function_definition" {
if let Some(m) = extract_function(&grandchild, source, Some(name)) {
methods.push(m);
}
}
}
}
_ => {}
}
}
}
Some(Symbol {
name: name.to_string(),
kind: SymbolKind::Class,
span: Span::from_node(node),
signature: extract_signature(node, "body", source),
doc_comment: extract_docstring(node, source),
parent: None,
children: methods,
})
}
fn extract_docstring(node: &tree_sitter::Node, source: &[u8]) -> Option<String> {
let body = node.child_by_field_name("body")?;
let first = body.child(0)?;
if first.kind() == "expression_statement" {
let expr = first.child(0)?;
if expr.kind() == "string" {
return Some(node_text(&expr, source).to_string());
}
}
extract_doc_comment(node, source)
}