ahab/src/cmd/argv.rs

use std::{
    fmt::Display,
    fs::File,
    os::unix::process::CommandExt,
    path::Path,
    process::{Command, ExitStatus, Output, Stdio},
};

use anyhow::{Context, Result, anyhow};

use crate::ctx::Ctx;

// a program and its arguments: the only thing in ahab that knows what argv looks like
#[derive(Default, Clone)]
pub struct Argv(Vec<String>);

impl Display for Argv {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.write_str(&self.quoted())
    }
}

impl Argv {
    pub fn new(program: &str) -> Self {
        Self(vec![program.to_string()])
    }

    pub fn arg(mut self, arg: impl AsRef<str>) -> Self {
        self.0.push(arg.as_ref().to_string());
        self
    }

    pub fn args<I, S>(mut self, args: I) -> Self
    where
        I: IntoIterator<Item = S>,
        S: AsRef<str>,
    {
        self.0
            .extend(args.into_iter().map(|arg| arg.as_ref().to_string()));
        self
    }

    // a long option and the value it takes
    pub fn flag(self, name: &str, value: impl AsRef<str>) -> Self {
        self.arg(name).arg(value)
    }

    pub fn program(&self) -> &str {
        self.0.first().map(String::as_str).unwrap_or_default()
    }

    pub fn words(&self) -> &[String] {
        &self.0
    }

    // one word list for `sh -c`, quoted so a value with a space survives the shell
    pub fn quoted(&self) -> String {
        // a nul byte is the only thing shlex refuses, and no argv can hold one
        shlex::try_join(self.0.iter().map(String::as_str)).unwrap_or_default()
    }

    pub fn run(&self, ctx: &Ctx) -> Result<()> {
        if self.skipped(ctx) {
            return Ok(());
        }

        let status = self.command(ctx)?.spawn()?.wait()?;

        self.check(status)
    }

    // reading changes nothing, so a dry run answers the question for real
    pub fn capture(&self, ctx: &Ctx) -> Result<String> {
        let out = self.command(ctx)?.output()?;

        if !out.status.success() {
            // output() holds stderr back, so the command's own complaint has to be
            // passed on here or it is lost
            let stderr = String::from_utf8_lossy(&out.stderr);

            return Err(match stderr.trim() {
                "" => self.failure(out.status),
                reason => anyhow!("`{self}` failed: {reason}"),
            });
        }

        Ok(String::from_utf8(out.stdout)?)
    }

    // replaces this process, so the command's exit code and signals become ours
    pub fn replace(&self, ctx: &Ctx) -> Result<()> {
        if self.skipped(ctx) {
            return Ok(());
        }

        let error = self.command(ctx)?.exec();

        Err(error).with_context(|| format!("running `{self}`"))
    }

    pub fn stream_to(&self, ctx: &Ctx, out: Stdio) -> Result<()> {
        if self.skipped(ctx) {
            return Ok(());
        }

        let status = self.command(ctx)?.stdout(out).spawn()?.wait()?;

        self.check(status)
    }

    // the status rather than an error, for callers with something better to say
    pub fn status(&self, ctx: &Ctx) -> Result<ExitStatus> {
        Ok(self.command(ctx)?.spawn()?.wait()?)
    }

    pub fn stdin_from(&self, ctx: &Ctx, input: &Path) -> Result<ExitStatus> {
        let stdin = self.opened(input)?;

        Ok(self.command(ctx)?.stdin(stdin).spawn()?.wait()?)
    }

    // both streams held back, for callers that read the command's complaints
    pub fn stdin_from_captured(&self, ctx: &Ctx, input: &Path) -> Result<Output> {
        let stdin = self.opened(input)?;

        Ok(self.command(ctx)?.stdin(stdin).output()?)
    }

    // whether it succeeded, without failure being an error
    pub fn quietly_succeeds(&self, ctx: &Ctx) -> Result<bool> {
        Ok(self
            .command(ctx)?
            .stdout(Stdio::null())
            .spawn()?
            .wait()?
            .success())
    }

    fn command(&self, ctx: &Ctx) -> Result<Command> {
        if ctx.verbose {
            eprintln!("running `{self}`");
        }

        let (program, rest) = self.0.split_first().context("empty command")?;
        let mut command = Command::new(program);
        command.args(rest);

        Ok(command)
    }

    fn skipped(&self, ctx: &Ctx) -> bool {
        if ctx.dry_run {
            eprintln!("would run `{self}`");
            return true;
        }

        false
    }

    fn opened(&self, path: &Path) -> Result<Stdio> {
        let file = File::open(path).with_context(|| format!("opening {}", path.display()))?;

        Ok(Stdio::from(file))
    }

    fn check(&self, status: ExitStatus) -> Result<()> {
        if status.success() {
            return Ok(());
        }

        Err(self.failure(status))
    }

    fn failure(&self, status: ExitStatus) -> anyhow::Error {
        match status.code() {
            Some(code) => anyhow!("`{self}` exited with {code}"),
            None => anyhow!("`{self}` was killed by a signal"),
        }
    }
}

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

    #[test]
    fn plain_words_are_left_as_they_are() {
        let argv = Argv::new("pg_dump")
            .flag("--username", "myproject")
            .arg("--data-only")
            .arg("myproject_db");

        assert_eq!(
            argv.quoted(),
            "pg_dump --username myproject --data-only myproject_db"
        );
    }

    #[test]
    fn anything_a_shell_would_read_as_more_than_one_word_is_quoted() {
        assert_eq!(Argv::new("psql").arg("my db").quoted(), "psql 'my db'");
        assert_eq!(Argv::new("sh").arg("a | b").quoted(), "sh 'a | b'");
        assert_eq!(Argv::new("echo").arg("").quoted(), "echo ''");
        assert_eq!(Argv::new("echo").arg("it's").quoted(), r#"echo "it's""#);
    }
}