Bring module's environment when activating overlay (#6425)

crates/nu-command/src/core_commands/overlay/use_.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::{Command, EngineState, Stack};
 use nu_protocol::{Category, Example, PipelineData, ShellError, Signature, Spanned, SyntaxShape};
@@ -49,11 +49,11 @@ impl Command for OverlayUse {
     fn run(
         &self,
         engine_state: &EngineState,
-        stack: &mut Stack,
+        caller_stack: &mut Stack,
         call: &Call,
-        _input: PipelineData,
+        input: PipelineData,
     ) -> Result<PipelineData, ShellError> {
-        let name_arg: Spanned<String> = call.req(engine_state, stack, 0)?;
+        let name_arg: Spanned<String> = call.req(engine_state, caller_stack, 0)?;
 
         let (overlay_name, overlay_name_span) = if let Some(kw_expression) = call.positional_nth(1)
         {
@@ -96,10 +96,10 @@ impl Command for OverlayUse {
         if let Some(overlay_id) = engine_state.find_overlay(overlay_name.as_bytes()) {
             let old_module_id = engine_state.get_overlay(overlay_id).origin;
 
-            stack.add_overlay(overlay_name.clone());
+            caller_stack.add_overlay(overlay_name.clone());
 
             if let Some(new_module_id) = engine_state.find_module(overlay_name.as_bytes(), &[]) {
-                if !stack.has_env_overlay(&overlay_name, engine_state)
+                if !caller_stack.has_env_overlay(&overlay_name, engine_state)
                     || (old_module_id != new_module_id)
                 {
                     // Add environment variables only if:
@@ -118,7 +118,7 @@ impl Command for OverlayUse {
 
                         let val = eval_block(
                             engine_state,
-                            stack,
+                            caller_stack,
                             block,
                             PipelineData::new(call.head),
                             false,
@@ -126,7 +126,24 @@ impl Command for OverlayUse {
                         )?
                         .into_value(call.head);
 
-                        stack.add_env_var(name, val);
+                        caller_stack.add_env_var(name, val);
+                    }
+
+                    // Evaluate the export-env block (if any) and keep its environment
+                    if let Some(block_id) = module.env_block {
+                        let block = engine_state.get_block(block_id);
+                        let mut callee_stack = caller_stack.gather_captures(&block.captures);
+
+                        let _ = eval_block(
+                            engine_state,
+                            &mut callee_stack,
+                            block,
+                            input,
+                            call.redirect_stdout,
+                            call.redirect_stderr,
+                        );
+
+                        redirect_env(engine_state, caller_stack, &callee_stack);
                     }
                 }
             }

crates/nu-parser/src/parse_keywords.rs

@@ -5,7 +5,7 @@ use nu_protocol::{
         ImportPatternMember, Pipeline,
     },
     engine::{StateWorkingSet, DEFAULT_OVERLAY_NAME},
-    span, Exportable, Module, PositionalArg, Span, Spanned, SyntaxShape, Type,
+    span, BlockId, Exportable, Module, PositionalArg, Span, Spanned, SyntaxShape, Type,
 };
 use std::collections::HashSet;
 use std::path::{Path, PathBuf};
@@ -1215,12 +1215,11 @@ pub fn parse_export_env(
     working_set: &mut StateWorkingSet,
     spans: &[Span],
     expand_aliases_denylist: &[usize],
-) -> (Pipeline, Option<ParseError>) {
+) -> (Pipeline, Option<BlockId>, Option<ParseError>) {
-    // Just used to be allowed inside modules, otherwise does nothing in the parser.
-
     if !spans.is_empty() && working_set.get_span_contents(spans[0]) != b"export-env" {
         return (
             garbage_pipeline(spans),
+            None,
             Some(ParseError::UnknownState(
                 "internal error: Wrong call name for 'export-env' command".into(),
                 span(spans),
@@ -1231,6 +1230,7 @@ pub fn parse_export_env(
     if spans.len() < 2 {
         return (
             garbage_pipeline(spans),
+            None,
             Some(ParseError::MissingPositional(
                 "block".into(),
                 span(spans),
@@ -1265,6 +1265,7 @@ pub fn parse_export_env(
                         ty: output,
                         custom_completion: None,
                     }]),
+                    None,
                     err,
                 );
             }
@@ -1274,6 +1275,7 @@ pub fn parse_export_env(
         None => {
             return (
                 garbage_pipeline(spans),
+                None,
                 Some(ParseError::UnknownState(
                     "internal error: 'export-env' declaration not found".into(),
                     span(spans),
@@ -1282,6 +1284,30 @@ pub fn parse_export_env(
         }
     };
 
+    let block_id = if let Some(block) = call.positional_nth(0) {
+        if let Some(block_id) = block.as_block() {
+            block_id
+        } else {
+            return (
+                garbage_pipeline(spans),
+                None,
+                Some(ParseError::UnknownState(
+                    "internal error: 'export-env' block is not a block".into(),
+                    block.span,
+                )),
+            );
+        }
+    } else {
+        return (
+            garbage_pipeline(spans),
+            None,
+            Some(ParseError::UnknownState(
+                "internal error: 'export-env' block is missing".into(),
+                span(spans),
+            )),
+        );
+    };
+
     let pipeline = Pipeline::from_vec(vec![Expression {
         expr: Expr::Call(call),
         span: span(spans),
@@ -1289,7 +1315,7 @@ pub fn parse_export_env(
         custom_completion: None,
     }]);
 
-    (pipeline, None)
+    (pipeline, Some(block_id), None)
 }
 
 pub fn parse_module_block(
@@ -1388,11 +1414,19 @@ pub fn parse_module_block(
 
                         (pipe, err)
                     }
-                    b"export-env" => parse_export_env(
+                    b"export-env" => {
-                        working_set,
+                        let (pipe, maybe_env_block, err) = parse_export_env(
-                        &pipeline.commands[0].parts,
+                            working_set,
-                        expand_aliases_denylist,
+                            &pipeline.commands[0].parts,
-                    ),
+                            expand_aliases_denylist,
+                        );
+
+                        if let Some(block_id) = maybe_env_block {
+                            module.add_env_block(block_id);
+                        }
+
+                        (pipe, err)
+                    }
                     _ => (
                         garbage_pipeline(&pipeline.commands[0].parts),
                         Some(ParseError::ExpectedKeyword(

crates/nu-protocol/src/module.rs

@@ -11,6 +11,7 @@ pub struct Module {
     pub decls: IndexMap<Vec<u8>, DeclId>,
     pub aliases: IndexMap<Vec<u8>, AliasId>,
     pub env_vars: IndexMap<Vec<u8>, BlockId>,
+    pub env_block: Option<BlockId>,
     pub span: Option<Span>,
 }
 
@@ -20,6 +21,7 @@ impl Module {
             decls: IndexMap::new(),
             aliases: IndexMap::new(),
             env_vars: IndexMap::new(),
+            env_block: None,
             span: None,
         }
     }
@@ -29,6 +31,7 @@ impl Module {
             decls: IndexMap::new(),
             aliases: IndexMap::new(),
             env_vars: IndexMap::new(),
+            env_block: None,
             span: Some(span),
         }
     }
@@ -45,6 +48,10 @@ impl Module {
         self.env_vars.insert(name, block_id)
     }
 
+    pub fn add_env_block(&mut self, block_id: BlockId) {
+        self.env_block = Some(block_id);
+    }
+
     pub fn extend(&mut self, other: &Module) {
         self.decls.extend(other.decls.clone());
         self.env_vars.extend(other.env_vars.clone());

tests/overlays/mod.rs

@@ -736,3 +736,34 @@ fn overlay_remove_and_add_renamed_overlay() {
     assert_eq!(actual.out, "foo");
     assert_eq!(actual_repl.out, "foo");
 }
+
+#[test]
+fn overlay_use_export_env() {
+    let inp = &[
+        r#"module spam { export-env { let-env FOO = 'foo' } }"#,
+        r#"overlay use spam"#,
+        r#"$env.FOO"#,
+    ];
+
+    let actual = nu!(cwd: "tests/overlays", pipeline(&inp.join("; ")));
+    let actual_repl = nu!(cwd: "tests/overlays", nu_repl_code(inp));
+
+    assert_eq!(actual.out, "foo");
+    assert_eq!(actual_repl.out, "foo");
+}
+
+#[test]
+fn overlay_use_export_env_hide() {
+    let inp = &[
+        r#"let-env FOO = 'foo'"#,
+        r#"module spam { export-env { hide-env FOO } }"#,
+        r#"overlay use spam"#,
+        r#"$env.FOO"#,
+    ];
+
+    let actual = nu!(cwd: "tests/overlays", pipeline(&inp.join("; ")));
+    let actual_repl = nu!(cwd: "tests/overlays", nu_repl_code(inp));
+
+    assert!(actual.err.contains("did you mean"));
+    assert!(actual_repl.err.contains("did you mean"));
+}