codesynapse_core/ts_extract/
python.rs

1use super::{
2    add_contains_edge, add_node_if_missing, make_file_node, run_query_matches_ranged,
3    run_query_named, strip_docstring,
4};
5use crate::error::Result;
6use crate::extract::{make_id, ImportNode, LanguageExtractor};
7use crate::types::{Edge, ExtractionFragment, Node};
8use std::collections::HashMap;
9use std::path::Path;
10
11/// Tree-sitter based Python extractor
12pub struct TsPythonExtractor;
13
14impl TsPythonExtractor {
15    pub fn extract(source: &[u8], path: &Path) -> Result<ExtractionFragment> {
16        let (file_id, _, file_node) = make_file_node(path);
17        let mut fragment = ExtractionFragment {
18            nodes: vec![file_node],
19            edges: vec![],
20        };
21
22        let lang = tree_sitter_python::LANGUAGE.into();
23
24        // Class definitions with bases and docstrings
25        let class_query = r#"
26            (class_definition
27                name: (identifier) @class.name
28                superclasses: (argument_list
29                    (identifier) @class.base
30                )?
31                body: (block .
32                    (expression_statement (string) @class.docstring)
33                )?
34            ) @class.node
35        "#;
36        if let Ok(matches) = run_query_matches_ranged(source, &lang, class_query) {
37            for (m, ranges) in &matches {
38                let name = match m.get("class.name") {
39                    Some(n) => n.clone(),
40                    None => continue,
41                };
42                let class_id = make_id(&[&file_id, &name]);
43                let docstring = m.get("class.docstring").and_then(|s| strip_docstring(s));
44                let source_location = ranges
45                    .get("class.node")
46                    .map(|(start, end)| format!("{}:{}", start, end));
47                let class_node = Node {
48                    id: class_id.clone(),
49                    label: name.clone(),
50                    file_type: "class".to_string(),
51                    source_file: path.to_string_lossy().to_string(),
52                    source_location,
53                    community: None,
54                    rationale: None,
55                    docstring,
56                    metadata: HashMap::new(),
57                };
58                add_node_if_missing(&mut fragment, class_node);
59                add_contains_edge(&mut fragment, &file_id, class_id.clone(), path);
60
61                if let Some(base_name) = m.get("class.base") {
62                    let base_id = make_id(&[base_name]);
63                    let base_node_ = Node {
64                        id: base_id.clone(),
65                        label: base_name.to_string(),
66                        file_type: "code".to_string(),
67                        source_file: path.to_string_lossy().to_string(),
68                        source_location: None,
69                        community: None,
70                        rationale: None,
71                        docstring: None,
72                        metadata: HashMap::new(),
73                    };
74                    add_node_if_missing(&mut fragment, base_node_);
75                    fragment.edges.push(Edge {
76                        source: class_id.clone(),
77                        target: base_id,
78                        relation: "inherits".to_string(),
79                        confidence: "EXTRACTED".to_string(),
80                        source_file: Some(path.to_string_lossy().to_string()),
81                        weight: 1.0,
82                        context: None,
83                    });
84                }
85            }
86        }
87
88        // Methods inside class bodies — must run before the general fn_query so that
89        // add_node_if_missing skips them when the general query runs.
90        let method_query = r#"
91            (class_definition
92                body: (block
93                    (function_definition
94                        name: (identifier) @method.name
95                        body: (block .
96                            (expression_statement (string) @method.docstring)
97                        )?
98                    ) @method.node
99                )
100            )
101        "#;
102        if let Ok(matches) = run_query_matches_ranged(source, &lang, method_query) {
103            for (m, ranges) in &matches {
104                let name = match m.get("method.name") {
105                    Some(n) => n.clone(),
106                    None => continue,
107                };
108                let fn_id = make_id(&[&file_id, &name, "()"]);
109                let docstring = m.get("method.docstring").and_then(|s| strip_docstring(s));
110                let source_location = ranges
111                    .get("method.node")
112                    .map(|(start, end)| format!("{}:{}", start, end));
113                let method_node = Node {
114                    id: fn_id.clone(),
115                    label: format!("{}()", name),
116                    file_type: "method".to_string(),
117                    source_file: path.to_string_lossy().to_string(),
118                    source_location,
119                    community: None,
120                    rationale: None,
121                    docstring,
122                    metadata: HashMap::new(),
123                };
124                add_node_if_missing(&mut fragment, method_node);
125                add_contains_edge(&mut fragment, &file_id, fn_id, path);
126            }
127        }
128
129        // Module-level function definitions with docstrings.
130        // Methods already added above are skipped by add_node_if_missing.
131        let fn_query = r#"
132            (function_definition
133                name: (identifier) @func.name
134                body: (block .
135                    (expression_statement (string) @func.docstring)
136                )?
137            ) @func.node
138        "#;
139        if let Ok(matches) = run_query_matches_ranged(source, &lang, fn_query) {
140            for (m, ranges) in &matches {
141                let name = match m.get("func.name") {
142                    Some(n) => n.clone(),
143                    None => continue,
144                };
145                let fn_id = make_id(&[&file_id, &name, "()"]);
146                let docstring = m.get("func.docstring").and_then(|s| strip_docstring(s));
147                let source_location = ranges
148                    .get("func.node")
149                    .map(|(start, end)| format!("{}:{}", start, end));
150                let fn_node = Node {
151                    id: fn_id.clone(),
152                    label: format!("{}()", name),
153                    file_type: "function".to_string(),
154                    source_file: path.to_string_lossy().to_string(),
155                    source_location,
156                    community: None,
157                    rationale: None,
158                    docstring,
159                    metadata: HashMap::new(),
160                };
161                add_node_if_missing(&mut fragment, fn_node);
162                add_contains_edge(&mut fragment, &file_id, fn_id, path);
163            }
164        }
165
166        // Import: import x, import x as y
167        let import_query = r#"
168            (import_statement
169                name: (dotted_name) @import.name
170                alias: (identifier)? @import.alias
171            )
172        "#;
173        if let Ok(mut captures) = run_query_named(source, &lang, import_query) {
174            let names = captures.remove("import.name").unwrap_or_default();
175            let aliases = captures.remove("import.alias").unwrap_or_default();
176            for name in &names {
177                let parts: Vec<&str> = name.split('.').collect();
178                let simple = parts.last().copied().unwrap_or(name.as_str());
179                let mod_id = make_id(&[simple]);
180                let mod_node = Node {
181                    id: mod_id.clone(),
182                    label: simple.to_string(),
183                    file_type: "code".to_string(),
184                    source_file: path.to_string_lossy().to_string(),
185                    source_location: None,
186                    community: None,
187                    rationale: None,
188                    docstring: None,
189                    metadata: HashMap::new(),
190                };
191                add_node_if_missing(&mut fragment, mod_node);
192                fragment.edges.push(Edge {
193                    source: file_id.clone(),
194                    target: mod_id,
195                    relation: "imports".to_string(),
196                    confidence: "EXTRACTED".to_string(),
197                    source_file: Some(path.to_string_lossy().to_string()),
198                    weight: 1.0,
199                    context: None,
200                });
201            }
202            for alias in &aliases {
203                let alias_id = make_id(&[&file_id, alias, "alias"]);
204                let alias_node = Node {
205                    id: alias_id.clone(),
206                    label: alias.to_string(),
207                    file_type: "code".to_string(),
208                    source_file: path.to_string_lossy().to_string(),
209                    source_location: None,
210                    community: None,
211                    rationale: None,
212                    docstring: None,
213                    metadata: HashMap::new(),
214                };
215                add_node_if_missing(&mut fragment, alias_node);
216                fragment.edges.push(Edge {
217                    source: file_id.clone(),
218                    target: alias_id,
219                    relation: "imports".to_string(),
220                    confidence: "EXTRACTED".to_string(),
221                    source_file: Some(path.to_string_lossy().to_string()),
222                    weight: 1.0,
223                    context: None,
224                });
225            }
226        }
227
228        // From-import: from x import y, from x import y as z
229        let from_query = r#"
230            (import_from_statement
231                module_name: [
232                    (dotted_name)
233                    (relative_import)
234                ] @from.module
235                name: (dotted_name) @from.name
236                alias: (identifier)? @from.alias
237            )
238        "#;
239        if let Ok(mut captures) = run_query_named(source, &lang, from_query) {
240            let modules = captures.remove("from.module").unwrap_or_default();
241            let names = captures.remove("from.name").unwrap_or_default();
242            let from_aliases = captures.remove("from.alias").unwrap_or_default();
243
244            for mod_name in &modules {
245                let mod_id = make_id(&[&file_id, mod_name]);
246                let mod_node = Node {
247                    id: mod_id.clone(),
248                    label: mod_name.to_string(),
249                    file_type: "code".to_string(),
250                    source_file: path.to_string_lossy().to_string(),
251                    source_location: None,
252                    community: None,
253                    rationale: None,
254                    docstring: None,
255                    metadata: HashMap::new(),
256                };
257                add_node_if_missing(&mut fragment, mod_node);
258                fragment.edges.push(Edge {
259                    source: file_id.clone(),
260                    target: mod_id,
261                    relation: "imports".to_string(),
262                    confidence: "EXTRACTED".to_string(),
263                    source_file: Some(path.to_string_lossy().to_string()),
264                    weight: 1.0,
265                    context: None,
266                });
267            }
268            for name in &names {
269                let import_id = make_id(&[&file_id, name]);
270                let import_node = Node {
271                    id: import_id.clone(),
272                    label: name.to_string(),
273                    file_type: "code".to_string(),
274                    source_file: path.to_string_lossy().to_string(),
275                    source_location: None,
276                    community: None,
277                    rationale: None,
278                    docstring: None,
279                    metadata: HashMap::new(),
280                };
281                add_node_if_missing(&mut fragment, import_node);
282                fragment.edges.push(Edge {
283                    source: file_id.clone(),
284                    target: import_id,
285                    relation: "imports".to_string(),
286                    confidence: "EXTRACTED".to_string(),
287                    source_file: Some(path.to_string_lossy().to_string()),
288                    weight: 1.0,
289                    context: None,
290                });
291            }
292            for alias in &from_aliases {
293                let alias_id = make_id(&[&file_id, alias, "alias"]);
294                let alias_node = Node {
295                    id: alias_id.clone(),
296                    label: alias.to_string(),
297                    file_type: "code".to_string(),
298                    source_file: path.to_string_lossy().to_string(),
299                    source_location: None,
300                    community: None,
301                    rationale: None,
302                    docstring: None,
303                    metadata: HashMap::new(),
304                };
305                add_node_if_missing(&mut fragment, alias_node);
306                fragment.edges.push(Edge {
307                    source: file_id.clone(),
308                    target: alias_id,
309                    relation: "imports".to_string(),
310                    confidence: "EXTRACTED".to_string(),
311                    source_file: Some(path.to_string_lossy().to_string()),
312                    weight: 1.0,
313                    context: None,
314                });
315            }
316        }
317
318        Ok(fragment)
319    }
320}
321
322/// Tree-sitter based JavaScript/TypeScript extractor
323impl LanguageExtractor for TsPythonExtractor {
324    fn file_extensions(&self) -> Vec<&'static str> {
325        vec!["py"]
326    }
327    fn extract(&self, source: &[u8], path: &Path) -> Result<ExtractionFragment> {
328        Self::extract(source, path)
329    }
330    fn resolve_imports(&self, _imports: &[ImportNode]) -> Vec<Edge> {
331        vec![]
332    }
333    fn collect_type_refs(&self, _fragment: &mut ExtractionFragment) {}
334}
codesynapse_core/ts_extract/python.rs

codesynapse_core/ts_extract/
python.rs
