     1
     2
     3
     4
     5
     6
     7
     8
     9
    10
    11
    12
    13
    14
    15
    16
    17
    18
    19
    20
    21
    22
    23
    24
    25
    26
    27
    28
    29
    30
    31
    32
    33
    34
    35
    36
    37
    38
    39
    40
    41
    42
    43
    44
    45
    46
    47
    48
    49
    50
    51
    52
    53
    54
    55
    56
    57
    58
    59
    60
    61
    62
    63
    64
    65
    66
    67
    68
    69
    70
    71
    72
    73
    74
    75
    76
    77
    78
    79
    80
    81
    82
    83
    84
    85
    86
    87
    88
    89
    90
    91
    92
    93
    94
    95

    #![allow(dead_code)]
    #![unstable(feature = "process_internals", issue = "none")]

    use crate::collections::BTreeMap;
    use crate::env;
    use crate::ffi::{OsStr, OsString};
    use crate::sys::process::EnvKey;

    // Stores a set of changes to an environment
    #[derive(Clone, Debug)]
    pub struct CommandEnv {
        clear: bool,
        saw_path: bool,
        vars: BTreeMap<EnvKey, Option<OsString>>,
    }

    impl Default for CommandEnv {
        fn default() -> Self {
            CommandEnv { clear: false, saw_path: false, vars: Default::default() }
        }
    }

    impl CommandEnv {
        // Capture the current environment with these changes applied
        pub fn capture(&self) -> BTreeMap<EnvKey, OsString> {
            let mut result = BTreeMap::<EnvKey, OsString>::new();
            if !self.clear {
                for (k, v) in env::vars_os() {
                    result.insert(k.into(), v);
                }
            }
            for (k, maybe_v) in &self.vars {
                if let &Some(ref v) = maybe_v {
                    result.insert(k.clone(), v.clone());
                } else {
                    result.remove(k);
                }
            }
            result
        }

        // Apply these changes directly to the current environment
        pub fn apply(&self) {
            if self.clear {
                for (k, _) in env::vars_os() {
                    env::remove_var(k);
                }
            }
            for (key, maybe_val) in self.vars.iter() {
                if let Some(ref val) = maybe_val {
                    env::set_var(key, val);
                } else {
                    env::remove_var(key);
                }
            }
        }

        pub fn is_unchanged(&self) -> bool {
            !self.clear && self.vars.is_empty()
        }

        pub fn capture_if_changed(&self) -> Option<BTreeMap<EnvKey, OsString>> {
            if self.is_unchanged() { None } else { Some(self.capture()) }
        }

        // The following functions build up changes
        pub fn set(&mut self, key: &OsStr, value: &OsStr) {
            self.maybe_saw_path(&key);
            self.vars.insert(key.to_owned().into(), Some(value.to_owned()));
        }

        pub fn remove(&mut self, key: &OsStr) {
            self.maybe_saw_path(&key);
            if self.clear {
                self.vars.remove(key);
            } else {
                self.vars.insert(key.to_owned().into(), None);
            }
        }

        pub fn clear(&mut self) {
            self.clear = true;
            self.vars.clear();
        }

        pub fn have_changed_path(&self) -> bool {
            self.saw_path || self.clear
        }

        fn maybe_saw_path(&mut self, key: &OsStr) {
            if !self.saw_path && key == "PATH" {
                self.saw_path = true;
            }
        }
    }