Allow keeping selected environment variables from removed overlay (#6007)

* Allow keeping selected env from removed overlay

* Remove some duplicate code

* Change --keep-all back to --keep-custom

Because, apparently, you cannot have a named flag called --keep-all,
otherwise tests fail?

* Fix missing line and wrong test value

crates/nu-command/src/core_commands/overlay/remove.rs

@@ -1,9 +1,7 @@
 use nu_engine::CallExt;
 use nu_protocol::ast::Call;
 use nu_protocol::engine::{Command, EngineState, Stack};
-use nu_protocol::{
+use nu_protocol::{Category, Example, PipelineData, ShellError, Signature, Spanned, SyntaxShape};
-    Category, Example, PipelineData, ShellError, Signature, Spanned, SyntaxShape, Value,
-};
 
 #[derive(Clone)]
 pub struct OverlayRemove;
@@ -22,9 +20,15 @@ impl Command for OverlayRemove {
             .optional("name", SyntaxShape::String, "Overlay to remove")
             .switch(
                 "keep-custom",
-                "Keep newly added symbols within the next activated overlay",
+                "Keep all newly added symbols within the next activated overlay",
                 Some('k'),
             )
+            .named(
+                "keep-env",
+                SyntaxShape::List(Box::new(SyntaxShape::String)),
+                "List of environment variables to keep from the removed overlay",
+                Some('e'),
+            )
             .category(Category::Core)
     }
 
@@ -60,30 +64,44 @@ impl Command for OverlayRemove {
             ));
         }
 
-        if call.has_flag("keep-custom") {
+        let keep_env: Option<Vec<Spanned<String>>> =
+            call.get_flag(engine_state, stack, "keep-env")?;
+
+        let env_vars_to_keep = if call.has_flag("keep-custom") {
             if let Some(overlay_id) = engine_state.find_overlay(overlay_name.item.as_bytes()) {
                 let overlay_frame = engine_state.get_overlay(overlay_id);
                 let origin_module = engine_state.get_module(overlay_frame.origin);
 
-                let env_vars_to_keep: Vec<(String, Value)> = stack
+                stack
                     .get_overlay_env_vars(engine_state, &overlay_name.item)
                     .into_iter()
                     .filter(|(name, _)| !origin_module.has_env_var(name.as_bytes()))
-                    .collect();
+                    .collect()
-
-                stack.remove_overlay(&overlay_name.item);
-
-                for (name, val) in env_vars_to_keep {
-                    stack.add_env_var(name, val);
-                }
             } else {
                 return Err(ShellError::OverlayNotFoundAtRuntime(
                     overlay_name.item,
                     overlay_name.span,
                 ));
             }
+        } else if let Some(env_var_names_to_keep) = keep_env {
+            let mut env_vars_to_keep = vec![];
+
+            for name in env_var_names_to_keep.into_iter() {
+                match stack.get_env_var(engine_state, &name.item) {
+                    Some(val) => env_vars_to_keep.push((name.item, val.clone())),
+                    None => return Err(ShellError::EnvVarNotFoundAtRuntime(name.item, name.span)),
+                }
+            }
+
+            env_vars_to_keep
         } else {
+            vec![]
+        };
+
         stack.remove_overlay(&overlay_name.item);
+
+        for (name, val) in env_vars_to_keep {
+            stack.add_env_var(name, val);
         }
 
         Ok(PipelineData::new(call.head))
@@ -112,6 +130,13 @@ impl Command for OverlayRemove {
     overlay remove"#,
                 result: None,
             },
+            Example {
+                description: "Keep the current working directory when removing an overlay",
+                example: r#"overlay new spam
+    cd some-dir
+    overlay remove --keep-env [ PWD ] spam"#,
+                result: None,
+            },
         ]
     }
 }

crates/nu-protocol/src/value/from_value.rs

@@ -238,6 +238,35 @@ impl FromValue for Vec<String> {
     }
 }
 
+impl FromValue for Vec<Spanned<String>> {
+    fn from_value(v: &Value) -> Result<Self, ShellError> {
+        // FIXME: we may want to fail a little nicer here
+        match v {
+            Value::List { vals, .. } => vals
+                .iter()
+                .map(|val| match val {
+                    Value::String { val, span } => Ok(Spanned {
+                        item: val.clone(),
+                        span: *span,
+                    }),
+                    c => Err(ShellError::CantConvert(
+                        "string".into(),
+                        c.get_type().to_string(),
+                        c.span()?,
+                        None,
+                    )),
+                })
+                .collect::<Result<Vec<Spanned<String>>, ShellError>>(),
+            v => Err(ShellError::CantConvert(
+                "string".into(),
+                v.get_type().to_string(),
+                v.span()?,
+                None,
+            )),
+        }
+    }
+}
+
 impl FromValue for Vec<bool> {
     fn from_value(v: &Value) -> Result<Self, ShellError> {
         match v {

tests/overlays/mod.rs

@@ -507,3 +507,19 @@ fn overlay_new() {
     assert_eq!(actual.out, "spam");
     assert_eq!(actual_repl.out, "spam");
 }
+
+#[test]
+fn overlay_keep_pwd() {
+    let inp = &[
+        r#"overlay new spam"#,
+        r#"cd samples"#,
+        r#"overlay remove --keep-env [ PWD ] spam"#,
+        r#"$env.PWD | path basename"#,
+    ];
+
+    let actual = nu!(cwd: "tests/overlays", pipeline(&inp.join("; ")));
+    let actual_repl = nu_repl("tests/overlays", inp);
+
+    assert_eq!(actual.out, "samples");
+    assert_eq!(actual_repl.out, "samples");
+}