Skip to main content

dk_engine/parser/langs/
python.rs

1//! Python language configuration for the query-driven parser.
2
3use crate::parser::lang_config::{CommentStyle, LanguageConfig};
4use dk_core::{Symbol, Visibility};
5use tree_sitter::Language;
6
7/// Python language configuration for [`QueryDrivenParser`](crate::parser::engine::QueryDrivenParser).
8pub struct PythonConfig;
9
10impl LanguageConfig for PythonConfig {
11    fn language(&self) -> Language {
12        tree_sitter_python::LANGUAGE.into()
13    }
14
15    fn extensions(&self) -> &'static [&'static str] {
16        &["py"]
17    }
18
19    fn symbols_query(&self) -> &'static str {
20        include_str!("../queries/python_symbols.scm")
21    }
22
23    fn calls_query(&self) -> &'static str {
24        include_str!("../queries/python_calls.scm")
25    }
26
27    fn imports_query(&self) -> &'static str {
28        include_str!("../queries/python_imports.scm")
29    }
30
31    fn comment_style(&self) -> CommentStyle {
32        CommentStyle::Hash
33    }
34
35    fn resolve_visibility(&self, _modifiers: Option<&str>, name: &str) -> Visibility {
36        if name.starts_with('_') {
37            Visibility::Private
38        } else {
39            Visibility::Public
40        }
41    }
42
43    fn adjust_symbol(&self, sym: &mut Symbol, node: &tree_sitter::Node, source: &[u8]) {
44        // 1. Decorated definitions: expand span to include decorator(s).
45        //    If this function/class is inside a `decorated_definition`, use
46        //    the parent's full span and its first line as the signature.
47        if let Some(parent) = node.parent() {
48            if parent.kind() == "decorated_definition" {
49                sym.span = dk_core::Span {
50                    start_byte: parent.start_byte() as u32,
51                    end_byte: parent.end_byte() as u32,
52                };
53                // Use the decorator's first line as the signature.
54                let text =
55                    std::str::from_utf8(&source[parent.start_byte()..parent.end_byte()])
56                        .unwrap_or("");
57                if let Some(first_line) = text.lines().next() {
58                    let trimmed = first_line.trim();
59                    if !trimmed.is_empty() {
60                        sym.signature = Some(trimmed.to_string());
61                    }
62                }
63            }
64        }
65
66        // 2. Docstrings: for functions and classes, check the body for a
67        //    leading expression_statement containing a string (triple-quoted).
68        //    Docstrings take priority over `# comment` blocks collected by
69        //    the engine's `collect_doc_comments`.
70        if matches!(
71            sym.kind,
72            dk_core::SymbolKind::Function | dk_core::SymbolKind::Class
73        ) {
74            if let Some(docstring) = Self::extract_docstring(node, source) {
75                sym.doc_comment = Some(docstring);
76            }
77        }
78    }
79
80    fn is_external_import(&self, module_path: &str) -> bool {
81        !module_path.starts_with('.')
82    }
83}
84
85impl PythonConfig {
86    /// Extract a docstring from a function or class body.
87    ///
88    /// In Python, a docstring is the first statement in the body if it is an
89    /// `expression_statement` containing a `string` node.
90    fn extract_docstring(node: &tree_sitter::Node, source: &[u8]) -> Option<String> {
91        let body = node.child_by_field_name("body")?;
92        let first_stmt = body.named_child(0)?;
93
94        if first_stmt.kind() == "expression_statement" {
95            let expr = first_stmt.child(0)?;
96            if expr.kind() == "string" {
97                let raw = std::str::from_utf8(&source[expr.start_byte()..expr.end_byte()])
98                    .unwrap_or("");
99                let content = raw
100                    .strip_prefix("\"\"\"")
101                    .and_then(|s| s.strip_suffix("\"\"\""))
102                    .or_else(|| raw.strip_prefix("'''").and_then(|s| s.strip_suffix("'''")))
103                    .unwrap_or(raw);
104                let trimmed = content.trim().to_string();
105                if !trimmed.is_empty() {
106                    return Some(trimmed);
107                }
108            }
109        }
110
111        None
112    }
113}