capsule_core/wasm/compiler/
python.rs

1use std::collections::HashMap;
2use std::fmt;
3use std::fs;
4use std::path::{Path, PathBuf};
5use std::process::Command;
6use std::process::Stdio;
7
8use crate::config::fingerprint::SourceFingerprint;
9use crate::wasm::utilities::introspection::scanner;
10use crate::wasm::utilities::wit_manager::WitManager;
11
12pub enum PythonWasmCompilerError {
13    CompileFailed(String),
14    FsError(String),
15}
16
17impl fmt::Display for PythonWasmCompilerError {
18    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
19        match self {
20            PythonWasmCompilerError::CompileFailed(msg) => {
21                write!(f, "Compilation failed > {}", msg)
22            }
23            PythonWasmCompilerError::FsError(msg) => write!(f, "File system error > {}", msg),
24        }
25    }
26}
27
28impl From<std::io::Error> for PythonWasmCompilerError {
29    fn from(err: std::io::Error) -> Self {
30        PythonWasmCompilerError::CompileFailed(err.to_string())
31    }
32}
33
34impl From<std::time::SystemTimeError> for PythonWasmCompilerError {
35    fn from(err: std::time::SystemTimeError) -> Self {
36        PythonWasmCompilerError::FsError(err.to_string())
37    }
38}
39
40pub struct PythonWasmCompiler {
41    pub source_path: PathBuf,
42    pub cache_dir: PathBuf,
43    pub output_wasm: PathBuf,
44}
45
46impl PythonWasmCompiler {
47    pub fn new(source_path: &Path) -> Result<Self, PythonWasmCompilerError> {
48        let source_path = source_path.canonicalize().map_err(|e| {
49            PythonWasmCompilerError::FsError(format!("Cannot resolve source path: {}", e))
50        })?;
51
52        let cache_dir = std::env::current_dir()
53            .map_err(|e| {
54                PythonWasmCompilerError::FsError(format!("Cannot get current directory: {}", e))
55            })?
56            .join(".capsule");
57
58        let output_wasm = cache_dir.join("capsule.wasm");
59
60        if !cache_dir.exists() {
61            fs::create_dir_all(&cache_dir)?;
62        }
63
64        Ok(Self {
65            source_path,
66            cache_dir,
67            output_wasm,
68        })
69    }
70
71    fn python_command() -> &'static str {
72        if Command::new("python")
73            .arg("--version")
74            .stdout(Stdio::null())
75            .stderr(Stdio::null())
76            .status()
77            .is_ok()
78        {
79            "python"
80        } else {
81            "python3"
82        }
83    }
84
85    fn normalize_path(path: &Path) -> PathBuf {
86        let path_str = path.to_string_lossy();
87        if let Some(stripped) = path_str.strip_prefix(r"\\?\") {
88            return PathBuf::from(stripped);
89        }
90        path.to_path_buf()
91    }
92
93    pub fn compile_wasm(&self) -> Result<PathBuf, PythonWasmCompilerError> {
94        let source_dir = self.source_path.parent().ok_or_else(|| {
95            PythonWasmCompilerError::FsError("Cannot determine source directory".to_string())
96        })?;
97
98        if !SourceFingerprint::needs_rebuild(
99            &self.cache_dir,
100            source_dir,
101            &self.output_wasm,
102            &["py", "toml"],
103            &["__pycache__"],
104        ) {
105            return Ok(self.output_wasm.clone());
106        }
107
108        let module_name = self
109            .source_path
110            .file_stem()
111            .ok_or(PythonWasmCompilerError::FsError(
112                "Invalid source file name".to_string(),
113            ))?
114            .to_str()
115            .ok_or(PythonWasmCompilerError::FsError(
116                "Invalid UTF-8 in file name".to_string(),
117            ))?;
118
119        let python_path = self
120            .source_path
121            .parent()
122            .ok_or(PythonWasmCompilerError::FsError(
123                "Cannot determine parent directory".to_string(),
124            ))?;
125
126        let wit_path = self.get_wit_path()?;
127        let sdk_path = self.get_sdk_path()?;
128
129        if !sdk_path.exists() {
130            return Err(PythonWasmCompilerError::FsError(format!(
131                "SDK directory not found: {}",
132                sdk_path.display()
133            )));
134        }
135
136        if !sdk_path.exists() {
137            return Err(PythonWasmCompilerError::FsError(format!(
138                "SDK directory not found: {}",
139                sdk_path.display()
140            )));
141        }
142
143        let bootloader_path = self.cache_dir.join("_capsule_boot.py");
144        let bootloader_content = format!(
145            r#"# Auto-generated bootloader for Capsule
146import {module_name}
147import capsule.app
148capsule.app._main_module = {module_name}
149from capsule.app import TaskRunner, exports
150"#,
151            module_name = module_name
152        );
153
154        fs::write(&bootloader_path, bootloader_content)?;
155
156        let wit_path_normalized = Self::normalize_path(&wit_path);
157        let cache_dir_normalized = Self::normalize_path(&self.cache_dir);
158        let python_path_normalized = Self::normalize_path(python_path);
159        let sdk_path_normalized = Self::normalize_path(&sdk_path);
160        let output_wasm_normalized = Self::normalize_path(&self.output_wasm);
161
162        let output = Command::new("componentize-py")
163            .arg("-d")
164            .arg(&wit_path_normalized)
165            .arg("-w")
166            .arg("capsule-agent")
167            .arg("componentize")
168            .arg("_capsule_boot")
169            .arg("-p")
170            .arg(&cache_dir_normalized)
171            .arg("-p")
172            .arg(&python_path_normalized)
173            .arg("-p")
174            .arg(&sdk_path_normalized)
175            .arg("-o")
176            .arg(&output_wasm_normalized)
177            .stdout(Stdio::piped())
178            .stderr(Stdio::piped())
179            .output()?;
180
181        if !output.status.success() {
182            return Err(PythonWasmCompilerError::CompileFailed(format!(
183                "Compilation failed: {}",
184                String::from_utf8_lossy(&output.stderr).trim()
185            )));
186        }
187
188        self.cleanup_pycache(python_path);
189
190        let _ = SourceFingerprint::update_after_build(
191            &self.cache_dir,
192            source_dir,
193            &["py", "toml"],
194            &["__pycache__"],
195        );
196
197        Ok(self.output_wasm.clone())
198    }
199
200    fn get_wit_path(&self) -> Result<PathBuf, PythonWasmCompilerError> {
201        let wit_dir = self.cache_dir.join("wit");
202
203        if !wit_dir.join("capsule.wit").exists() {
204            WitManager::import_wit_deps(&wit_dir)?;
205        }
206
207        Ok(wit_dir)
208    }
209
210    fn get_sdk_path(&self) -> Result<PathBuf, PythonWasmCompilerError> {
211        if let Ok(sdk_path) = self.find_python_sdk_path() {
212            return Ok(sdk_path);
213        }
214
215        if let Ok(exe_path) = std::env::current_exe()
216            && let Some(project_root) = exe_path
217                .parent()
218                .and_then(|p| p.parent())
219                .and_then(|p| p.parent())
220                .and_then(|p| p.parent())
221                .and_then(|p| p.parent())
222        {
223            let sdk_path = project_root.join("crates/capsule-sdk/python/src");
224            if sdk_path.exists() {
225                return Ok(sdk_path);
226            }
227        }
228
229        Err(PythonWasmCompilerError::FsError(
230            "Cannot find SDK. Make sure to install capsule package.".to_string(),
231        ))
232    }
233
234    fn find_python_sdk_path(&self) -> Result<PathBuf, PythonWasmCompilerError> {
235        let python_cmd = Self::python_command();
236        let output = Command::new(python_cmd)
237            .arg("-c")
238            .arg("import capsule; import os; print(os.path.dirname(os.path.dirname(capsule.__file__)), end='')")
239            .output()
240            .map_err(|e| {
241                PythonWasmCompilerError::FsError(format!(
242                    "Failed to execute '{}': {}. Make sure Python is installed and 'pip install capsule-run' was run.",
243                    python_cmd, e
244                ))
245            })?;
246
247        if !output.status.success() {
248            let stderr = String::from_utf8_lossy(&output.stderr);
249            return Err(PythonWasmCompilerError::FsError(format!(
250                "Cannot find 'capsule' module. Run 'pip install capsule-run' first. Python error: {}",
251                stderr.trim()
252            )));
253        }
254
255        let path_str = String::from_utf8_lossy(&output.stdout).trim().to_string();
256
257        if path_str.is_empty() {
258            return Err(PythonWasmCompilerError::FsError(
259                "Python returned empty path for capsule package".to_string(),
260            ));
261        }
262
263        let sdk_path = PathBuf::from(&path_str);
264
265        if !sdk_path.exists() {
266            return Err(PythonWasmCompilerError::FsError(format!(
267                "SDK path from Python does not exist: {}",
268                sdk_path.display()
269            )));
270        }
271
272        Ok(sdk_path)
273    }
274
275    fn cleanup_pycache(&self, source_dir: &Path) {
276        Self::remove_pycache_recursive(source_dir);
277        Self::remove_pycache_recursive(&self.cache_dir);
278    }
279
280    fn remove_pycache_recursive(dir: &Path) {
281        if let Ok(entries) = fs::read_dir(dir) {
282            for entry in entries.flatten() {
283                let path = entry.path();
284                if path.is_dir() {
285                    if path.file_name().is_some_and(|n| n == "__pycache__") {
286                        let _ = fs::remove_dir_all(&path);
287                    } else {
288                        Self::remove_pycache_recursive(&path);
289                    }
290                }
291            }
292        }
293    }
294
295    pub fn introspect_task_registry(&self) -> Option<HashMap<String, serde_json::Value>> {
296        let source_dir = self.source_path.parent()?;
297        scanner::scan_python_tasks(source_dir)
298    }
299}
capsule_core/wasm/compiler/python.rs

capsule_core/wasm/compiler/
python.rs
