shx/
cmd.rs

1//! Command-wrapper that simplifies piping and other operations.
2use std::{
3    borrow::Cow,
4    io::{ErrorKind, Read, Write},
5    ops::{Deref, DerefMut},
6    process::{Command, Stdio},
7};
8
9use thiserror::Error;
10
11/// Builder object for [`QCmd`]
12#[derive(Debug)]
13pub struct CmdBuilder<'source, 'sink> {
14    cmd: Command,
15    source: Source<'source>,
16    sink: Sink<'sink>,
17}
18
19/// Command stdin sources.
20#[derive(Debug, Default)]
21pub enum Source<'source> {
22    /// Inherit the stdin pipe
23    #[default]
24    Stdin,
25    /// Pipe the input from a string-like object
26    Str(Cow<'source, str>),
27    /// Pipe the input from a bytes-like object
28    Bytes(Cow<'source, [u8]>),
29}
30
31/// Command stdout sinks
32#[derive(Debug, Default)]
33pub enum Sink<'sink> {
34    /// Inherit the stdout pipe
35    #[default]
36    Stdout,
37    /// Append output to a [`String`]
38    Str(&'sink mut String),
39    /// Append output to a [`Vec<u8>`]
40    Bytes(&'sink mut Vec<u8>),
41}
42
43impl<'source, 'sink> CmdBuilder<'source, 'sink> {
44    /// Creates the builder with the given Command.
45    pub fn new(cmd: Command) -> Self {
46        Self {
47            cmd,
48            source: Source::Stdin,
49            sink: Sink::Stdout,
50        }
51    }
52
53    /// Sets the source
54    pub fn source(&mut self, source: impl Into<Source<'source>>) -> &mut Self {
55        self.source = source.into();
56        self
57    }
58
59    /// Sets the sink
60    pub fn sink(&mut self, sink: impl Into<Sink<'sink>>) -> &mut Self {
61        self.sink = sink.into();
62        self
63    }
64
65    /// Builds the [`QCmd`], consuming the builder.
66    pub fn build(self) -> Cmd<'source, 'sink> {
67        Cmd::new(self.cmd, self.source, self.sink)
68    }
69}
70
71/// A "quick" command that holds references to the source and sink.
72///
73/// The canonical way to construct this is with the [`cmd!`](crate::cmd) macro.
74#[derive(Debug)]
75pub struct Cmd<'source, 'sink> {
76    cmd: Command,
77    source: Source<'source>,
78    sink: Sink<'sink>,
79}
80
81/// Error executing the command
82#[derive(Error, Debug)]
83pub enum Error {
84    /// IO Error
85    #[error("IO Error while executing the command")]
86    Io(#[from] std::io::Error),
87    /// Non-zero status code
88    #[error("Command returned a non okay status")]
89    StatusFailure(i32),
90    /// Unexpected termination by signal
91    #[error("Process was terminated by a signal")]
92    UnexpectedTermination,
93    /// Output data is not UTF8
94    #[error("Piped output does not conform to UTF-8")]
95    NotUtf8,
96}
97
98impl Deref for Cmd<'_, '_> {
99    type Target = Command;
100    fn deref(&self) -> &Self::Target {
101        &self.cmd
102    }
103}
104impl DerefMut for Cmd<'_, '_> {
105    fn deref_mut(&mut self) -> &mut Self::Target {
106        &mut self.cmd
107    }
108}
109
110impl<'source, 'sink> Cmd<'source, 'sink> {
111    /// Creates a new QCmd with the optional source and sink.
112    pub fn new(mut cmd: Command, source: Source<'source>, sink: Sink<'sink>) -> Self {
113        match &source {
114            Source::Stdin => cmd.stdin(Stdio::inherit()),
115            Source::Str(_) | Source::Bytes(_) => cmd.stdin(Stdio::piped()),
116        };
117        match &sink {
118            Sink::Stdout => cmd.stdout(Stdio::inherit()),
119            Sink::Str(_) | Sink::Bytes(_) => cmd.stdout(Stdio::piped()),
120        };
121        Self { cmd, source, sink }
122    }
123
124    /// Executes the command piping the I/O as requested.
125    pub fn exec(&mut self) -> Result<(), Error> {
126        let mut child = self.cmd.spawn()?;
127        match &self.source {
128            Source::Stdin => {}
129            Source::Str(source) => {
130                let mut stdin = child.stdin.take().unwrap();
131                stdin.write_all(source.as_bytes())?;
132            }
133            Source::Bytes(source) => {
134                let mut stdin = child.stdin.take().unwrap();
135                stdin.write_all(source.as_ref())?;
136            }
137        }
138        let status = child.wait().expect("Child process wasn't running?");
139        match status.code() {
140            Some(i) if i != 0 => return Err(Error::StatusFailure(i)),
141            Some(zero) => debug_assert!(zero == 0),
142            None => return Err(Error::UnexpectedTermination),
143        }
144        match &mut self.sink {
145            Sink::Stdout => {}
146            Sink::Str(ref mut sink) => {
147                // This can't fail due to QCmd::new
148                let mut stdout = child.stdout.take().unwrap();
149                if let Err(e) = stdout.read_to_string(sink) {
150                    match e.kind() {
151                        ErrorKind::InvalidData => return Err(Error::NotUtf8),
152                        _ => return Err(Error::Io(e)),
153                    }
154                }
155            }
156            Sink::Bytes(ref mut sink) => {
157                // This can't fail due to QCmd::new
158                let mut stdout = child.stdout.take().unwrap();
159                stdout.read_to_end(sink)?;
160            }
161        }
162        Ok(())
163    }
164}
165
166impl<'source> From<&'source str> for Source<'source> {
167    fn from(value: &'source str) -> Self {
168        Self::Str(Cow::Borrowed(value))
169    }
170}
171impl<'source> From<String> for Source<'source> {
172    fn from(value: String) -> Self {
173        Self::Str(Cow::Owned(value))
174    }
175}
176impl<'source> From<&'source [u8]> for Source<'source> {
177    fn from(value: &'source [u8]) -> Self {
178        Self::Bytes(Cow::Borrowed(value))
179    }
180}
181impl<'source> From<Vec<u8>> for Source<'source> {
182    fn from(value: Vec<u8>) -> Self {
183        Self::Bytes(Cow::Owned(value))
184    }
185}
186
187impl<'sink> From<&'sink mut String> for Sink<'sink> {
188    fn from(value: &'sink mut String) -> Self {
189        Self::Str(value)
190    }
191}
192impl<'sink> From<&'sink mut Vec<u8>> for Sink<'sink> {
193    fn from(value: &'sink mut Vec<u8>) -> Self {
194        Self::Bytes(value)
195    }
196}