From bc6f7a69f3c474cf68330cff1043e2098ef11a89 Mon Sep 17 00:00:00 2001
From: Jun Wu <quark@meta.com>
Date: Mon, 18 Mar 2024 16:40:23 -0700
Subject: [PATCH] spawn-ext: add checked_output

Summary:
Similar to Python's checked_output, return an error if the process exited non-zero.
The error format is similar to what `git.py` does:

   program exited with code 1:
     git foo bar ...
       stdout + stderr

Reviewed By: muirdm

Differential Revision: D54218774

fbshipit-source-id: d9f2931b059f53cbd9f16616a77dbd5f7e9523b2
---
 eden/scm/lib/spawn-ext/src/lib.rs | 130 ++++++++++++++++++++++++++++++
 1 file changed, 130 insertions(+)

diff --git a/eden/scm/lib/spawn-ext/src/lib.rs b/eden/scm/lib/spawn-ext/src/lib.rs
index 7d707e6af210a..3ed8c3bd79090 100644
--- a/eden/scm/lib/spawn-ext/src/lib.rs
+++ b/eden/scm/lib/spawn-ext/src/lib.rs
@@ -11,9 +11,14 @@
 //!   Unix.
 //! - `spawn_detached` is a quicker way to spawn and forget.
 
+use std::error::Error;
+use std::ffi::OsStr;
+use std::fmt;
 use std::io;
 use std::process::Child;
 use std::process::Command;
+use std::process::ExitStatus;
+use std::process::Output;
 use std::process::Stdio;
 
 /// Extensions to `std::process::Command`.
@@ -30,6 +35,12 @@ pub trait CommandExt {
     /// Return the process id.
     fn spawn_detached(&mut self) -> io::Result<Child>;
 
+    /// Similar to `Output` but reports as an error for non-zero exit code.
+    fn checked_output(&mut self) -> io::Result<Output>;
+
+    /// Similar to `status` but reports an error for non-zero exits.
+    fn checked_run(&mut self) -> io::Result<ExitStatus>;
+
     /// Create a `Command` to run `shell_cmd` through system's shell. This uses "cmd.exe"
     /// on Windows and "/bin/sh" otherwise. Do not add more args to the returned
     /// `Command`. On Windows, you do not need to use the shell to run batch files (the
@@ -37,6 +48,101 @@ pub trait CommandExt {
     fn new_shell(shell_cmd: impl AsRef<str>) -> Command;
 }
 
+#[derive(Debug)]
+struct CommandError {
+    title: String,
+    command: String,
+    output: String,
+    source: Option<io::Error>,
+}
+
+impl fmt::Display for CommandError {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        // title
+        //   command
+        //     output
+        let title = if self.title.is_empty() {
+            "CommandError:"
+        } else {
+            self.title.as_str()
+        };
+        write!(f, "{}\n  {}\n", title, &self.command)?;
+        for line in self.output.lines() {
+            write!(f, "    {line}\n")?;
+        }
+        Ok(())
+    }
+}
+
+impl Error for CommandError {
+    fn source(&self) -> Option<&(dyn Error + 'static)> {
+        self.source.as_ref().map(|s| s as &dyn Error)
+    }
+}
+
+fn os_str_to_naive_quoted_str(s: &OsStr) -> String {
+    let debug_format = format!("{:?}", s);
+    if debug_format.len() == s.len() + 2
+        && debug_format.split_ascii_whitespace().take(2).count() == 1
+    {
+        debug_format[1..debug_format.len() - 1].to_string()
+    } else {
+        debug_format
+    }
+}
+
+impl CommandError {
+    fn new(command: &Command, source: Option<io::Error>) -> Self {
+        let arg0 = os_str_to_naive_quoted_str(command.get_program());
+        let args = command
+            .get_args()
+            .map(os_str_to_naive_quoted_str)
+            .collect::<Vec<String>>()
+            .join(" ");
+        let command = format!("{arg0} {args}");
+        Self {
+            title: Default::default(),
+            output: Default::default(),
+            command,
+            source,
+        }
+    }
+
+    fn with_output(mut self, output: &Output) -> Self {
+        for out in [&output.stdout, &output.stderr] {
+            self.output.push_str(&String::from_utf8_lossy(out));
+        }
+        self.with_status(&output.status)
+    }
+
+    fn with_status(mut self, exit: &ExitStatus) -> Self {
+        match exit.code() {
+            None =>
+            {
+                #[cfg(unix)]
+                match std::os::unix::process::ExitStatusExt::signal(exit) {
+                    Some(sig) => self.title = format!("Command terminated by signal {}", sig),
+                    None => {}
+                }
+            }
+            Some(code) => {
+                if code != 0 {
+                    self.title = format!("Command exited with code {}", code);
+                }
+            }
+        }
+        self
+    }
+
+    fn into_io_error(self: CommandError) -> io::Error {
+        let kind = match self.source.as_ref() {
+            None => io::ErrorKind::Other,
+            Some(e) => e.kind(),
+        };
+        io::Error::new(kind, self)
+    }
+}
+
 impl CommandExt for Command {
     fn avoid_inherit_handles(&mut self) -> &mut Self {
         #[cfg(unix)]
@@ -67,6 +173,30 @@ impl CommandExt for Command {
             .spawn()
     }
 
+    fn checked_output(&mut self) -> io::Result<Output> {
+        let out = self
+            .output()
+            .map_err(|e| CommandError::new(self, Some(e)).into_io_error())?;
+        if !out.status.success() {
+            return Err(CommandError::new(self, None)
+                .with_output(&out)
+                .into_io_error());
+        }
+        Ok(out)
+    }
+
+    fn checked_run(&mut self) -> io::Result<ExitStatus> {
+        let status = self
+            .status()
+            .map_err(|e| CommandError::new(self, Some(e)).into_io_error())?;
+        if !status.success() {
+            return Err(CommandError::new(self, None)
+                .with_status(&status)
+                .into_io_error());
+        }
+        Ok(status)
+    }
+
     fn new_shell(shell_cmd: impl AsRef<str>) -> Command {
         #[cfg(unix)]
         return unix::new_shell(shell_cmd);