markdown-code-runner 0.1.0

//! Code execution for Python and Bash code blocks.

use anyhow::{Context, Result};
use std::fs;
use std::path::Path;
use std::process::Command;

/// Language for code execution.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Language {
    Python,
    Bash,
}

impl Language {
    /// Parse a language string into a Language enum.
    pub fn parse(s: &str) -> Option<Self> {
        match s.to_lowercase().as_str() {
            "python" | "python3" | "py" => Some(Language::Python),
            "bash" | "sh" | "shell" => Some(Language::Bash),
            _ => None,
        }
    }
}

/// Execute code and return the output lines.
///
/// If `output_file` is provided, the code is written to the file instead of being executed.
pub fn execute_code(
    code: &[String],
    language: Language,
    output_file: Option<&Path>,
    verbose: bool,
) -> Result<Vec<String>> {
    let full_code = code.join("\n");

    if verbose {
        eprintln!("\n\x1b[1mExecuting code {:?} block:\x1b[0m", language);
        eprintln!("\n{}\n", full_code);
    }

    let output = if let Some(path) = output_file {
        fs::write(path, &full_code)
            .with_context(|| format!("Failed to write code to file: {:?}", path))?;
        Vec::new()
    } else {
        match language {
            Language::Python => execute_python(&full_code)?,
            Language::Bash => execute_bash(&full_code)?,
        }
    };

    if verbose {
        eprintln!("\x1b[1mOutput:\x1b[0m");
        eprintln!("\n{:?}\n", output);
    }

    Ok(output)
}

/// Execute Python code and return the output lines.
fn execute_python(code: &str) -> Result<Vec<String>> {
    let output = Command::new("python3")
        .arg("-c")
        .arg(code)
        .output()
        .context("Failed to execute python3")?;

    let stdout = String::from_utf8_lossy(&output.stdout);
    Ok(stdout.split('\n').map(|s| s.to_string()).collect())
}

/// Execute Bash code and return the output lines.
fn execute_bash(code: &str) -> Result<Vec<String>> {
    let output = Command::new("bash")
        .arg("-c")
        .arg(code)
        .output()
        .context("Failed to execute bash")?;

    let stdout = String::from_utf8_lossy(&output.stdout);
    Ok(stdout.split('\n').map(|s| s.to_string()).collect())
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_execute_python() {
        let code = vec!["print('Hello, world!')".to_string()];
        let output = execute_code(&code, Language::Python, None, false).unwrap();
        assert_eq!(output, vec!["Hello, world!", ""]);
    }

    #[test]
    fn test_execute_bash() {
        let code = vec!["echo \"Hello, world!\"".to_string()];
        let output = execute_code(&code, Language::Bash, None, false).unwrap();
        assert_eq!(output, vec!["Hello, world!", ""]);
    }

    #[test]
    fn test_execute_python_multiline() {
        let code = vec![
            "a = 1".to_string(),
            "b = 2".to_string(),
            "print(a + b)".to_string(),
        ];
        let output = execute_code(&code, Language::Python, None, false).unwrap();
        assert_eq!(output, vec!["3", ""]);
    }

    #[test]
    fn test_language_from_str() {
        assert_eq!(Language::parse("python"), Some(Language::Python));
        assert_eq!(Language::parse("Python"), Some(Language::Python));
        assert_eq!(Language::parse("python3"), Some(Language::Python));
        assert_eq!(Language::parse("bash"), Some(Language::Bash));
        assert_eq!(Language::parse("sh"), Some(Language::Bash));
        assert_eq!(Language::parse("rust"), None);
    }
}