From 3fbbe6e322d6c5ae5dae2e8b9c35787cd4038a82 Mon Sep 17 00:00:00 2001
From: Sam Hedin <sam.hedin@gmail.com>
Date: Sun, 14 Jun 2020 10:57:16 +0200
Subject: [PATCH] Refactor restore function

---
 .../src/env/directory_specific_environment.rs | 50 +++++++------------
 crates/nu-cli/src/env/environment.rs          | 23 ++++++---
 2 files changed, 34 insertions(+), 39 deletions(-)

diff --git a/crates/nu-cli/src/env/directory_specific_environment.rs b/crates/nu-cli/src/env/directory_specific_environment.rs
index 9edc2f41c9..c8513d7ddc 100644
--- a/crates/nu-cli/src/env/directory_specific_environment.rs
+++ b/crates/nu-cli/src/env/directory_specific_environment.rs
@@ -51,31 +51,25 @@ impl DirectorySpecificEnvironment {
         let current_dir = std::env::current_dir()?;
         let mut working_dir = Some(current_dir.as_path());
 
-
-        let mut file = OpenOptions::new()
-            .write(true)
-            .append(true)
-            .create(true)
-            .open("restore.txt").unwrap();
-
-        write!(&mut file, "about to restore: {:?}\n", self.overwritten_env_values).unwrap();
-
+        let mut new_overwritten_env_values: IndexMap<PathBuf, IndexMap<String, OsString>> =
+            IndexMap::new();
         let mut keyvals_to_restore = IndexMap::new();
-        let mut new_overwritten_env_values = IndexMap::new();
-        //If we are not in wdir or its subdir, remove its vals
-        self.overwritten_env_values
-            .iter()
-            .for_each(|(directory, keyvals)| {
-                while let Some(wdir) = working_dir {
-                    if &wdir == directory {
-                        keyvals.iter().for_each(|(k, v)| {
-                            keyvals_to_restore.insert(k.clone(), v.to_str().unwrap().to_string());
-                        });
-                    }
-                    working_dir = working_dir.expect("This directory has no parent").parent();
-                }
-                new_overwritten_env_values.insert(directory.clone(), keyvals.clone());
-            });
+
+        while let Some(wdir) = working_dir {
+            if let Some(val) = self.overwritten_env_values.get(wdir) {
+                new_overwritten_env_values.insert(wdir.to_path_buf(), val.clone());
+            }
+            working_dir = working_dir.unwrap().parent();
+        }
+
+        for (dir, keyvals) in &self.overwritten_env_values {
+            if !new_overwritten_env_values.contains_key(dir) {
+                keyvals.iter().for_each(|(k, v)| {
+                    keyvals_to_restore.insert(k.clone(), v.to_str().unwrap().to_string());
+                });
+                self.added_env_vars.remove(dir);
+            }
+        }
 
         self.overwritten_env_values = new_overwritten_env_values;
         Ok(keyvals_to_restore)
@@ -121,14 +115,6 @@ impl DirectorySpecificEnvironment {
                     });
             }
 
-            let mut file = OpenOptions::new()
-                .write(true)
-                .append(true)
-                .create(true)
-                .open("restore.txt").unwrap();
-
-            write!(&mut file, "overwritten: {:?}\n\n", self.overwritten_env_values).unwrap();
-
             working_dir = working_dir //Keep going up in the directory structure with .parent()
                 .expect("This directory has no parent")
                 .parent();
diff --git a/crates/nu-cli/src/env/environment.rs b/crates/nu-cli/src/env/environment.rs
index d2b1520531..b84cdcd93d 100644
--- a/crates/nu-cli/src/env/environment.rs
+++ b/crates/nu-cli/src/env/environment.rs
@@ -2,8 +2,9 @@ use crate::data::config::Conf;
 use crate::env::directory_specific_environment::*;
 use indexmap::{indexmap, IndexSet};
 use nu_protocol::{UntaggedValue, Value};
+use std::io::Write;
 use std::ffi::OsString;
-use std::fmt::Debug;
+use std::{fs::OpenOptions, fmt::Debug};
 
 pub trait Env: Debug + Send {
     fn env(&self) -> Option<Value>;
@@ -63,12 +64,20 @@ impl Environment {
             self.remove_env(&k);
         });
 
-        // self.direnv
-        //     .overwritten_values_to_restore()?
-        //     .iter()
-        //     .for_each(|(k, v)| {
-        //         self.add_env(&k, &v, true);
-        //     });
+        // let mut file = OpenOptions::new()
+        //     .write(true)
+        //     .append(true)
+        //     .create(true)
+        //     .open("restore.txt")
+        //     .unwrap();
+
+        self.direnv
+            .overwritten_values_to_restore()?
+            .iter()
+            .for_each(|(k, v)| {
+                // write!(&mut file, "restoring: {:?}\n", (k, v)).unwrap();
+                self.add_env(&k, &v, true);
+            });
 
         self.direnv.env_vars_to_add()?.iter().for_each(|(k, v)| {
             self.add_env(&k, &v, true);