From 8dce6912c28fd4a7a82b36ed49db500c8ab6841f Mon Sep 17 00:00:00 2001
From: Charlie Marsh <charlie.r.marsh@gmail.com>
Date: Tue, 3 Sep 2024 23:26:47 -0400
Subject: [PATCH] Reflect exit code in uv tool run and uv run

---
 crates/uv/src/commands/mod.rs         |  4 ++++
 crates/uv/src/commands/project/run.rs | 16 +++++++++++++---
 crates/uv/src/commands/tool/run.rs    | 16 +++++++++++++---
 crates/uv/tests/run.rs                | 19 +++++++++++++++++++
 crates/uv/tests/tool_run.rs           |  2 +-
 5 files changed, 50 insertions(+), 7 deletions(-)

diff --git a/crates/uv/src/commands/mod.rs b/crates/uv/src/commands/mod.rs
index d013d48b5e65..cafc8276d7d2 100644
--- a/crates/uv/src/commands/mod.rs
+++ b/crates/uv/src/commands/mod.rs
@@ -80,6 +80,9 @@ pub(crate) enum ExitStatus {
 
     /// The command failed with an unexpected error.
     Error,
+
+    /// The command's exit status is propagated from an external command.
+    External(u8),
 }
 
 impl From<ExitStatus> for ExitCode {
@@ -88,6 +91,7 @@ impl From<ExitStatus> for ExitCode {
             ExitStatus::Success => Self::from(0),
             ExitStatus::Failure => Self::from(1),
             ExitStatus::Error => Self::from(2),
+            ExitStatus::External(code) => Self::from(code),
         }
     }
 }
diff --git a/crates/uv/src/commands/project/run.rs b/crates/uv/src/commands/project/run.rs
index 4047f6c48483..a2094f3cc06d 100644
--- a/crates/uv/src/commands/project/run.rs
+++ b/crates/uv/src/commands/project/run.rs
@@ -679,10 +679,20 @@ pub(crate) async fn run(
     let status = handle.wait().await.context("Child process disappeared")?;
 
     // Exit based on the result of the command
-    // TODO(zanieb): Do we want to exit with the code of the child process? Probably.
-    if status.success() {
-        Ok(ExitStatus::Success)
+    if let Some(code) = status.code() {
+        debug!("Command exited with code: {code}");
+        if let Ok(code) = u8::try_from(code) {
+            Ok(ExitStatus::External(code))
+        } else {
+            #[allow(clippy::exit)]
+            std::process::exit(code);
+        }
     } else {
+        #[cfg(unix)]
+        {
+            use std::os::unix::process::ExitStatusExt;
+            debug!("Command exited with signal: {:?}", status.signal());
+        }
         Ok(ExitStatus::Failure)
     }
 }
diff --git a/crates/uv/src/commands/tool/run.rs b/crates/uv/src/commands/tool/run.rs
index 55d9042386b4..8f65c561e37a 100644
--- a/crates/uv/src/commands/tool/run.rs
+++ b/crates/uv/src/commands/tool/run.rs
@@ -221,10 +221,20 @@ pub(crate) async fn run(
     let status = handle.wait().await.context("Child process disappeared")?;
 
     // Exit based on the result of the command
-    // TODO(zanieb): Do we want to exit with the code of the child process? Probably.
-    if status.success() {
-        Ok(ExitStatus::Success)
+    if let Some(code) = status.code() {
+        debug!("Command exited with code: {code}");
+        if let Ok(code) = u8::try_from(code) {
+            Ok(ExitStatus::External(code))
+        } else {
+            #[allow(clippy::exit)]
+            std::process::exit(code);
+        }
     } else {
+        #[cfg(unix)]
+        {
+            use std::os::unix::process::ExitStatusExt;
+            debug!("Command exited with signal: {:?}", status.signal());
+        }
         Ok(ExitStatus::Failure)
     }
 }
diff --git a/crates/uv/tests/run.rs b/crates/uv/tests/run.rs
index e393260598c9..c12fd5cd6442 100644
--- a/crates/uv/tests/run.rs
+++ b/crates/uv/tests/run.rs
@@ -1774,3 +1774,22 @@ fn run_compiled_python_file() -> Result<()> {
 
     Ok(())
 }
+
+#[test]
+fn run_exit_code() -> Result<()> {
+    let context = TestContext::new("3.12");
+
+    let test_script = context.temp_dir.child("script.py");
+    test_script.write_str(indoc! { r#"
+        # /// script
+        # requires-python = ">=3.11"
+        # ///
+
+        exit(42)
+       "#
+    })?;
+
+    context.run().arg("script.py").assert().code(42);
+
+    Ok(())
+}
diff --git a/crates/uv/tests/tool_run.rs b/crates/uv/tests/tool_run.rs
index 2012edb9e3cd..14608fc2fa6e 100644
--- a/crates/uv/tests/tool_run.rs
+++ b/crates/uv/tests/tool_run.rs
@@ -258,7 +258,7 @@ fn tool_run_warn_executable_not_in_from() {
         .env("UV_TOOL_DIR", tool_dir.as_os_str())
         .env("XDG_BIN_HOME", bin_dir.as_os_str()), @r###"
     success: false
-    exit_code: 1
+    exit_code: 2
     ----- stdout -----
 
     ----- stderr -----