From 90ba39184ae2c26218d0e64fb978b3f4e47aed33 Mon Sep 17 00:00:00 2001
From: pwygab <88221256+merelymyself@users.noreply.github.com>
Date: Thu, 13 Oct 2022 17:04:34 +0800
Subject: [PATCH] allow for `$in` to affect environment (#6649)

* allow for `$in` to affect environment

* fix for vars, overlays, env_hidden

* fmt

* carry over variables properly

* add test

* modify name, remove redundant

* fmt
---
 crates/nu-command/src/filters/collect.rs | 34 +++++++++++++++++++-----
 crates/nu-command/tests/commands/cd.rs   | 14 ++++++++++
 crates/nu-parser/src/parser.rs           |  9 +++++++
 3 files changed, 51 insertions(+), 6 deletions(-)

diff --git a/crates/nu-command/src/filters/collect.rs b/crates/nu-command/src/filters/collect.rs
index 76b4428fcb..bd78a34dbb 100644
--- a/crates/nu-command/src/filters/collect.rs
+++ b/crates/nu-command/src/filters/collect.rs
@@ -1,4 +1,4 @@
-use nu_engine::{eval_block, CallExt};
+use nu_engine::{eval_block, redirect_env, CallExt};
 use nu_protocol::ast::Call;
 use nu_protocol::engine::{CaptureBlock, Command, EngineState, Stack};
 use nu_protocol::{
@@ -20,6 +20,11 @@ impl Command for Collect {
                 SyntaxShape::Block(Some(vec![SyntaxShape::Any])),
                 "the block to run once the stream is collected",
             )
+            .switch(
+                "keep-env",
+                "let the block affect environment variables",
+                None,
+            )
             .category(Category::Filters)
     }
 
@@ -37,26 +42,43 @@ impl Command for Collect {
         let capture_block: CaptureBlock = call.req(engine_state, stack, 0)?;
 
         let block = engine_state.get_block(capture_block.block_id).clone();
-        let mut stack = stack.captures_to_stack(&capture_block.captures);
+        let mut stack_captures = stack.captures_to_stack(&capture_block.captures);
 
         let metadata = input.metadata();
         let input: Value = input.into_value(call.head);
 
+        let mut saved_positional = None;
         if let Some(var) = block.signature.get_positional(0) {
             if let Some(var_id) = &var.var_id {
-                stack.add_var(*var_id, input.clone());
+                stack_captures.add_var(*var_id, input.clone());
+                saved_positional = Some(*var_id);
             }
         }
 
-        eval_block(
+        let result = eval_block(
             engine_state,
-            &mut stack,
+            &mut stack_captures,
             &block,
             input.into_pipeline_data(),
             call.redirect_stdout,
             call.redirect_stderr,
         )
-        .map(|x| x.set_metadata(metadata))
+        .map(|x| x.set_metadata(metadata));
+
+        if call.has_flag("keep-env") {
+            redirect_env(engine_state, stack, &stack_captures);
+            // for when we support `data | let x = $in;`
+            // remove the variables added earlier
+            for var_id in capture_block.captures.keys() {
+                stack_captures.vars.remove(var_id);
+            }
+            if let Some(u) = saved_positional {
+                stack_captures.vars.remove(&u);
+            }
+            // add any new variables to the stack
+            stack.vars.extend(stack_captures.vars);
+        }
+        result
     }
 
     fn examples(&self) -> Vec<Example> {
diff --git a/crates/nu-command/tests/commands/cd.rs b/crates/nu-command/tests/commands/cd.rs
index c23e8e456f..89ddbfed3e 100644
--- a/crates/nu-command/tests/commands/cd.rs
+++ b/crates/nu-command/tests/commands/cd.rs
@@ -3,6 +3,20 @@ use nu_test_support::nu;
 use nu_test_support::playground::Playground;
 use std::path::PathBuf;
 
+#[test]
+fn cd_works_with_in_var() {
+    Playground::setup("cd_test_1", |dirs, _| {
+        let actual = nu!(
+            cwd: dirs.root(),
+            r#"
+                "cd_test_1" | cd $in; $env.PWD | path split | last
+            "#
+        );
+
+        assert_eq!("cd_test_1", actual.out);
+    })
+}
+
 // FIXME: jt: needs more work
 #[ignore]
 #[test]
diff --git a/crates/nu-parser/src/parser.rs b/crates/nu-parser/src/parser.rs
index 42d0b1e6c4..662b20a86d 100644
--- a/crates/nu-parser/src/parser.rs
+++ b/crates/nu-parser/src/parser.rs
@@ -5367,6 +5367,15 @@ fn wrap_expr_with_collect(working_set: &mut StateWorkingSet, expr: &Expression)
             custom_completion: None,
         }));
 
+        output.push(Argument::Named((
+            Spanned {
+                item: "keep-env".to_string(),
+                span: Span::new(0, 0),
+            },
+            None,
+            None,
+        )));
+
         // The containing, synthetic call to `collect`.
         // We don't want to have a real span as it will confuse flattening
         // The args are where we'll get the real info