Skip to main content

codesynapse_core/ts_extract/
powershell.rs

1use super::make_file_node;
2use crate::error::{CodeSynapseError, Result};
3use crate::extract::{make_id, ImportNode, LanguageExtractor};
4use crate::types::{Edge, ExtractionFragment, Node};
5use std::collections::HashMap;
6use std::path::Path;
7use tree_sitter::{Node as TsNode, Parser};
8
9pub struct PowerShellExtractor;
10
11fn ps_text(source: &[u8], node: &TsNode<'_>) -> String {
12    std::str::from_utf8(&source[node.start_byte()..node.end_byte()])
13        .unwrap_or("")
14        .to_string()
15}
16
17fn ps_node(id: String, label: String, file_type: &str, path: &Path) -> Node {
18    Node {
19        id,
20        label,
21        file_type: file_type.to_string(),
22        source_file: path.to_string_lossy().to_string(),
23        source_location: None,
24        community: None,
25        rationale: None,
26        docstring: None,
27        metadata: HashMap::new(),
28    }
29}
30
31fn ps_edge(fragment: &mut ExtractionFragment, src: &str, tgt: &str, rel: &str, path: &Path) {
32    fragment.edges.push(Edge {
33        source: src.to_string(),
34        target: tgt.to_string(),
35        relation: rel.to_string(),
36        confidence: "EXTRACTED".to_string(),
37        source_file: Some(path.to_string_lossy().to_string()),
38        weight: 1.0,
39        context: None,
40    });
41}
42
43fn walk_ps<'t>(
44    node: TsNode<'t>,
45    source: &[u8],
46    file_id: &str,
47    stem: &str,
48    path: &Path,
49    fragment: &mut ExtractionFragment,
50    parent_class: Option<String>,
51) {
52    match node.kind() {
53        "function_statement" => {
54            for i in 0..node.child_count() {
55                if let Some(child) = node.child(i) {
56                    if child.kind() == "function_name" {
57                        let name = ps_text(source, &child);
58                        if !name.is_empty() {
59                            let id = make_id(&[stem, &name]);
60                            fragment.nodes.push(ps_node(
61                                id.clone(),
62                                format!("{}()", name),
63                                "function",
64                                path,
65                            ));
66                            ps_edge(fragment, file_id, &id, "contains", path);
67                        }
68                        break;
69                    }
70                }
71            }
72        }
73        "class_statement" => {
74            let mut class_nid: Option<String> = None;
75            for i in 0..node.child_count() {
76                if let Some(child) = node.child(i) {
77                    if child.kind() == "simple_name" {
78                        let name = ps_text(source, &child);
79                        if !name.is_empty() {
80                            let id = make_id(&[stem, &name]);
81                            fragment
82                                .nodes
83                                .push(ps_node(id.clone(), name, "class", path));
84                            ps_edge(fragment, file_id, &id, "contains", path);
85                            class_nid = Some(id);
86                            break;
87                        }
88                    }
89                }
90            }
91            for i in 0..node.child_count() {
92                if let Some(child) = node.child(i) {
93                    walk_ps(
94                        child,
95                        source,
96                        file_id,
97                        stem,
98                        path,
99                        fragment,
100                        class_nid.clone(),
101                    );
102                }
103            }
104        }
105        "class_method_definition" => {
106            for i in 0..node.child_count() {
107                if let Some(child) = node.child(i) {
108                    if child.kind() == "simple_name" {
109                        let name = ps_text(source, &child);
110                        if !name.is_empty() {
111                            let parent = parent_class.as_deref().unwrap_or(file_id);
112                            let id = make_id(&[parent, &name]);
113                            fragment.nodes.push(ps_node(
114                                id.clone(),
115                                format!(".{}()", name),
116                                "function",
117                                path,
118                            ));
119                            ps_edge(fragment, parent, &id, "method", path);
120                        }
121                        break;
122                    }
123                }
124            }
125        }
126        _ => {
127            for i in 0..node.child_count() {
128                if let Some(child) = node.child(i) {
129                    walk_ps(
130                        child,
131                        source,
132                        file_id,
133                        stem,
134                        path,
135                        fragment,
136                        parent_class.clone(),
137                    );
138                }
139            }
140        }
141    }
142}
143
144impl PowerShellExtractor {
145    pub fn extract(source: &[u8], path: &Path) -> Result<ExtractionFragment> {
146        let (file_id, _, file_node) = make_file_node(path);
147        let stem = file_id.clone();
148        let mut fragment = ExtractionFragment {
149            nodes: vec![file_node],
150            edges: vec![],
151        };
152
153        let lang: tree_sitter::Language = tree_sitter_powershell::LANGUAGE.into();
154        let mut parser = Parser::new();
155        parser
156            .set_language(&lang)
157            .map_err(|e| CodeSynapseError::Parse(format!("ps1 set_language: {e}")))?;
158        let tree = parser
159            .parse(source, None)
160            .ok_or_else(|| CodeSynapseError::Parse("ps1 parse failed".to_string()))?;
161
162        walk_ps(
163            tree.root_node(),
164            source,
165            &file_id,
166            &stem,
167            path,
168            &mut fragment,
169            None,
170        );
171
172        Ok(fragment)
173    }
174}
175
176impl LanguageExtractor for PowerShellExtractor {
177    fn file_extensions(&self) -> Vec<&'static str> {
178        vec!["ps1", "psm1", "psd1"]
179    }
180    fn extract(&self, source: &[u8], path: &Path) -> Result<ExtractionFragment> {
181        Self::extract(source, path)
182    }
183    fn resolve_imports(&self, _imports: &[ImportNode]) -> Vec<Edge> {
184        vec![]
185    }
186    fn collect_type_refs(&self, _fragment: &mut ExtractionFragment) {}
187}