add save method

Signed-off-by: Freyskeyd <[email protected]>

Author: Freyskeyd

Cargo.toml

@@ -11,3 +11,4 @@ readme = "README.md"
 keywords = ["environment", "env"]
 
 [dependencies]
+regex = "0.2"

src/lib.rs

@@ -1,3 +1,5 @@
+extern crate regex;
+use regex::Regex;
 use std::ffi::OsString;
 
 /// Structure to deal with environment variables
@@ -7,13 +9,15 @@ pub struct Environment {
     vars: Vec<(OsString, OsString)>,
     /// Define if the structure must inherit
     inherit: bool,
+    save_pattern: Option<String>,
 }
 
 impl Default for Environment {
     fn default() -> Self {
         Self {
             vars: vec![],
             inherit: false,
+            save_pattern: None,
         }
     }
 }
@@ -36,6 +40,7 @@ impl Environment {
         Self {
             vars: vec![],
             inherit: true,
+            save_pattern: None,
         }
     }
 
@@ -70,12 +75,53 @@ impl Environment {
         self
     }
 
+    /// Define a regex that will be used to save some existing variables.
+    /// It's only used when having an inherit equal to false.
+    ///
+    /// # Examples
+    ///
+    /// ```rust
+    /// extern crate environment;
+    ///
+    /// use std::ffi::OsString;
+    ///
+    /// ::std::env::set_var("BAR", "BAZ");
+    ///
+    /// let e = environment::Environment::empty().insert("foo", "bar").compile();
+    ///
+    /// assert_eq!(e, vec![(OsString::from("foo"), OsString::from("bar"))]);
+    ///
+    /// let e = environment::Environment::empty().save("BAR").insert("foo", "bar").compile();
+    ///
+    /// assert_eq!(
+    ///     e,
+    ///     vec![
+    ///         (OsString::from("BAR"), OsString::from("BAZ")),
+    ///         (OsString::from("foo"), OsString::from("bar"))
+    ///     ]
+    /// );
+    /// ```
+    pub fn save<S: Into<String>>(mut self, pattern: S) -> Self {
+        self.save_pattern = Some(pattern.into());
+
+        self
+    }
+
     /// Compile Environment object
     pub fn compile(self) -> Vec<(OsString, OsString)> {
         if self.inherit {
             ::std::env::vars_os().chain(self.vars).collect()
         } else {
-            self.vars
+            match self.save_pattern {
+                Some(v) => {
+                    let re = Regex::new(&v).unwrap();
+                    ::std::env::vars_os()
+                        .filter(|&(ref k, _)| re.is_match(k.to_str().unwrap()))
+                        .chain(self.vars)
+                        .collect()
+                }
+                None => self.vars
+            }
         }
     }
 }
@@ -118,6 +164,7 @@ where
         Self {
             vars: v.into_iter().map(|k| k.to_environment_tuple()).collect(),
             inherit: false,
+            save_pattern: None,
         }
     }
 }
@@ -233,4 +280,21 @@ mod test {
         assert!(output.contains("bar=vv"));
         assert!(!output.contains("bar=baz"));
     }
+
+    #[test]
+    fn save_pattern() {
+        let e: Environment = vec![("foo", "bar")].into();
+
+        let mut c = Command::new("printenv");
+
+        let output = c.env_clear()
+            .envs(e.clone().save("PAT.*").compile())
+            .output()
+            .expect("failed to execute command");
+
+        let output = String::from_utf8_lossy(&output.stdout);
+
+        assert!(output.contains("foo=bar"));
+        assert!(output.contains("PATH"));
+    }
 }