Add export-env eval to use command (#6572)

crates/nu-command/src/core_commands/use_.rs

@@ -1,4 +1,4 @@
-use nu_engine::eval_block;
+use nu_engine::{eval_block, find_in_dirs_env, redirect_env};
 use nu_protocol::ast::{Call, Expr, Expression, ImportPatternMember};
 use nu_protocol::engine::{Command, EngineState, Stack};
 use nu_protocol::{
@@ -35,9 +35,9 @@ impl Command for Use {
     fn run(
         &self,
         engine_state: &EngineState,
-        stack: &mut Stack,
+        caller_stack: &mut Stack,
         call: &Call,
-        _input: PipelineData,
+        input: PipelineData,
     ) -> Result<PipelineData, ShellError> {
         let import_pattern = if let Some(Expression {
             expr: Expr::ImportPattern(pat),
@@ -107,7 +107,7 @@ impl Command for Use {
 
                 let val = eval_block(
                     engine_state,
-                    stack,
+                    caller_stack,
                     block,
                     PipelineData::new(call.head),
                     false,
@@ -115,11 +115,50 @@ impl Command for Use {
                 )?
                 .into_value(call.head);
 
-                stack.add_env_var(name, val);
+                caller_stack.add_env_var(name, val);
+            }
+
+            // Evaluate the export-env block if there is one
+            if let Some(block_id) = module.env_block {
+                let block = engine_state.get_block(block_id);
+
+                // See if the module is a file
+                let module_arg_str = String::from_utf8_lossy(
+                    engine_state.get_span_contents(&import_pattern.head.span),
+                );
+                let maybe_parent = if let Some(path) =
+                    find_in_dirs_env(&module_arg_str, engine_state, caller_stack)?
+                {
+                    path.parent().map(|p| p.to_path_buf()).or(None)
+                } else {
+                    None
+                };
+
+                let mut callee_stack = caller_stack.gather_captures(&block.captures);
+
+                // If so, set the currently evaluated directory (file-relative PWD)
+                if let Some(parent) = maybe_parent {
+                    let file_pwd = Value::String {
+                        val: parent.to_string_lossy().to_string(),
+                        span: call.head,
+                    };
+                    callee_stack.add_env_var("FILE_PWD".to_string(), file_pwd);
+                }
+
+                // Run the block (discard the result)
+                let _ = eval_block(
+                    engine_state,
+                    &mut callee_stack,
+                    block,
+                    input,
+                    call.redirect_stdout,
+                    call.redirect_stderr,
+                )?;
+
+                // Merge the block's environment to the current stack
+                redirect_env(engine_state, caller_stack, &callee_stack);
             }
         } else {
-            // TODO: This is a workaround since call.positional[0].span points at 0 for some reason
-            // when this error is triggered
             return Err(ShellError::GenericError(
                 format!(
                     "Could not import from '{}'",

crates/nu-command/tests/commands/use_.rs

@@ -1,4 +1,5 @@
-use nu_test_support::fs::{AbsolutePath, Stub::FileWithContent};
+use nu_test_support::fs::AbsolutePath;
+use nu_test_support::fs::Stub::{FileWithContent, FileWithContentToBeTrimmed};
 use nu_test_support::nu;
 use nu_test_support::pipeline;
 use nu_test_support::playground::Playground;
@@ -63,3 +64,122 @@ fn use_keeps_doc_comments() {
         assert!(actual.out.contains("this is an x parameter"));
     })
 }
+
+#[test]
+fn use_eval_export_env() {
+    Playground::setup("use_eval_export_env", |dirs, sandbox| {
+        sandbox.with_files(vec![FileWithContentToBeTrimmed(
+            "spam.nu",
+            r#"
+                export-env { let-env FOO = 'foo' }
+            "#,
+        )]);
+
+        let inp = &[r#"use spam.nu"#, r#"$env.FOO"#];
+
+        let actual = nu!(cwd: dirs.test(), pipeline(&inp.join("; ")));
+
+        assert_eq!(actual.out, "foo");
+    })
+}
+
+#[test]
+fn use_eval_export_env_hide() {
+    Playground::setup("use_eval_export_env", |dirs, sandbox| {
+        sandbox.with_files(vec![FileWithContentToBeTrimmed(
+            "spam.nu",
+            r#"
+                export-env { hide-env FOO }
+            "#,
+        )]);
+
+        let inp = &[r#"let-env FOO = 'foo'"#, r#"use spam.nu"#, r#"$env.FOO"#];
+
+        let actual = nu!(cwd: dirs.test(), pipeline(&inp.join("; ")));
+
+        assert!(actual.err.contains("did you mean"));
+    })
+}
+
+#[test]
+fn use_do_cd() {
+    Playground::setup("use_do_cd", |dirs, sandbox| {
+        sandbox
+            .mkdir("test1/test2")
+            .with_files(vec![FileWithContentToBeTrimmed(
+                "test1/test2/spam.nu",
+                r#"
+                    export-env { cd test1/test2 }
+                "#,
+            )]);
+
+        let inp = &[r#"use test1/test2/spam.nu"#, r#"$env.PWD | path basename"#];
+
+        let actual = nu!(cwd: dirs.test(), pipeline(&inp.join("; ")));
+
+        assert_eq!(actual.out, "test2");
+    })
+}
+
+#[test]
+fn use_do_cd_file_relative() {
+    Playground::setup("use_do_cd_file_relative", |dirs, sandbox| {
+        sandbox
+            .mkdir("test1/test2")
+            .with_files(vec![FileWithContentToBeTrimmed(
+                "test1/test2/spam.nu",
+                r#"
+                    export-env { cd ($env.FILE_PWD | path join '..') }
+                "#,
+            )]);
+
+        let inp = &[r#"use test1/test2/spam.nu"#, r#"$env.PWD | path basename"#];
+
+        let actual = nu!(cwd: dirs.test(), pipeline(&inp.join("; ")));
+
+        assert_eq!(actual.out, "test1");
+    })
+}
+
+#[test]
+fn use_dont_cd_overlay() {
+    Playground::setup("use_dont_cd_overlay", |dirs, sandbox| {
+        sandbox
+            .mkdir("test1/test2")
+            .with_files(vec![FileWithContentToBeTrimmed(
+                "test1/test2/spam.nu",
+                r#"
+                    export-env {
+                        overlay new spam
+                        cd test1/test2
+                        overlay hide spam
+                    }
+                "#,
+            )]);
+
+        let inp = &[r#"use test1/test2/spam.nu"#, r#"$env.PWD | path basename"#];
+
+        let actual = nu!(cwd: dirs.test(), pipeline(&inp.join("; ")));
+
+        assert_eq!(actual.out, "use_dont_cd_overlay");
+    })
+}
+
+#[test]
+fn use_export_env_combined() {
+    Playground::setup("use_is_scoped", |dirs, sandbox| {
+        sandbox.with_files(vec![FileWithContentToBeTrimmed(
+            "spam.nu",
+            r#"
+                alias bar = foo
+                export-env { let-env FOO = bar }
+                def foo [] { 'foo' }
+            "#,
+        )]);
+
+        let inp = &[r#"use spam.nu"#, r#"$env.FOO"#];
+
+        let actual = nu!(cwd: dirs.test(), pipeline(&inp.join("; ")));
+        assert_eq!(actual.out, "foo");
+    })
+}