markdown_code_runner/
executor.rs

1//! Code execution for Python and Bash code blocks.
2
3use anyhow::{Context, Result};
4use std::fs;
5use std::path::Path;
6use std::process::Command;
7
8/// Language for code execution.
9#[derive(Debug, Clone, Copy, PartialEq, Eq)]
10pub enum Language {
11    Python,
12    Bash,
13}
14
15impl Language {
16    /// Parse a language string into a Language enum.
17    pub fn parse(s: &str) -> Option<Self> {
18        match s.to_lowercase().as_str() {
19            "python" | "python3" | "py" => Some(Language::Python),
20            "bash" | "sh" | "shell" => Some(Language::Bash),
21            _ => None,
22        }
23    }
24}
25
26/// Execute code and return the output lines.
27///
28/// If `output_file` is provided, the code is written to the file instead of being executed.
29pub fn execute_code(
30    code: &[String],
31    language: Language,
32    output_file: Option<&Path>,
33    verbose: bool,
34) -> Result<Vec<String>> {
35    let full_code = code.join("\n");
36
37    if verbose {
38        eprintln!("\n\x1b[1mExecuting code {:?} block:\x1b[0m", language);
39        eprintln!("\n{}\n", full_code);
40    }
41
42    let output = if let Some(path) = output_file {
43        fs::write(path, &full_code)
44            .with_context(|| format!("Failed to write code to file: {:?}", path))?;
45        Vec::new()
46    } else {
47        match language {
48            Language::Python => execute_python(&full_code)?,
49            Language::Bash => execute_bash(&full_code)?,
50        }
51    };
52
53    if verbose {
54        eprintln!("\x1b[1mOutput:\x1b[0m");
55        eprintln!("\n{:?}\n", output);
56    }
57
58    Ok(output)
59}
60
61/// Execute Python code and return the output lines.
62fn execute_python(code: &str) -> Result<Vec<String>> {
63    let output = Command::new("python3")
64        .arg("-c")
65        .arg(code)
66        .output()
67        .context("Failed to execute python3")?;
68
69    let stdout = String::from_utf8_lossy(&output.stdout);
70    Ok(stdout.split('\n').map(|s| s.to_string()).collect())
71}
72
73/// Execute Bash code and return the output lines.
74fn execute_bash(code: &str) -> Result<Vec<String>> {
75    let output = Command::new("bash")
76        .arg("-c")
77        .arg(code)
78        .output()
79        .context("Failed to execute bash")?;
80
81    let stdout = String::from_utf8_lossy(&output.stdout);
82    Ok(stdout.split('\n').map(|s| s.to_string()).collect())
83}
84
85#[cfg(test)]
86mod tests {
87    use super::*;
88
89    #[test]
90    fn test_execute_python() {
91        let code = vec!["print('Hello, world!')".to_string()];
92        let output = execute_code(&code, Language::Python, None, false).unwrap();
93        assert_eq!(output, vec!["Hello, world!", ""]);
94    }
95
96    #[test]
97    fn test_execute_bash() {
98        let code = vec!["echo \"Hello, world!\"".to_string()];
99        let output = execute_code(&code, Language::Bash, None, false).unwrap();
100        assert_eq!(output, vec!["Hello, world!", ""]);
101    }
102
103    #[test]
104    fn test_execute_python_multiline() {
105        let code = vec![
106            "a = 1".to_string(),
107            "b = 2".to_string(),
108            "print(a + b)".to_string(),
109        ];
110        let output = execute_code(&code, Language::Python, None, false).unwrap();
111        assert_eq!(output, vec!["3", ""]);
112    }
113
114    #[test]
115    fn test_language_from_str() {
116        assert_eq!(Language::parse("python"), Some(Language::Python));
117        assert_eq!(Language::parse("Python"), Some(Language::Python));
118        assert_eq!(Language::parse("python3"), Some(Language::Python));
119        assert_eq!(Language::parse("bash"), Some(Language::Bash));
120        assert_eq!(Language::parse("sh"), Some(Language::Bash));
121        assert_eq!(Language::parse("rust"), None);
122    }
123}