use crate::parser::lang_config::{CommentStyle, LanguageConfig};
use dk_core::{Symbol, Visibility};
use tree_sitter::Language;
pub struct PythonConfig;
impl LanguageConfig for PythonConfig {
fn language(&self) -> Language {
tree_sitter_python::LANGUAGE.into()
}
fn extensions(&self) -> &'static [&'static str] {
&["py"]
}
fn symbols_query(&self) -> &'static str {
include_str!("../queries/python_symbols.scm")
}
fn calls_query(&self) -> &'static str {
include_str!("../queries/python_calls.scm")
}
fn imports_query(&self) -> &'static str {
include_str!("../queries/python_imports.scm")
}
fn comment_style(&self) -> CommentStyle {
CommentStyle::Hash
}
fn resolve_visibility(&self, _modifiers: Option<&str>, name: &str) -> Visibility {
if name.starts_with('_') {
Visibility::Private
} else {
Visibility::Public
}
}
fn adjust_symbol(&self, sym: &mut Symbol, node: &tree_sitter::Node, source: &[u8]) {
if let Some(parent) = node.parent() {
if parent.kind() == "decorated_definition" {
sym.span = dk_core::Span {
start_byte: parent.start_byte() as u32,
end_byte: parent.end_byte() as u32,
};
let text =
std::str::from_utf8(&source[parent.start_byte()..parent.end_byte()])
.unwrap_or("");
if let Some(first_line) = text.lines().next() {
let trimmed = first_line.trim();
if !trimmed.is_empty() {
sym.signature = Some(trimmed.to_string());
}
}
}
}
if matches!(
sym.kind,
dk_core::SymbolKind::Function | dk_core::SymbolKind::Class
) {
if let Some(docstring) = Self::extract_docstring(node, source) {
sym.doc_comment = Some(docstring);
}
}
}
fn is_external_import(&self, module_path: &str) -> bool {
!module_path.starts_with('.')
}
}
impl PythonConfig {
fn extract_docstring(node: &tree_sitter::Node, source: &[u8]) -> Option<String> {
let body = node.child_by_field_name("body")?;
let first_stmt = body.named_child(0)?;
if first_stmt.kind() == "expression_statement" {
let expr = first_stmt.child(0)?;
if expr.kind() == "string" {
let raw = std::str::from_utf8(&source[expr.start_byte()..expr.end_byte()])
.unwrap_or("");
let content = raw
.strip_prefix("\"\"\"")
.and_then(|s| s.strip_suffix("\"\"\""))
.or_else(|| raw.strip_prefix("'''").and_then(|s| s.strip_suffix("'''")))
.unwrap_or(raw);
let trimmed = content.trim().to_string();
if !trimmed.is_empty() {
return Some(trimmed);
}
}
}
None
}
}