library/std/src/sys_common/process.rs

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