capsule_core/wasm/compiler/
python.rs1use 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}