honey-biscuit-workshop/workshop-engine/src/upm/python.rs

use async_trait::async_trait;
use serde::Deserialize;
use std::collections::HashMap;
use std::path::PathBuf;

use crate::error::DependencyError;
use crate::types::{DeltaExecutionContext, Engine, EventSender, ExecutionContext};
use crate::upm::UPMSysDeps;

use super::UserPackageManager;

#[derive(Clone)]
pub struct Python;

fn default_type() -> String {
    "pip".to_string()
}

// ── PythonAction (one entry in the list) ────────────────────────────

/// A single Python package-management operation.
///
/// `type` selects the backend: ``pip``, ``uv``, ``poetry``, or ``pipx``.
/// Everything else applies to that backend's invocation.
#[derive(Debug, Clone, Deserialize)]
struct PythonAction {
    /// Backend tool name.
    #[serde(default = "default_type", rename = "type")]
    type_: String,

    /// Python interpreter to use (default: system ``python3``).
    #[serde(default)]
    python: Option<String>,

    /// Virtualenv directory — created if set.
    #[serde(default)]
    directory: Option<PathBuf>,

    /// PyPI index URL (``-i`` / ``--index-url``).
    #[serde(default)]
    index: Option<String>,

    /// Package names.
    #[serde(default)]
    packages: Vec<String>,

    /// Requirements file or pyproject.toml.
    #[serde(default)]
    manifest: Option<PathBuf>,

    /// Extra CLI arguments forwarded to the backend.
    #[serde(default)]
    args: Vec<String>,

    /// Environment variables injected into the child process.
    #[serde(default, rename = "env")]
    env_: HashMap<String, Option<String>>,
}

// ── PythonConfig (top-level, transparent vec) ───────────────────────

#[derive(Debug, Clone, Deserialize)]
#[serde(transparent)]
struct PythonConfig(Vec<PythonAction>);

// ── helpers ─────────────────────────────────────────────────────────

fn python_bin(action: &PythonAction) -> &str {
    action.python.as_deref().unwrap_or("python3")
}

fn index_flag(action: &PythonAction) -> String {
    match &action.index {
        Some(idx) => format!(" -i '{}'", idx),
        None => String::new(),
    }
}

fn env_delta(action: &PythonAction) -> DeltaExecutionContext {
    DeltaExecutionContext {
        env_vars: action.env_.clone(),
        ..Default::default()
    }
}

fn venv_pip(action: &PythonAction) -> String {
    match &action.directory {
        Some(dir) => format!("{}/bin/pip", dir.display()),
        None => format!("{} -m pip", python_bin(action)),
    }
}

// ── UserPackageManager impl ─────────────────────────────────────────

#[async_trait]
impl UserPackageManager for Python {
    fn name(&self) -> &'static str {
        "python"
    }

    fn system_dependency(&self, config: &serde_yaml::Value) -> UPMSysDeps {
        let mut pkgs = vec!["python3".to_string()];

        if let Ok(cfg) = serde_yaml::from_value::<PythonConfig>(config.clone()) {
            for a in &cfg.0 {
                match a.type_.as_str() {
                    "uv" => pkgs.push("uv".into()),
                    "poetry" => pkgs.push("poetry".into()),
                    "pipx" => pkgs.push("pipx".into()),
                    _ => {}
                }
            }
            pkgs.sort();
            pkgs.dedup();
        }

        UPMSysDeps::new(pkgs, false)
    }

    async fn main(
        &self,
        config: &serde_yaml::Value,
        ctx: &ExecutionContext,
        event_tx: &EventSender,
    ) -> Result<(), DependencyError> {
        let cfg: PythonConfig =
            serde_yaml::from_value(config.clone()).map_err(|e| {
                log::error!("Failed to parse user.python config: {}", e);
                DependencyError::ParseError()
            })?;

        let engine = Engine::new();

        for action in &cfg.0 {
            self.run_action(&engine, action, ctx, event_tx).await?;
        }

        Ok(())
    }
}

impl Python {
    async fn run_action(
        &self,
        engine: &Engine,
        action: &PythonAction,
        ctx: &ExecutionContext,
        event_tx: &EventSender,
    ) -> Result<(), DependencyError> {
        let delta = env_delta(action);
        let idx = index_flag(action);

        // ── 1. venv ──────────────────────────────────────────────
        if let Some(ref dir) = action.directory {
            let venv_cmd =
                format!("{} -m venv '{}'", python_bin(action), dir.display());
            log::info!("Python UPM creating venv: {}", venv_cmd);
            let r = engine
                .execute_command_with_delta(&venv_cmd, ctx, event_tx, &delta)
                .await;
            if let Err(e) = r {
                log::error!("Failed to create venv: {}", e);
                return Err(DependencyError::CustomFailed(e));
            }
        }

        // ── 2. backend command ──────────────────────────────────
        match action.type_.as_str() {
            "pip" => {
                if !action.packages.is_empty() {
                    let cmd = format!(
                        "{} install{}{} {}",
                        venv_pip(action),
                        idx,
                        action.args.iter().fold(String::new(), |a, b| a + " " + b),
                        action.packages.join(" "),
                    );
                    log::info!("Python UPM executing: {}", cmd);
                    engine
                        .execute_command_with_delta(&cmd, ctx, event_tx, &delta)
                        .await
                        .map_err(DependencyError::CustomFailed)?;
                }
                if let Some(ref manifest) = action.manifest {
                    let cmd = format!(
                        "{} install{} -r '{}'{}",
                        venv_pip(action),
                        idx,
                        manifest.display(),
                        action.args.iter().fold(String::new(), |a, b| a + " " + b),
                    );
                    log::info!("Python UPM executing: {}", cmd);
                    engine
                        .execute_command_with_delta(&cmd, ctx, event_tx, &delta)
                        .await
                        .map_err(DependencyError::CustomFailed)?;
                }
            }

            "uv" => {
                // uv pip install <packages> / uv pip sync <manifest>
                if !action.packages.is_empty() {
                    let cmd = format!(
                        "uv pip install{}{} {}",
                        idx,
                        action.args.iter().fold(String::new(), |a, b| a + " " + b),
                        action.packages.join(" "),
                    );
                    log::info!("Python UPM executing: {}", cmd);
                    engine
                        .execute_command_with_delta(&cmd, ctx, event_tx, &delta)
                        .await
                        .map_err(DependencyError::CustomFailed)?;
                }
                if let Some(ref manifest) = action.manifest {
                    let cmd = format!(
                        "uv pip sync '{}'{}",
                        manifest.display(),
                        action.args.iter().fold(String::new(), |a, b| a + " " + b),
                    );
                    log::info!("Python UPM executing: {}", cmd);
                    engine
                        .execute_command_with_delta(&cmd, ctx, event_tx, &delta)
                        .await
                        .map_err(DependencyError::CustomFailed)?;
                }
            }

            "poetry" => {
                let cmd = format!(
                    "poetry install{}",
                    action.args.iter().fold(String::new(), |a, b| a + " " + b),
                );
                log::info!("Python UPM executing: {}", cmd);
                engine
                    .execute_command_with_delta(&cmd, ctx, event_tx, &delta)
                    .await
                    .map_err(DependencyError::CustomFailed)?;
            }

            "pipx" => {
                if !action.packages.is_empty() {
                    let cmd = format!(
                        "pipx install{}{} {}",
                        idx,
                        action.args.iter().fold(String::new(), |a, b| a + " " + b),
                        action.packages.join(" "),
                    );
                    log::info!("Python UPM executing: {}", cmd);
                    engine
                        .execute_command_with_delta(&cmd, ctx, event_tx, &delta)
                        .await
                        .map_err(DependencyError::CustomFailed)?;
                }
            }

            other => {
                log::error!("Unknown Python backend type: {}", other);
                return Err(DependencyError::ParseError());
            }
        }

        Ok(())
    }
}

// ── Tests ──────────────────────────────────────────────────────────

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

    #[test]
    fn test_name() {
        assert_eq!(Python.name(), "python");
    }

    #[test]
    fn test_system_dependency_basic() {
        let yaml = "- pip: {}";
        let v: serde_yaml::Value = serde_yaml::from_str(yaml).unwrap();
        let deps = Python.system_dependency(&v);
        assert_eq!(deps.packages, vec!["python3"]);
    }

    #[test]
    fn test_system_dependency_uv() {
        let yaml = r#"
- type: uv
  packages: [pytest]
"#;
        let v: serde_yaml::Value = serde_yaml::from_str(yaml).unwrap();
        let deps = Python.system_dependency(&v);
        assert!(deps.packages.contains(&"python3".to_string()));
        assert!(deps.packages.contains(&"uv".to_string()));
    }

    #[test]
    fn test_parse_pip_action() {
        let yaml = r#"
- type: pip
  packages: [pytest, requests]
"#;
        let cfg: PythonConfig = serde_yaml::from_str(yaml).unwrap();
        assert_eq!(cfg.0.len(), 1);
        let a = &cfg.0[0];
        assert_eq!(a.type_, "pip");
        assert_eq!(a.packages, vec!["pytest", "requests"]);
    }

    #[test]
    fn test_parse_full_action() {
        let yaml = r#"
- type: uv
  python: "python3.14"
  directory: ".venv"
  index: "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple"
  packages: [pytest, requests]
  manifest: "requirements.txt"
  args: ["--upgrade"]
  env:
    UV_PYTHON: "python3.14"
"#;
        let cfg: PythonConfig = serde_yaml::from_str(yaml).unwrap();
        let a = &cfg.0[0];
        assert_eq!(a.type_, "uv");
        assert_eq!(a.python.as_deref(), Some("python3.14"));
        assert_eq!(a.directory, Some(PathBuf::from(".venv")));
        assert_eq!(
            a.index.as_deref(),
            Some("https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple"),
        );
        assert_eq!(a.packages, vec!["pytest", "requests"]);
        assert_eq!(a.manifest, Some(PathBuf::from("requirements.txt")));
        assert_eq!(a.args, vec!["--upgrade"]);
        assert_eq!(
            a.env_.get("UV_PYTHON").and_then(|v| v.as_deref()),
            Some("python3.14"),
        );
    }

    #[test]
    fn test_parse_multiple_actions() {
        let yaml = r#"
- type: pip
  packages: [pytest]
- type: uv
  directory: ".venv"
  manifest: "requirements.txt"
"#;
        let cfg: PythonConfig = serde_yaml::from_str(yaml).unwrap();
        assert_eq!(cfg.0.len(), 2);
        assert_eq!(cfg.0[0].type_, "pip");
        assert_eq!(cfg.0[1].type_, "uv");
    }
}