xref: /drstd/src/std/sys_common/process.rs (revision 9670759b785600bf6315e4173e46a602f16add7a)
1 #![allow(dead_code)]
2 
3 use crate::std::collections::BTreeMap;
4 use crate::std::env;
5 use crate::std::ffi::{OsStr, OsString};
6 use crate::std::fmt;
7 use crate::std::io;
8 use crate::std::sys::pipe::read2;
9 use crate::std::sys::process::{EnvKey, ExitStatus, Process, StdioPipes};
10 
11 // Stores a set of changes to an environment
12 #[derive(Clone)]
13 pub struct CommandEnv {
14     clear: bool,
15     saw_path: bool,
16     vars: BTreeMap<EnvKey, Option<OsString>>,
17 }
18 
19 impl Default for CommandEnv {
default() -> Self20     fn default() -> Self {
21         CommandEnv {
22             clear: false,
23             saw_path: false,
24             vars: Default::default(),
25         }
26     }
27 }
28 
29 impl fmt::Debug for CommandEnv {
fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result30     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
31         let mut debug_command_env = f.debug_struct("CommandEnv");
32         debug_command_env
33             .field("clear", &self.clear)
34             .field("vars", &self.vars);
35         debug_command_env.finish()
36     }
37 }
38 
39 impl CommandEnv {
40     // Capture the current environment with these changes applied
capture(&self) -> BTreeMap<EnvKey, OsString>41     pub fn capture(&self) -> BTreeMap<EnvKey, OsString> {
42         let mut result = BTreeMap::<EnvKey, OsString>::new();
43         if !self.clear {
44             for (k, v) in env::vars_os() {
45                 result.insert(k.into(), v);
46             }
47         }
48         for (k, maybe_v) in &self.vars {
49             if let &Some(ref v) = maybe_v {
50                 result.insert(k.clone(), v.clone());
51             } else {
52                 result.remove(k);
53             }
54         }
55         result
56     }
57 
is_unchanged(&self) -> bool58     pub fn is_unchanged(&self) -> bool {
59         !self.clear && self.vars.is_empty()
60     }
61 
capture_if_changed(&self) -> Option<BTreeMap<EnvKey, OsString>>62     pub fn capture_if_changed(&self) -> Option<BTreeMap<EnvKey, OsString>> {
63         if self.is_unchanged() {
64             None
65         } else {
66             Some(self.capture())
67         }
68     }
69 
70     // The following functions build up changes
set(&mut self, key: &OsStr, value: &OsStr)71     pub fn set(&mut self, key: &OsStr, value: &OsStr) {
72         let key = EnvKey::from(key);
73         self.maybe_saw_path(&key);
74         self.vars.insert(key, Some(value.to_owned()));
75     }
76 
remove(&mut self, key: &OsStr)77     pub fn remove(&mut self, key: &OsStr) {
78         let key = EnvKey::from(key);
79         self.maybe_saw_path(&key);
80         if self.clear {
81             self.vars.remove(&key);
82         } else {
83             self.vars.insert(key, None);
84         }
85     }
86 
clear(&mut self)87     pub fn clear(&mut self) {
88         self.clear = true;
89         self.vars.clear();
90     }
91 
have_changed_path(&self) -> bool92     pub fn have_changed_path(&self) -> bool {
93         self.saw_path || self.clear
94     }
95 
maybe_saw_path(&mut self, key: &EnvKey)96     fn maybe_saw_path(&mut self, key: &EnvKey) {
97         if !self.saw_path && key == "PATH" {
98             self.saw_path = true;
99         }
100     }
101 
iter(&self) -> CommandEnvs<'_>102     pub fn iter(&self) -> CommandEnvs<'_> {
103         let iter = self.vars.iter();
104         CommandEnvs { iter }
105     }
106 }
107 
108 /// An iterator over the command environment variables.
109 ///
110 /// This struct is created by
111 /// [`Command::get_envs`][crate::std::process::Command::get_envs]. See its
112 /// documentation for more.
113 #[must_use = "iterators are lazy and do nothing unless consumed"]
114 #[derive(Debug)]
115 pub struct CommandEnvs<'a> {
116     iter: crate::std::collections::btree_map::Iter<'a, EnvKey, Option<OsString>>,
117 }
118 
119 impl<'a> Iterator for CommandEnvs<'a> {
120     type Item = (&'a OsStr, Option<&'a OsStr>);
next(&mut self) -> Option<Self::Item>121     fn next(&mut self) -> Option<Self::Item> {
122         self.iter
123             .next()
124             .map(|(key, value)| (key.as_ref(), value.as_deref()))
125     }
size_hint(&self) -> (usize, Option<usize>)126     fn size_hint(&self) -> (usize, Option<usize>) {
127         self.iter.size_hint()
128     }
129 }
130 
131 impl<'a> ExactSizeIterator for CommandEnvs<'a> {
len(&self) -> usize132     fn len(&self) -> usize {
133         self.iter.len()
134     }
is_empty(&self) -> bool135     fn is_empty(&self) -> bool {
136         self.iter.is_empty()
137     }
138 }
139 
wait_with_output( mut process: Process, mut pipes: StdioPipes, ) -> io::Result<(ExitStatus, Vec<u8>, Vec<u8>)>140 pub fn wait_with_output(
141     mut process: Process,
142     mut pipes: StdioPipes,
143 ) -> io::Result<(ExitStatus, Vec<u8>, Vec<u8>)> {
144     drop(pipes.stdin.take());
145 
146     let (mut stdout, mut stderr) = (Vec::new(), Vec::new());
147     match (pipes.stdout.take(), pipes.stderr.take()) {
148         (None, None) => {}
149         (Some(out), None) => {
150             let res = out.read_to_end(&mut stdout);
151             res.unwrap();
152         }
153         (None, Some(err)) => {
154             let res = err.read_to_end(&mut stderr);
155             res.unwrap();
156         }
157         (Some(out), Some(err)) => {
158             let res = read2(out, &mut stdout, err, &mut stderr);
159             res.unwrap();
160         }
161     }
162 
163     let status = process.wait()?;
164     Ok((status, stdout, stderr))
165 }
166