From fefd5fef12702e40911dffd7bc7b198045b2ff1a Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jakub=20=C5=BD=C3=A1dn=C3=ADk?= <kubouch@gmail.com>
Date: Fri, 4 Feb 2022 20:02:03 +0200
Subject: [PATCH] Allow def-env to hide environment variables (#921)

---
 crates/nu-engine/src/eval.rs           | 16 ++++++++++++-
 crates/nu-protocol/src/engine/stack.rs | 31 ++++++++++++++++++++++++++
 src/tests/test_engine.rs               | 16 +++++++++++++
 3 files changed, 62 insertions(+), 1 deletion(-)

diff --git a/crates/nu-engine/src/eval.rs b/crates/nu-engine/src/eval.rs
index 6e7f037316..c827ebf341 100644
--- a/crates/nu-engine/src/eval.rs
+++ b/crates/nu-engine/src/eval.rs
@@ -133,14 +133,28 @@ fn eval_call(
                 }
             }
         }
+
         let result = eval_block(engine_state, &mut callee_stack, block, input);
+
         if block.redirect_env {
+            let caller_env_vars = caller_stack.get_env_var_names(engine_state);
+
+            // remove env vars that are present in the caller but not in the callee
+            // (the callee hid them)
+            for var in caller_env_vars.iter() {
+                if !callee_stack.has_env_var(engine_state, var) {
+                    caller_stack.remove_env_var(engine_state, var);
+                }
+            }
+
+            // add new env vars from callee to caller
             for env_vars in callee_stack.env_vars {
                 for (var, value) in env_vars {
-                    caller_stack.add_env_var(var, value)
+                    caller_stack.add_env_var(var, value);
                 }
             }
         }
+
         result
     } else {
         // We pass caller_stack here with the knowledge that internal commands
diff --git a/crates/nu-protocol/src/engine/stack.rs b/crates/nu-protocol/src/engine/stack.rs
index 07faaad85c..94c96fc8b0 100644
--- a/crates/nu-protocol/src/engine/stack.rs
+++ b/crates/nu-protocol/src/engine/stack.rs
@@ -140,6 +140,23 @@ impl Stack {
         result
     }
 
+    /// Same as get_env_vars, but returns only the names as a HashSet
+    pub fn get_env_var_names(&self, engine_state: &EngineState) -> HashSet<String> {
+        let mut result: HashSet<String> = engine_state
+            .env_vars
+            .keys()
+            .filter(|k| !self.env_hidden.contains(*k))
+            .cloned()
+            .collect();
+
+        for scope in &self.env_vars {
+            let scope_keys: HashSet<String> = scope.keys().cloned().collect();
+            result.extend(scope_keys);
+        }
+
+        result
+    }
+
     pub fn get_env_var(&self, engine_state: &EngineState, name: &str) -> Option<Value> {
         for scope in self.env_vars.iter().rev() {
             if let Some(v) = scope.get(name) {
@@ -154,6 +171,20 @@ impl Stack {
         }
     }
 
+    pub fn has_env_var(&self, engine_state: &EngineState, name: &str) -> bool {
+        for scope in self.env_vars.iter().rev() {
+            if scope.contains_key(name) {
+                return true;
+            }
+        }
+
+        if self.env_hidden.contains(name) {
+            false
+        } else {
+            engine_state.env_vars.contains_key(name)
+        }
+    }
+
     pub fn remove_env_var(&mut self, engine_state: &EngineState, name: &str) -> Option<Value> {
         for scope in self.env_vars.iter_mut().rev() {
             if let Some(v) = scope.remove(name) {
diff --git a/src/tests/test_engine.rs b/src/tests/test_engine.rs
index a1df9f1641..aae132b880 100644
--- a/src/tests/test_engine.rs
+++ b/src/tests/test_engine.rs
@@ -207,6 +207,22 @@ fn not_def_env() -> TestResult {
     )
 }
 
+#[test]
+fn def_env_hiding_something() -> TestResult {
+    fail_test(
+        r#"let-env FOO = "foo"; def-env bob [] { hide FOO }; bob; $env.FOO"#,
+        "did you mean",
+    )
+}
+
+#[test]
+fn def_env_then_hide() -> TestResult {
+    fail_test(
+        r#"def-env bob [] { let-env BOB = "bob" }; def-env un-bob [] { hide BOB }; bob; un-bob; $env.BOB"#,
+        "did you mean",
+    )
+}
+
 #[test]
 fn export_def_env() -> TestResult {
     run_test(