From f33b1927a664cdd5beda70e6dea0054fddbdd09d Mon Sep 17 00:00:00 2001
From: Freyskeyd <simon.paitrault@gmail.com>
Date: Fri, 13 Oct 2017 10:01:48 +0200
Subject: [PATCH] add save method

Signed-off-by: Freyskeyd <simon.paitrault@gmail.com>
---
 Cargo.toml |  1 +
 src/lib.rs | 62 +++++++++++++++++++++++++++++++++++++++++++++++++++++-
 2 files changed, 62 insertions(+), 1 deletion(-)

diff --git a/Cargo.toml b/Cargo.toml
index 7226d36..571111e 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -11,3 +11,4 @@ readme = "README.md"
 keywords = ["environment", "env"]
 
 [dependencies]
+regex = "0.2"
diff --git a/src/lib.rs b/src/lib.rs
index bc0e97c..ef65c55 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -1,3 +1,5 @@
+extern crate regex;
+use regex::Regex;
 use std::ffi::OsString;
 
 /// Structure to deal with environment variables
@@ -7,6 +9,7 @@ pub struct Environment {
     vars: Vec<(OsString, OsString)>,
     /// Define if the structure must inherit
     inherit: bool,
+    save_pattern: Option<String>,
 }
 
 impl Default for Environment {
@@ -14,6 +17,7 @@ impl Default for Environment {
         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,49 @@ 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 +160,7 @@ where
         Self {
             vars: v.into_iter().map(|k| k.to_environment_tuple()).collect(),
             inherit: false,
+            save_pattern: None,
         }
     }
 }
@@ -233,4 +276,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"));
+    }
 }