src/libstd/sys_common/process.rs

   1 // Copyright 2014 The Rust Project Developers. See the COPYRIGHT
   2 // file at the top-level directory of this distribution and at
   3 // http://rust-lang.org/COPYRIGHT.
   4 //
   5 // Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
   6 // http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
   7 // <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
   8 // option. This file may not be copied, modified, or distributed
   9 // except according to those terms.
  10
  11 #![allow(dead_code)]
  12 #![unstable(feature = "process_internals", issue = "0")]
  13
  14 use ffi::{OsStr, OsString};
  15 use env;
  16 use collections::BTreeMap;
  17 use borrow::Borrow;
  18
  19 pub trait EnvKey:
  20     From<OsString> + Into<OsString> +
  21     Borrow<OsStr> + Borrow<Self> + AsRef<OsStr> +
  22     Ord + Clone {}
  23
  24 // Implement a case-sensitive environment variable key
  25 #[derive(Clone, Debug, Eq, PartialEq, Ord, PartialOrd)]
  26 pub struct DefaultEnvKey(OsString);
  27
  28 impl From<OsString> for DefaultEnvKey {
  29     fn from(k: OsString) -> Self { DefaultEnvKey(k) }
  30 }
  31
  32 impl From<DefaultEnvKey> for OsString {
  33     fn from(k: DefaultEnvKey) -> Self { k.0 }
  34 }
  35
  36 impl Borrow<OsStr> for DefaultEnvKey {
  37     fn borrow(&self) -> &OsStr { &self.0 }
  38 }
  39
  40 impl AsRef<OsStr> for DefaultEnvKey {
  41     fn as_ref(&self) -> &OsStr { &self.0 }
  42 }
  43
  44 impl EnvKey for DefaultEnvKey {}
  45
  46 // Stores a set of changes to an environment
  47 #[derive(Clone, Debug)]
  48 pub struct CommandEnv<K> {
  49     clear: bool,
  50     saw_path: bool,
  51     vars: BTreeMap<K, Option<OsString>>
  52 }
  53
  54 impl<K: EnvKey> Default for CommandEnv<K> {
  55     fn default() -> Self {
  56         CommandEnv {
  57             clear: false,
  58             saw_path: false,
  59             vars: Default::default()
  60         }
  61     }
  62 }
  63
  64 impl<K: EnvKey> CommandEnv<K> {
  65     // Capture the current environment with these changes applied
  66     pub fn capture(&self) -> BTreeMap<K, OsString> {
  67         let mut result = BTreeMap::<K, OsString>::new();
  68         if !self.clear {
  69             for (k, v) in env::vars_os() {
  70                 result.insert(k.into(), v);
  71             }
  72         }
  73         for (k, maybe_v) in &self.vars {
  74             if let &Some(ref v) = maybe_v {
  75                 result.insert(k.clone(), v.clone());
  76             } else {
  77                 result.remove(k);
  78             }
  79         }
  80         result
  81     }
  82
  83     // Apply these changes directly to the current environment
  84     pub fn apply(&self) {
  85         if self.clear {
  86             for (k, _) in env::vars_os() {
  87                 env::remove_var(k);
  88             }
  89         }
  90         for (key, maybe_val) in self.vars.iter() {
  91             if let &Some(ref val) = maybe_val {
  92                 env::set_var(key, val);
  93             } else {
  94                 env::remove_var(key);
  95             }
  96         }
  97     }
  98
  99     pub fn is_unchanged(&self) -> bool {
 100         !self.clear && self.vars.is_empty()
 101     }
 102
 103     pub fn capture_if_changed(&self) -> Option<BTreeMap<K, OsString>> {
 104         if self.is_unchanged() {
 105             None
 106         } else {
 107             Some(self.capture())
 108         }
 109     }
 110
 111     // The following functions build up changes
 112     pub fn set(&mut self, key: &OsStr, value: &OsStr) {
 113         self.maybe_saw_path(&key);
 114         self.vars.insert(key.to_owned().into(), Some(value.to_owned()));
 115     }
 116     pub fn remove(&mut self, key: &OsStr) {
 117         self.maybe_saw_path(&key);
 118         if self.clear {
 119             self.vars.remove(key);
 120         } else {
 121             self.vars.insert(key.to_owned().into(), None);
 122         }
 123     }
 124     pub fn clear(&mut self) {
 125         self.clear = true;
 126         self.vars.clear();
 127     }
 128     pub fn have_changed_path(&self) -> bool {
 129         self.saw_path || self.clear
 130     }
 131     fn maybe_saw_path(&mut self, key: &OsStr) {
 132         if !self.saw_path && key == "PATH" {
 133             self.saw_path = true;
 134         }
 135     }
 136 }