From 6974eb099499c2c4bdc4040035837d3e0957530b Mon Sep 17 00:00:00 2001
From: Sam Hedin <sam.hedin@gmail.com>
Date: Sat, 6 Jun 2020 08:01:50 +0200
Subject: [PATCH] Restore env values

---
 .../src/env/directory_specific_environment.rs | 58 ++++++++++++++++++-
 crates/nu-cli/src/env/environment.rs          |  4 ++
 2 files changed, 59 insertions(+), 3 deletions(-)

diff --git a/crates/nu-cli/src/env/directory_specific_environment.rs b/crates/nu-cli/src/env/directory_specific_environment.rs
index 74ea07da17..bd8f4d88cd 100644
--- a/crates/nu-cli/src/env/directory_specific_environment.rs
+++ b/crates/nu-cli/src/env/directory_specific_environment.rs
@@ -1,6 +1,7 @@
 use indexmap::IndexMap;
 use nu_protocol::{Primitive, UntaggedValue, Value};
-use std::{fmt::Debug, path::PathBuf};
+use std::io::Write;
+use std::{ffi::OsString, fmt::Debug, path::PathBuf, fs::OpenOptions};
 
 #[derive(Debug, Default)]
 pub struct DirectorySpecificEnvironment {
@@ -10,7 +11,7 @@ pub struct DirectorySpecificEnvironment {
     pub added_env_vars: IndexMap<PathBuf, Vec<String>>,
 
     //Directory -> (env_key, value). If a .nu overwrites some existing environment variables, they are added here so that they can be restored later.
-    pub overwritten_env_values: IndexMap<PathBuf, Vec<(String, String)>>,
+    pub overwritten_env_values: IndexMap<PathBuf, Vec<(String, OsString)>>,
 }
 
 impl DirectorySpecificEnvironment {
@@ -44,6 +45,35 @@ impl DirectorySpecificEnvironment {
         }
     }
 
+    // pub overwritten_env_values: IndexMap<PathBuf, Vec<(String, OsString)>>,
+    //overwritten_env_values maps a directory with a .nu file to some environment variables that it overwrote.
+    //if we are not in that directory, we re-add those variables.
+    pub fn overwritten_values_to_restore(&mut self) -> std::io::Result<IndexMap<String, String>> {
+        let current_dir = std::env::current_dir()?;
+
+        let mut keyvals_to_restore = IndexMap::new();
+        for (directory, keyvals) in &self.overwritten_env_values {
+            let mut working_dir = Some(current_dir.as_path());
+
+            let mut readd = true;
+            while let Some(wdir) = working_dir {
+                if wdir == directory.as_path() {
+                    readd = false;
+                    break;
+                } else {
+                    working_dir = working_dir.unwrap().parent();
+                }
+            }
+            if readd {
+                for (k, v) in keyvals {
+                    keyvals_to_restore.insert(k.clone(), v.to_str().unwrap().to_string());
+                }
+            }
+        }
+
+        Ok(keyvals_to_restore)
+    }
+
     pub fn env_vars_to_add(&mut self) -> std::io::Result<IndexMap<String, String>> {
         let current_dir = std::env::current_dir()?;
 
@@ -60,7 +90,7 @@ impl DirectorySpecificEnvironment {
                         .unwrap();
                     let vars_in_current_file = toml_doc.get("env").unwrap().as_table().unwrap();
 
-                    let keys_in_file = vars_in_current_file
+                    let keys_in_file: Vec<String> = vars_in_current_file
                         .iter()
                         .map(|(k, v)| {
                             vars_to_add.insert(k.clone(), v.as_str().unwrap().to_string()); //This is to add the keys and values to the environment
@@ -68,7 +98,29 @@ impl DirectorySpecificEnvironment {
                         })
                         .collect();
 
+                    self.overwritten_env_values.insert(
+                        wdir.to_path_buf(),
+                        keys_in_file.iter().fold(vec![], |mut keyvals, key| {
+                            if let Some(val) = std::env::var_os(key) {
+
+                                let mut file = OpenOptions::new()
+                                    .write(true)
+                                    .append(true)
+                                    .create(true)
+                                    .open("restore.txt").unwrap();
+
+                                write!(&mut file, "about to overwrite: {:?}\n", val).unwrap();
+
+                                keyvals.push((key.clone(), val));
+                                keyvals
+                            } else {
+                                keyvals
+                            }
+                        }),
+                    );
+
                     self.added_env_vars.insert(wdir.to_path_buf(), keys_in_file);
+
                     break;
                 } else {
                     working_dir = working_dir.unwrap().parent();
diff --git a/crates/nu-cli/src/env/environment.rs b/crates/nu-cli/src/env/environment.rs
index 0e987e06d1..3d4df5e7fd 100644
--- a/crates/nu-cli/src/env/environment.rs
+++ b/crates/nu-cli/src/env/environment.rs
@@ -66,6 +66,10 @@ impl Environment {
         self.direnv.env_vars_to_delete()?.iter().for_each(|v| {
             self.remove_env(v);
         });
+
+        self.direnv.overwritten_values_to_restore()?.iter().for_each(|(k, v)| {
+            self.add_env(&k, &v, true);
+        });
         Ok(())
     }