Skip to main content

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_for_command(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
128        let sdk_path = self.get_sdk_path()?;
129
130        if !sdk_path.exists() {
131            return Err(PythonWasmCompilerError::FsError(format!(
132                "SDK directory not found: {}",
133                sdk_path.display()
134            )));
135        }
136
137        if !sdk_path.exists() {
138            return Err(PythonWasmCompilerError::FsError(format!(
139                "SDK directory not found: {}",
140                sdk_path.display()
141            )));
142        }
143
144        let bootloader_path = self.cache_dir.join("_capsule_boot.py");
145        let bootloader_content = format!(
146            r#"# Auto-generated bootloader for Capsule
147import {module_name}
148import capsule.app
149capsule.app._main_module = {module_name}
150from capsule.app import TaskRunner, exports
151"#,
152            module_name = module_name
153        );
154
155        fs::write(&bootloader_path, bootloader_content)?;
156
157        let wit_path_normalized = Self::normalize_path_for_command(&wit_path);
158        let cache_dir_normalized = Self::normalize_path_for_command(&self.cache_dir);
159        let python_path_normalized = Self::normalize_path_for_command(python_path);
160        let sdk_path_normalized = Self::normalize_path_for_command(&sdk_path);
161        let output_wasm_normalized = Self::normalize_path_for_command(&self.output_wasm);
162
163        let output = Command::new("componentize-py")
164            .arg("-d")
165            .arg(&wit_path_normalized)
166            .arg("-w")
167            .arg("capsule-agent")
168            .arg("componentize")
169            .arg("_capsule_boot")
170            .arg("-p")
171            .arg(&cache_dir_normalized)
172            .arg("-p")
173            .arg(&python_path_normalized)
174            .arg("-p")
175            .arg(&sdk_path_normalized)
176            .arg("-o")
177            .arg(&output_wasm_normalized)
178            .stdout(Stdio::piped())
179            .stderr(Stdio::piped())
180            .output()?;
181
182        if !output.status.success() {
183            return Err(PythonWasmCompilerError::CompileFailed(format!(
184                "Compilation failed: {}",
185                String::from_utf8_lossy(&output.stderr).trim()
186            )));
187        }
188
189        self.cleanup_pycache(python_path);
190
191        let _ = SourceFingerprint::update_after_build(
192            &self.cache_dir,
193            source_dir,
194            &["py", "toml"],
195            &["__pycache__"],
196        );
197
198        Ok(self.output_wasm.clone())
199    }
200
201    fn get_wit_path(&self) -> Result<PathBuf, PythonWasmCompilerError> {
202        if let Ok(path) = std::env::var("CAPSULE_WIT_PATH") {
203            let wit_path = PathBuf::from(path);
204            if wit_path.exists() {
205                return Ok(wit_path);
206            }
207        }
208
209        let wit_dir = self.cache_dir.join("wit");
210
211        if !wit_dir.join("capsule.wit").exists() {
212            WitManager::import_wit_deps(&wit_dir)?;
213        }
214
215        Ok(wit_dir)
216    }
217
218    fn get_sdk_path(&self) -> Result<PathBuf, PythonWasmCompilerError> {
219        if let Ok(path) = std::env::var("CAPSULE_SDK_PATH") {
220            let sdk_path = PathBuf::from(path);
221            if sdk_path.exists() {
222                return Ok(sdk_path);
223            }
224        }
225
226        if let Ok(sdk_path) = self.find_sdk_via_python() {
227            return Ok(sdk_path);
228        }
229
230        if let Ok(exe_path) = std::env::current_exe()
231            && let Some(project_root) = exe_path
232                .parent()
233                .and_then(|p| p.parent())
234                .and_then(|p| p.parent())
235                .and_then(|p| p.parent())
236                .and_then(|p| p.parent())
237        {
238            let sdk_path = project_root.join("crates/capsule-sdk/python/src");
239            if sdk_path.exists() {
240                return Ok(sdk_path);
241            }
242        }
243
244        Err(PythonWasmCompilerError::FsError(
245            "Cannot find SDK. Make sure to install capsule package.".to_string(),
246        ))
247    }
248
249    fn find_sdk_via_python(&self) -> Result<PathBuf, PythonWasmCompilerError> {
250        let python_cmd = Self::python_command();
251        let output = Command::new(python_cmd)
252            .arg("-c")
253            .arg("import capsule; import os; print(os.path.dirname(os.path.dirname(capsule.__file__)), end='')")
254            .output()
255            .map_err(|e| {
256                PythonWasmCompilerError::FsError(format!(
257                    "Failed to execute '{}': {}. Make sure Python is installed and 'pip install capsule-run' was run.",
258                    python_cmd, e
259                ))
260            })?;
261
262        if !output.status.success() {
263            let stderr = String::from_utf8_lossy(&output.stderr);
264            return Err(PythonWasmCompilerError::FsError(format!(
265                "Cannot find 'capsule' module. Run 'pip install capsule-run' first. Python error: {}",
266                stderr.trim()
267            )));
268        }
269
270        let path_str = String::from_utf8_lossy(&output.stdout).trim().to_string();
271
272        if path_str.is_empty() {
273            return Err(PythonWasmCompilerError::FsError(
274                "Python returned empty path for capsule package".to_string(),
275            ));
276        }
277
278        let sdk_path = PathBuf::from(&path_str);
279
280        if !sdk_path.exists() {
281            return Err(PythonWasmCompilerError::FsError(format!(
282                "SDK path from Python does not exist: {}",
283                sdk_path.display()
284            )));
285        }
286
287        Ok(sdk_path)
288    }
289
290    fn cleanup_pycache(&self, source_dir: &Path) {
291        Self::remove_pycache_recursive(source_dir);
292        Self::remove_pycache_recursive(&self.cache_dir);
293    }
294
295    fn remove_pycache_recursive(dir: &Path) {
296        if let Ok(entries) = fs::read_dir(dir) {
297            for entry in entries.flatten() {
298                let path = entry.path();
299                if path.is_dir() {
300                    if path.file_name().is_some_and(|n| n == "__pycache__") {
301                        let _ = fs::remove_dir_all(&path);
302                    } else {
303                        Self::remove_pycache_recursive(&path);
304                    }
305                }
306            }
307        }
308    }
309
310    pub fn introspect_task_registry(&self) -> Option<HashMap<String, serde_json::Value>> {
311        let source_dir = self.source_path.parent()?;
312        scanner::scan_python_tasks(source_dir)
313    }
314}