Skip to main content

insta_cmd/
spawn.rs

1use std::collections::BTreeMap;
2use std::env;
3use std::ffi::OsStr;
4use std::io::Write;
5use std::mem;
6use std::ops::{Deref, DerefMut};
7use std::path::Path;
8use std::process::{Command, Output, Stdio};
9
10use serde::Serialize;
11
12#[derive(Serialize)]
13pub struct Info {
14    program: String,
15    args: Vec<String>,
16    #[serde(skip_serializing_if = "BTreeMap::is_empty")]
17    env: BTreeMap<String, String>,
18    #[serde(skip_serializing_if = "Option::is_none")]
19    stdin: Option<String>,
20}
21
22fn describe_program(cmd: &OsStr) -> String {
23    let filename = Path::new(cmd).file_name().unwrap();
24    let name = filename.to_string_lossy();
25    let mut name = &name as &str;
26    if !env::consts::EXE_SUFFIX.is_empty() {
27        name = name.strip_suffix(env::consts::EXE_SUFFIX).unwrap_or(name);
28    }
29    name.into()
30}
31
32impl Info {
33    fn from_std_command(cmd: &Command, stdin: Option<&[u8]>) -> Info {
34        Info {
35            program: describe_program(cmd.get_program()),
36            args: cmd
37                .get_args()
38                .map(|x| x.to_string_lossy().into_owned())
39                .collect(),
40            env: cmd
41                .get_envs()
42                .filter_map(|(k, v)| {
43                    v.map(|val| {
44                        (
45                            k.to_string_lossy().into_owned(),
46                            val.to_string_lossy().into_owned(),
47                        )
48                    })
49                })
50                .collect(),
51            stdin: stdin.as_ref().map(|x| String::from_utf8_lossy(x).into()),
52        }
53    }
54}
55
56/// Implemented by different types that can be spawned and snapshotted.
57pub trait Spawn {
58    #[doc(hidden)]
59    fn spawn_with_info(&mut self, stdin: Option<Vec<u8>>) -> (Info, Output);
60}
61
62/// Utility methods for spawning.
63pub trait SpawnExt {
64    /// This passes the given input to stdin before spawning the command.
65    fn pass_stdin(&mut self, stdin: impl Into<Vec<u8>>) -> SpawnWithStdin<'_>;
66}
67
68impl<T: Spawn> SpawnExt for T {
69    fn pass_stdin(&mut self, stdin: impl Into<Vec<u8>>) -> SpawnWithStdin<'_> {
70        SpawnWithStdin {
71            spawn: self,
72            stdin: stdin.into(),
73        }
74    }
75}
76
77pub struct SpawnWithStdin<'a> {
78    spawn: &'a mut dyn Spawn,
79    stdin: Vec<u8>,
80}
81
82impl<'a> Spawn for SpawnWithStdin<'a> {
83    fn spawn_with_info(&mut self, stdin: Option<Vec<u8>>) -> (Info, Output) {
84        self.spawn
85            .spawn_with_info(Some(stdin.unwrap_or(mem::take(&mut self.stdin))))
86    }
87}
88
89impl Spawn for Command {
90    fn spawn_with_info(&mut self, stdin: Option<Vec<u8>>) -> (Info, Output) {
91        let info = Info::from_std_command(self, stdin.as_deref());
92        let output = if let Some(stdin) = stdin {
93            self.stdin(Stdio::piped());
94            self.stdout(Stdio::piped());
95            self.stderr(Stdio::piped());
96            let mut child = self.spawn().unwrap();
97            let mut child_stdin = child.stdin.take().expect("Failed to open stdin");
98            std::thread::spawn(move || {
99                child_stdin
100                    .write_all(&stdin)
101                    .expect("Failed to write to stdin");
102            });
103            child.wait_with_output().unwrap()
104        } else {
105            self.output().unwrap()
106        };
107        (info, output)
108    }
109}
110
111impl<'a> Spawn for &'a mut Command {
112    fn spawn_with_info(&mut self, stdin: Option<Vec<u8>>) -> (Info, Output) {
113        <Command as Spawn>::spawn_with_info(self, stdin)
114    }
115}
116
117/// Like [`Command`] but sends some input to stdin.
118#[deprecated = "Use .pass_stdin(...) instead"]
119pub struct StdinCommand {
120    command: Command,
121    stdin: Vec<u8>,
122}
123
124#[allow(deprecated)]
125impl StdinCommand {
126    /// Creates a new command that also gets some input value fed to stdin.
127    pub fn new<S: AsRef<OsStr>, I: Into<Vec<u8>>>(program: S, stdin: I) -> StdinCommand {
128        let mut command = Command::new(program);
129        command.stdin(Stdio::piped());
130        command.stdout(Stdio::piped());
131        command.stderr(Stdio::piped());
132        StdinCommand {
133            command,
134            stdin: stdin.into(),
135        }
136    }
137}
138
139#[allow(deprecated)]
140impl Deref for StdinCommand {
141    type Target = Command;
142
143    fn deref(&self) -> &Self::Target {
144        &self.command
145    }
146}
147
148#[allow(deprecated)]
149impl DerefMut for StdinCommand {
150    fn deref_mut(&mut self) -> &mut Self::Target {
151        &mut self.command
152    }
153}
154
155#[allow(deprecated)]
156impl Spawn for StdinCommand {
157    fn spawn_with_info(&mut self, stdin: Option<Vec<u8>>) -> (Info, Output) {
158        Command::spawn_with_info(
159            &mut self.command,
160            Some(stdin.unwrap_or(mem::take(&mut self.stdin))),
161        )
162    }
163}
164
165impl<'a, T: AsRef<OsStr>> Spawn for &'a [T] {
166    fn spawn_with_info(&mut self, stdin: Option<Vec<u8>>) -> (Info, Output) {
167        let mut cmd = Command::new(self.first().expect("expected program name as first item"));
168        for arg in &self[1..] {
169            cmd.arg(arg);
170        }
171        cmd.spawn_with_info(stdin)
172    }
173}
174
175impl<T: AsRef<OsStr>, const N: usize> Spawn for [T; N] {
176    fn spawn_with_info(&mut self, stdin: Option<Vec<u8>>) -> (Info, Output) {
177        (&self[..]).spawn_with_info(stdin)
178    }
179}
180
181impl<T: AsRef<OsStr>> Spawn for Vec<T> {
182    fn spawn_with_info(&mut self, stdin: Option<Vec<u8>>) -> (Info, Output) {
183        (&self[..]).spawn_with_info(stdin)
184    }
185}