srgn/scoping/langs/
python.rs

1use std::fmt::Debug;
2
3use clap::ValueEnum;
4use const_format::formatcp;
5
6use super::{Find, LanguageScoper, QuerySource, TSLanguage, TSQuery, TSQueryError};
7use crate::scoping::langs::IGNORE;
8
9/// A compiled query for the Python language.
10#[derive(Debug)]
11pub struct CompiledQuery(super::CompiledQuery);
12
13impl TryFrom<QuerySource> for CompiledQuery {
14    type Error = TSQueryError;
15
16    /// Create a new compiled query for the Python language.
17    ///
18    /// # Errors
19    ///
20    /// See the concrete type of the [`TSQueryError`](tree_sitter::QueryError)variant for when this method errors.
21    fn try_from(query: QuerySource) -> Result<Self, Self::Error> {
22        let q = super::CompiledQuery::from_source(&tree_sitter_python::LANGUAGE.into(), &query)?;
23        Ok(Self(q))
24    }
25}
26
27impl From<PreparedQuery> for CompiledQuery {
28    fn from(query: PreparedQuery) -> Self {
29        Self(super::CompiledQuery::from_prepared_query(
30            &tree_sitter_python::LANGUAGE.into(),
31            query.as_str(),
32        ))
33    }
34}
35
36/// Prepared tree-sitter queries for Python.
37#[derive(Debug, Clone, Copy, ValueEnum)]
38pub enum PreparedQuery {
39    /// Comments.
40    Comments,
41    /// Strings (raw, byte, f-strings; interpolation not included).
42    Strings,
43    /// Module names in imports (incl. periods; excl. `import`/`from`/`as`/`*`).
44    Imports,
45    /// Docstrings (not including multi-line strings).
46    DocStrings,
47    /// Function names, at the definition site.
48    FunctionNames,
49    /// Function calls.
50    FunctionCalls,
51    /// Class definitions (in their entirety).
52    Class,
53    /// Function definitions (*all* `def` block in their entirety).
54    Def,
55    /// Async function definitions (*all* `async def` block in their entirety).
56    AsyncDef,
57    /// Function definitions inside `class` bodies.
58    Methods,
59    /// Function definitions decorated as `classmethod` (excl. the decorator).
60    ClassMethods,
61    /// Function definitions decorated as `staticmethod` (excl. the decorator).
62    StaticMethods,
63    /// `with` blocks (in their entirety).
64    With,
65    /// `try` blocks (in their entirety).
66    Try,
67    /// `lambda` statements (in their entirety).
68    Lambda,
69    /// Global, i.e. module-level variables.
70    Globals,
71    /// Identifiers for variables (left-hand side of assignments).
72    VariableIdentifiers,
73    /// Types in type hints.
74    Types,
75    /// Identifiers (variable names, ...).
76    Identifiers,
77}
78
79impl PreparedQuery {
80    #[expect(clippy::too_many_lines)]
81    const fn as_str(self) -> &'static str {
82        match self {
83            Self::Comments => "(comment) @comment",
84            Self::Strings => "(string_content) @string",
85            Self::Imports => {
86                r"[
87                    (import_statement
88                            name: (dotted_name) @dn)
89                    (import_from_statement
90                            module_name: (dotted_name) @dn)
91                    (import_from_statement
92                            module_name: (dotted_name) @dn
93                                (wildcard_import))
94                    (import_statement(
95                        aliased_import
96                            name: (dotted_name) @dn))
97                    (import_from_statement
98                        module_name: (relative_import) @ri)
99                ]"
100            }
101            Self::DocStrings => {
102                // Triple-quotes are also used for multi-line strings. So look only
103                // for stand-alone expressions, which are not part of some variable
104                // assignment.
105                formatcp!(
106                    "
107                    (
108                        (expression_statement
109                            (string
110                                (string_start) @{0}
111                                (string_content) @string
112                                (#match? @{0} \"\\^\\\"\\\"\\\"\")
113                            )
114                        )
115                    )
116                    ",
117                    IGNORE
118                )
119            }
120            Self::FunctionNames => {
121                r"
122                (function_definition
123                    name: (identifier) @function-name
124                )
125                "
126            }
127            Self::FunctionCalls => {
128                r"
129                (call
130                    function: (identifier) @function-name
131                )
132                "
133            }
134            Self::Class => "(class_definition) @class",
135            Self::Def => "(function_definition) @def",
136            Self::AsyncDef => r#"((function_definition) @def (#match? @def "^async "))"#,
137            Self::Methods => {
138                r"
139                (class_definition
140                    body: (block
141                        [
142                            (function_definition) @method
143                            (decorated_definition definition: (function_definition)) @method
144                        ]
145                    )
146                )
147                "
148            }
149            Self::ClassMethods => {
150                formatcp!(
151                    "
152                    (class_definition
153                        body: (block
154                            (decorated_definition
155                                (decorator (identifier) @{0})
156                                definition: (function_definition) @method
157                                (#eq? @{0} \"classmethod\")
158                            )
159                        )
160                    )",
161                    IGNORE
162                )
163            }
164            Self::StaticMethods => {
165                formatcp!(
166                    "
167                    (class_definition
168                        body: (block
169                            (decorated_definition
170                                (decorator (identifier) @{0})
171                                definition: (function_definition) @method
172                                (#eq? @{0} \"staticmethod\")
173                            )
174                        )
175                    )",
176                    IGNORE
177                )
178            }
179            Self::With => "(with_statement) @with",
180            Self::Try => "(try_statement) @try",
181            Self::Lambda => "(lambda) @lambda",
182            Self::Globals => {
183                "(module (expression_statement (assignment left: (identifier) @global)))"
184            }
185            Self::VariableIdentifiers => "(assignment left: (identifier) @identifier)",
186            Self::Types => "(type) @type",
187            Self::Identifiers => "(identifier) @identifier",
188        }
189    }
190}
191
192impl LanguageScoper for CompiledQuery {
193    fn lang() -> TSLanguage {
194        tree_sitter_python::LANGUAGE.into()
195    }
196
197    fn pos_query(&self) -> &TSQuery {
198        &self.0.positive_query
199    }
200
201    fn neg_query(&self) -> Option<&TSQuery> {
202        self.0.negative_query.as_ref()
203    }
204}
205
206impl Find for CompiledQuery {
207    fn extensions(&self) -> &'static [&'static str] {
208        &["py"]
209    }
210
211    fn interpreters(&self) -> Option<&'static [&'static str]> {
212        Some(&["python", "python3"])
213    }
214}
srgn/scoping/langs/python.rs

srgn/scoping/langs/
python.rs
