rust_script_ext/
cmd.rs

1use crate::prelude::{anyhow, Context, Result};
2use flume::{unbounded, Sender};
3use itertools::Itertools;
4use std::ffi::OsStr;
5use std::io::{Read, Write};
6use std::path::Path;
7use std::process::*;
8
9/// Describes the handling of a command execution for implementors of [`CommandExecute`].
10#[derive(Copy, Clone, Default)]
11pub enum Output {
12    /// Do not print stdout or stderr.
13    Quiet,
14    /// Print stdout.
15    Stdout,
16    /// Print stderr.
17    Stderr,
18    /// Print stdout and stderr. This is the default option.
19    #[default]
20    Verbose,
21}
22
23/// Execute a command.
24///
25/// This trait is intended to endow [`Command`] with `execute` and `execute_str`, handling the
26/// output of execution for easy use. See the
27/// [implementation on `Command`](#impl-CommandExecute-for-Command)
28/// for more details.
29pub trait CommandExecute {
30    /// Execute and collect output into a byte buffer.
31    fn execute(self, output: Output) -> Result<Vec<u8>>;
32
33    /// Execute and collect output into string.
34    fn execute_str(self, output: Output) -> Result<String>
35    where
36        Self: CommandString + Sized,
37    {
38        let cstr = self.cmd_str();
39        self.execute(output).and_then(|x| {
40            String::from_utf8(x)
41                .context("failed to encode stdout to UTF8 string")
42                .with_context(|| format!("cmd str: {cstr}"))
43        })
44    }
45
46    /// Run the command with no capturing IO.
47    fn run(self) -> Result<()>;
48}
49
50/// Run a [`Command`] to completion and handle the output.
51///
52/// Execution provides a simple way to run a command to completion and capture the outputs.
53/// Both stdout and stderr are captured, the `output` argument describes how they should be
54/// directed to the parent stdio.
55/// By default, output is [`Output::Verbose`] which prints both the stdout and stderr to the terminal.
56///
57/// The result of the execution is the raw stdout bytes. Use `execute_str` to try to encode this
58/// into a `String`.
59/// If the command exits with an error (ie [`ExitStatus::success`] is `false`), an error is
60/// constructed which includes the captured stderr.
61///
62/// ```rust,no_run
63/// # use rust_script_ext::prelude::*;
64/// let ls = cmd!(ls).execute_str(Verbose).unwrap();
65/// assert_eq!(&ls, "Cargo.lock
66/// Cargo.toml
67/// LICENSE
68/// local.rs
69/// README.md
70/// src
71/// target
72/// template.rs
73/// ");
74/// ```
75impl CommandExecute for Command {
76    fn execute(mut self, output: Output) -> Result<Vec<u8>> {
77        // pipe both
78        let mut child = self
79            .stdout(Stdio::piped())
80            .stderr(Stdio::piped())
81            .spawn()
82            .with_context(|| format!("failed to start cmd: {}", self.cmd_str()))?;
83
84        let stdout = child.stdout.take().expect("stdout piped");
85        let stderr = child.stderr.take().expect("stderr piped");
86
87        let (tx_so, rx_so) = unbounded();
88        let (tx_se, rx_se) = unbounded();
89
90        fn fwd(
91            tx: Sender<Vec<u8>>,
92            mut rdr: impl Read + Send + 'static,
93            print: impl Fn(&[u8]) + Send + 'static,
94        ) {
95            std::thread::spawn(move || {
96                let buf: &mut [u8] = &mut *Box::new([0u8; 1024 * 4]);
97                while let Ok(len) = rdr.read(buf) {
98                    if len == 0 {
99                        break;
100                    }
101
102                    let buf = buf[..len].to_vec();
103                    print(&buf);
104                    let _ = tx.send(buf);
105                }
106            });
107        }
108
109        fwd(tx_so, stdout, move |buf| {
110            if matches!(output, Output::Verbose | Output::Stdout) {
111                let _ = std::io::stdout().write_all(buf);
112            }
113        });
114        fwd(tx_se, stderr, move |buf| {
115            if matches!(output, Output::Verbose | Output::Stderr) {
116                let _ = std::io::stderr().write_all(buf);
117            }
118        });
119
120        let xs = child
121            .wait()
122            .with_context(|| format!("failed to execute cmd: {}", self.cmd_str()))?;
123
124        if xs.success() {
125            Ok(rx_so.into_iter().flatten().collect_vec())
126        } else {
127            let se = rx_se.into_iter().flatten().collect_vec();
128            let se = String::from_utf8_lossy(&se).to_string();
129            Err(anyhow!(se)).with_context(|| format!("failed to execute cmd: {}", self.cmd_str(),))
130        }
131    }
132
133    /// Run a command but do not capture IO.
134    ///
135    /// This provides an error message displaying the command run.
136    ///
137    /// Use this method when the command being run uses stdio for progress bars/updates.
138    fn run(mut self) -> Result<()> {
139        self.status().map_err(anyhow::Error::from).and_then(|x| {
140            if x.success() {
141                Ok(())
142            } else {
143                Err(anyhow!("cmd exited with code {}: {}", x, self.cmd_str()))
144            }
145        })
146    }
147}
148
149/// Methods on [`Command`] which take `self`.
150///
151/// This is useful with [`cargs!`](crate::prelude::cargs).
152///
153/// # Example
154/// ```rust
155/// # use rust_script_ext::prelude::*;
156/// cmd!(ls)
157///     .with_args(cargs!(foo/bar, zog))
158///     .run()
159///     .ok();
160/// ```
161pub trait CommandBuilder {
162    /// Akin to [`Command::arg`].
163    fn with_arg<S: AsRef<OsStr>>(self, arg: S) -> Self;
164    /// Akin to [`Command::args`].
165    fn with_args<I, S>(mut self, args: I) -> Self
166    where
167        Self: Sized,
168        I: IntoIterator<Item = S>,
169        S: AsRef<OsStr>,
170    {
171        for a in args {
172            self = self.with_arg(a);
173        }
174        self
175    }
176
177    /// Add the argument if `apply` is `true`.
178    fn maybe_with_arg<S>(self, apply: bool, arg: S) -> Self
179    where
180        Self: Sized,
181        S: AsRef<OsStr>,
182    {
183        if apply {
184            self.with_arg(arg)
185        } else {
186            self
187        }
188    }
189
190    /// Add the arguments if `apply` is `true`.
191    fn maybe_with_args<I, S>(self, apply: bool, args: I) -> Self
192    where
193        Self: Sized,
194        I: IntoIterator<Item = S>,
195        S: AsRef<OsStr>,
196    {
197        if apply {
198            self.with_args(args)
199        } else {
200            self
201        }
202    }
203
204    /// Akin to [`Command::env`].
205    fn with_env<K, V>(self, key: K, val: V) -> Self
206    where
207        K: AsRef<OsStr>,
208        V: AsRef<OsStr>;
209
210    /// Akin to [`Command::envs`].
211    fn with_envs<I, K, V>(mut self, vars: I) -> Self
212    where
213        Self: Sized,
214        I: IntoIterator<Item = (K, V)>,
215        K: AsRef<OsStr>,
216        V: AsRef<OsStr>,
217    {
218        for (k, v) in vars {
219            self = self.with_env(k, v);
220        }
221        self
222    }
223
224    /// Akin to [`Command::current_dir`].
225    fn with_current_dir<P: AsRef<Path>>(self, path: P) -> Self;
226
227    /// Pipe `stdout` of _this_ into `next` command.
228    fn pipe(self, next: Command) -> Result<Self>
229    where
230        Self: Sized;
231
232    /// Pipe `stderr` of _this_ into `next` command.
233    fn pipe_stderr(self, next: Command) -> Result<Self>
234    where
235        Self: Sized;
236}
237
238impl CommandBuilder for Command {
239    fn with_arg<S: AsRef<OsStr>>(mut self, arg: S) -> Self {
240        self.arg(arg);
241        self
242    }
243
244    fn with_env<K, V>(mut self, key: K, val: V) -> Self
245    where
246        K: AsRef<OsStr>,
247        V: AsRef<OsStr>,
248    {
249        self.env(key, val);
250        self
251    }
252
253    fn with_current_dir<P: AsRef<Path>>(mut self, dir: P) -> Self {
254        self.current_dir(dir);
255        self
256    }
257
258    fn pipe(mut self, mut next: Command) -> Result<Self> {
259        let cmd = self
260            .stdout(Stdio::piped())
261            .spawn()
262            .map_err(|e| anyhow!("encountered error with command {}: {e}", self.cmd_str()))?;
263
264        let out = cmd.stdout.expect("piped so should exist");
265        let stdin = Stdio::from(out);
266
267        next.stdin(stdin);
268        Ok(next)
269    }
270
271    fn pipe_stderr(mut self, mut next: Command) -> Result<Self> {
272        let cmd = self
273            .stderr(Stdio::piped())
274            .spawn()
275            .map_err(|e| anyhow!("encountered error with command {}: {e}", self.cmd_str()))?;
276
277        let out = cmd.stderr.expect("piped so should exist");
278        let stdin = Stdio::from(out);
279
280        next.stdin(stdin);
281        Ok(next)
282    }
283}
284
285/// Output [`Command`] as a text string, useful for debugging.
286pub trait CommandString {
287    /// Format the command like a bash string.
288    fn cmd_str(&self) -> String;
289
290    /// Print the command string to stderr.
291    fn debug_print(self) -> Self
292    where
293        Self: Sized,
294    {
295        eprintln!("{}", self.cmd_str());
296        self
297    }
298}
299
300impl CommandString for Command {
301    fn cmd_str(&self) -> String {
302        // note that the debug format is unstable and need careful testing/handling
303        let x = format!("{self:#?}");
304        // eprintln!("{x}");
305
306        let prg = if cfg!(windows) {
307            x.split_once(' ')
308                .map(|x| x.0)
309                .unwrap_or(&x)
310                .trim_matches('"')
311        } else {
312            x.split_once("program:")
313                .expect("known format")
314                .1
315                .split_once(',')
316                .expect("known format")
317                .0
318                .trim()
319                .trim_matches('"')
320        };
321
322        // eprintln!("{prg}");
323
324        self.get_args()
325            .fold(prg.to_string(), |s, a| s + " " + &*a.to_string_lossy())
326    }
327}
328
329#[cfg(test)]
330mod tests {
331    use super::Output::*;
332    use super::*;
333    use crate::prelude::*;
334    use crate::pretty_print_err;
335    use insta::assert_snapshot;
336
337    #[test]
338    fn cmd_macro_output() {
339        let x = cmd!(ls).cmd_str();
340        assert_eq!(&x, "ls");
341
342        let x = cmd!(ls: foo, bar).cmd_str();
343        assert_eq!(&x, "ls foo bar");
344
345        let x = cmd!(ls: {"foo"}, bar).cmd_str();
346        assert_eq!(&x, "ls foo bar");
347
348        let x = cmd!(ls: "foo bar").cmd_str();
349        assert_eq!(&x, r#"ls "foo bar""#);
350
351        let x = cmd!(./script.sh: "foo bar").cmd_str();
352        assert_eq!(&x, r#"./script.sh "foo bar""#);
353    }
354
355    #[test]
356    fn cmd_execute() {
357        let x = cmd!(ls).execute_str(Quiet).unwrap();
358        let mut x = x.trim().split('\n').collect::<Vec<_>>();
359        x.sort();
360
361        assert_eq!(
362            &x,
363            &[
364                "Cargo.lock",
365                "Cargo.toml",
366                "LICENSE",
367                "README.md",
368                "macros",
369                "src",
370                "target",
371                "template-cargo-script.rs",
372                "template-rust-script.rs",
373            ]
374        );
375
376        let x = cmd!(ls: "foo").execute_str(Verbose).unwrap_err();
377        assert_snapshot!("execute-err", pretty_print_err(x));
378
379        let x = cmd!(watcmd: "foo").execute_str(Verbose).unwrap_err();
380        assert_snapshot!("unknown-cmd", pretty_print_err(x));
381    }
382
383    #[test]
384    fn cmd_naming_with_env() {
385        let x = cmd!(ls).with_env("YO", "zog").cmd_str();
386        assert_eq!(&x, "ls");
387
388        let x = cmd!(ls: foo, bar).with_env("YO", "zog").cmd_str();
389        assert_eq!(&x, "ls foo bar");
390
391        let x = cmd!(ls: foo, bar)
392            .with_envs([("YO", "zog"), ("JO", "bar")])
393            .cmd_str();
394        assert_eq!(&x, "ls foo bar");
395    }
396
397    #[test]
398    fn cmd_piping() {
399        let x = cmd!(ls)
400            .pipe(cmd!(grep: Cargo.*))
401            .unwrap()
402            .execute_str(Quiet)
403            .unwrap();
404        let mut x = x.trim().split('\n').collect::<Vec<_>>();
405        x.sort();
406
407        assert_eq!(&x, &["Cargo.lock", "Cargo.toml",]);
408
409        let x = cmd!(ls)
410            .pipe(cmd!(grep: Cargo.*))
411            .unwrap()
412            .pipe(cmd!(grep: toml))
413            .unwrap()
414            .execute_str(Quiet)
415            .unwrap();
416        let mut x = x.trim().split('\n').collect::<Vec<_>>();
417        x.sort();
418
419        assert_eq!(&x, &["Cargo.toml",]);
420
421        let x = cmd!(ls: foo)
422            .pipe_stderr(cmd!(grep: foo))
423            .unwrap()
424            .execute_str(Quiet)
425            .unwrap();
426        let mut x = x.trim().split('\n').collect::<Vec<_>>();
427        x.sort();
428
429        assert_eq!(&x, &["ls: cannot access 'foo': No such file or directory",]);
430    }
431}
rust_script_ext/cmd.rs

rust_script_ext/
cmd.rs
