From 5002d87af471685b5bde1b866d3e4a74b87c2f88 Mon Sep 17 00:00:00 2001
From: Wind <WindSoilder@outlook.com>
Date: Thu, 10 Oct 2024 20:54:00 +0800
Subject: [PATCH] fix $env.FILE_PWD and $env.CURRENT_FILE inside `use` (#13958)

# Description
Fixes: #13425

Similar to `source-env`, `use` command should also remove `FILE_PWD` and
`CURRENT_FILE` after evaluating code block in the module file.

And user input can be a directory, in this case, we need to use the
return value of `find_in_dirs_env` carefully, so in case, I renamed
`maybe_file_path` to `maybe_file_path_or_dir` to emphasize it.

# User-Facing Changes
`$env.FILE_PWD` and `$env.CURRENT_FILE` will be more reliable to use.

# Tests + Formatting
Added 2 test cases.

# After Submitting
NaN
---
 crates/nu-cmd-lang/src/core_commands/use_.rs | 30 +++++++++----
 crates/nu-command/tests/commands/use_.rs     | 44 ++++++++++++++++++++
 2 files changed, 67 insertions(+), 7 deletions(-)

diff --git a/crates/nu-cmd-lang/src/core_commands/use_.rs b/crates/nu-cmd-lang/src/core_commands/use_.rs
index 8035cc81fe..d822e0c955 100644
--- a/crates/nu-cmd-lang/src/core_commands/use_.rs
+++ b/crates/nu-cmd-lang/src/core_commands/use_.rs
@@ -98,15 +98,21 @@ This command is a parser keyword. For details, check:
                     engine_state.get_span_contents(import_pattern.head.span),
                 );
 
-                let maybe_file_path = find_in_dirs_env(
+                let maybe_file_path_or_dir = find_in_dirs_env(
                     &module_arg_str,
                     engine_state,
                     caller_stack,
                     get_dirs_var_from_call(caller_stack, call),
                 )?;
-                let maybe_parent = maybe_file_path
-                    .as_ref()
-                    .and_then(|path| path.parent().map(|p| p.to_path_buf()));
+                // module_arg_str maybe a directory, in this case
+                // find_in_dirs_env returns a directory.
+                let maybe_parent = maybe_file_path_or_dir.as_ref().and_then(|path| {
+                    if path.is_dir() {
+                        Some(path.to_path_buf())
+                    } else {
+                        path.parent().map(|p| p.to_path_buf())
+                    }
+                });
 
                 let mut callee_stack = caller_stack
                     .gather_captures(engine_state, &block.captures)
@@ -118,9 +124,15 @@ This command is a parser keyword. For details, check:
                     callee_stack.add_env_var("FILE_PWD".to_string(), file_pwd);
                 }
 
-                if let Some(file_path) = maybe_file_path {
-                    let file_path = Value::string(file_path.to_string_lossy(), call.head);
-                    callee_stack.add_env_var("CURRENT_FILE".to_string(), file_path);
+                if let Some(path) = maybe_file_path_or_dir {
+                    let module_file_path = if path.is_dir() {
+                        // the existence of `mod.nu` is verified in parsing time
+                        // so it's safe to use it here.
+                        Value::string(path.join("mod.nu").to_string_lossy(), call.head)
+                    } else {
+                        Value::string(path.to_string_lossy(), call.head)
+                    };
+                    callee_stack.add_env_var("CURRENT_FILE".to_string(), module_file_path);
                 }
 
                 let eval_block = get_eval_block(engine_state);
@@ -130,6 +142,10 @@ This command is a parser keyword. For details, check:
 
                 // Merge the block's environment to the current stack
                 redirect_env(engine_state, caller_stack, &callee_stack);
+
+                // File-reative PWD is useless after eval file block.
+                caller_stack.remove_env_var(engine_state, "FILE_PWD");
+                caller_stack.remove_env_var(engine_state, "CURRENT_FILE");
             }
         } else {
             return Err(ShellError::GenericError {
diff --git a/crates/nu-command/tests/commands/use_.rs b/crates/nu-command/tests/commands/use_.rs
index 2895e1635f..2822be83c7 100644
--- a/crates/nu-command/tests/commands/use_.rs
+++ b/crates/nu-command/tests/commands/use_.rs
@@ -308,3 +308,47 @@ fn can_use_sub_subname_from_submodule() {
     let actual = nu!(inp);
     assert_eq!(actual.out, "bar")
 }
+
+#[test]
+fn test_use_with_printing_file_pwd() {
+    Playground::setup("use_with_printing_file_pwd", |dirs, nu| {
+        let file = dirs.test().join("mod.nu");
+        nu.with_files(&[FileWithContent(
+            file.as_os_str().to_str().unwrap(),
+            r#"
+                export-env {
+                    print $env.FILE_PWD
+                }
+            "#,
+        )]);
+
+        let actual = nu!(
+            cwd: dirs.test(),
+            "use ."
+        );
+
+        assert_eq!(actual.out, dirs.test().to_string_lossy());
+    });
+}
+
+#[test]
+fn test_use_with_printing_current_file() {
+    Playground::setup("use_with_printing_current_file", |dirs, nu| {
+        let file = dirs.test().join("mod.nu");
+        nu.with_files(&[FileWithContent(
+            file.as_os_str().to_str().unwrap(),
+            r#"
+                export-env {
+                    print $env.CURRENT_FILE
+                }
+            "#,
+        )]);
+
+        let actual = nu!(
+            cwd: dirs.test(),
+            "use ."
+        );
+
+        assert_eq!(actual.out, dirs.test().join("mod.nu").to_string_lossy());
+    });
+}