sqry_lang_python/
lib.rs

1//! Python language plugin
2//!
3//! Provides AST parsing, scope extraction, and graph builder integration.
4
5pub mod relations;
6
7pub use relations::PythonGraphBuilder;
8
9use sqry_core::ast::{Scope, ScopeId, link_nested_scopes};
10use sqry_core::plugin::{
11    LanguageMetadata, LanguagePlugin,
12    error::{ParseError, ScopeError},
13};
14use std::path::Path;
15use streaming_iterator::StreamingIterator;
16use tree_sitter::{Language, Parser, Query, QueryCursor, Tree};
17
18const PLUGIN_ID: &str = "python";
19const TREE_SITTER_VERSION: &str = "0.23";
20
21struct ScopeCapture {
22    scope_type: String,
23    scope_name: String,
24    start: tree_sitter::Point,
25    end: tree_sitter::Point,
26}
27
28/// Python language plugin implementation
29pub struct PythonPlugin {
30    graph_builder: PythonGraphBuilder,
31}
32
33impl PythonPlugin {
34    /// Creates a new Python plugin instance.
35    #[must_use]
36    pub fn new() -> Self {
37        Self {
38            graph_builder: PythonGraphBuilder::default(),
39        }
40    }
41}
42
43impl Default for PythonPlugin {
44    fn default() -> Self {
45        Self::new()
46    }
47}
48
49impl LanguagePlugin for PythonPlugin {
50    fn metadata(&self) -> LanguageMetadata {
51        LanguageMetadata {
52            id: PLUGIN_ID,
53            name: "Python",
54            version: env!("CARGO_PKG_VERSION"),
55            author: "Verivus Pty Ltd",
56            description: "Python language support for sqry",
57            tree_sitter_version: TREE_SITTER_VERSION,
58        }
59    }
60
61    fn extensions(&self) -> &'static [&'static str] {
62        &["py", "pyi"]
63    }
64
65    fn language(&self) -> Language {
66        tree_sitter_python::LANGUAGE.into()
67    }
68
69    fn parse_ast(&self, content: &[u8]) -> Result<Tree, ParseError> {
70        let mut parser = Parser::new();
71        parser
72            .set_language(&self.language())
73            .map_err(|e| ParseError::LanguageSetFailed(e.to_string()))?;
74        parser
75            .parse(content, None)
76            .ok_or(ParseError::TreeSitterFailed)
77    }
78
79    fn extract_scopes(
80        &self,
81        tree: &Tree,
82        content: &[u8],
83        file_path: &Path,
84    ) -> Result<Vec<Scope>, ScopeError> {
85        Self::extract_python_scopes(tree, content, file_path)
86    }
87
88    fn graph_builder(&self) -> Option<&dyn sqry_core::graph::GraphBuilder> {
89        Some(&self.graph_builder)
90    }
91}
92
93impl PythonPlugin {
94    fn extract_python_scopes(
95        tree: &Tree,
96        content: &[u8],
97        file_path: &Path,
98    ) -> Result<Vec<Scope>, ScopeError> {
99        let root_node = tree.root_node();
100        let language = tree_sitter_python::LANGUAGE.into();
101
102        let scope_query = Self::scope_query_source();
103        let query = Query::new(&language, scope_query)
104            .map_err(|e| ScopeError::QueryCompilationFailed(e.to_string()))?;
105
106        let mut scopes = Vec::new();
107        let mut cursor = QueryCursor::new();
108        let mut query_matches = cursor.matches(&query, root_node, content);
109
110        while let Some(m) = query_matches.next() {
111            if let Some(scope) = Self::scope_from_match(&query, m, content, file_path) {
112                scopes.push(scope);
113            }
114        }
115
116        scopes.sort_by_key(|s| (s.start_line, s.start_column));
117
118        link_nested_scopes(&mut scopes);
119        Ok(scopes)
120    }
121
122    fn scope_from_match(
123        query: &Query,
124        match_: &tree_sitter::QueryMatch<'_, '_>,
125        content: &[u8],
126        file_path: &Path,
127    ) -> Option<Scope> {
128        let capture = Self::scope_capture_from_match(query, match_, content)?;
129        Some(Self::build_scope_from_capture(capture, file_path))
130    }
131
132    fn scope_capture_from_match(
133        query: &Query,
134        match_: &tree_sitter::QueryMatch<'_, '_>,
135        content: &[u8],
136    ) -> Option<ScopeCapture> {
137        let mut scope_type = None;
138        let mut scope_name = None;
139        let mut scope_start = None;
140        let mut scope_end = None;
141
142        for capture in match_.captures {
143            let capture_name = query.capture_names()[capture.index as usize];
144            let node = capture.node;
145
146            if let Some((prefix, suffix)) = capture_name.rsplit_once('.') {
147                match suffix {
148                    "type" => {
149                        scope_type = Some(prefix.to_string());
150                        scope_start = Some(node.start_position());
151                        scope_end = Some(node.end_position());
152                    }
153                    "name" => {
154                        scope_name = node
155                            .utf8_text(content)
156                            .ok()
157                            .map(std::string::ToString::to_string);
158                    }
159                    _ => {}
160                }
161            }
162        }
163
164        Some(ScopeCapture {
165            scope_type: scope_type?,
166            scope_name: scope_name?,
167            start: scope_start?,
168            end: scope_end?,
169        })
170    }
171
172    fn build_scope_from_capture(capture: ScopeCapture, file_path: &Path) -> Scope {
173        Scope {
174            id: ScopeId::new(0),
175            scope_type: capture.scope_type,
176            name: capture.scope_name,
177            file_path: file_path.to_path_buf(),
178            start_line: capture.start.row + 1,
179            start_column: capture.start.column,
180            end_line: capture.end.row + 1,
181            end_column: capture.end.column,
182            parent_id: None,
183        }
184    }
185
186    fn scope_query_source() -> &'static str {
187        r"
188; Function scopes
189(function_definition
190  name: (identifier) @function.name
191) @function.type
192
193; Class scopes
194(class_definition
195  name: (identifier) @class.name
196) @class.type
197"
198    }
199}
sqry_lang_python/lib.rs

sqry_lang_python/
lib.rs
