Use suffix on uvx binary when searching for uv binary (#12923)

This is a rebased and updated version of #11925 based on my review (I
didn't have permission to push to their branch).

For posterity I've preserved their commits but my final commit
essentially rewrites the whole thing anyway.

Fixes #11637

---------

Co-authored-by: Chris Lieb <[email protected]>

Author: Gankra

crates/uv/src/bin/uvx.rs

@@ -1,4 +1,5 @@
 use std::convert::Infallible;
+use std::path::{Path, PathBuf};
 use std::{
     ffi::OsString,
     process::{Command, ExitCode, ExitStatus},
@@ -22,6 +23,57 @@ fn exec_spawn(cmd: &mut Command) -> std::io::Result<Infallible> {
     }
 }
 
+/// Assuming the binary is called something like `[email protected](.exe)`, compute the `@1.2.3(.exe)` part
+/// so that we can preferentially find `[email protected](.exe)`, for folks who like managing multiple
+/// installs in this way.
+fn get_uvx_suffix(current_exe: &Path) -> Option<&str> {
+    let os_file_name = current_exe.file_name()?;
+    let file_name_str = os_file_name.to_str()?;
+    file_name_str.strip_prefix("uvx")
+}
+
+/// Gets the path to `uv`, given info about `uvx`
+fn get_uv_path(current_exe_parent: &Path, uvx_suffix: Option<&str>) -> std::io::Result<PathBuf> {
+    // First try to find a matching suffixed `uv`, e.g. `[email protected](.exe)`
+    let uv_with_suffix = uvx_suffix.map(|suffix| current_exe_parent.join(format!("uv{suffix}")));
+    if let Some(uv_with_suffix) = &uv_with_suffix {
+        #[allow(clippy::print_stderr, reason = "printing a very rare warning")]
+        match uv_with_suffix.try_exists() {
+            Ok(true) => return Ok(uv_with_suffix.to_owned()),
+            Ok(false) => { /* definitely not there, proceed to fallback */ }
+            Err(err) => {
+                // We don't know if `[email protected]` exists, something errored when checking.
+                // We *could* blindly use `[email protected]` in this case, as the code below does, however
+                // in this extremely narrow corner case it's *probably* better to default to `uv`,
+                // since we don't want to mess up existing users who weren't using suffixes?
+                eprintln!(
+                    "warning: failed to determine if `{}` exists, trying `uv` instead: {err}",
+                    uv_with_suffix.display()
+                );
+            }
+        }
+    }
+
+    // Then just look for good ol' `uv`
+    let uv = current_exe_parent.join(format!("uv{}", std::env::consts::EXE_SUFFIX));
+    // If we are sure the `uv` binary does not exist, display a clearer error message.
+    // If we're not certain if uv exists (try_exists == Err), keep going and hope it works.
+    if matches!(uv.try_exists(), Ok(false)) {
+        let message = if let Some(uv_with_suffix) = uv_with_suffix {
+            format!(
+                "Could not find the `uv` binary at either of:\n  {}\n  {}",
+                uv_with_suffix.display(),
+                uv.display(),
+            )
+        } else {
+            format!("Could not find the `uv` binary at: {}", uv.display())
+        };
+        Err(std::io::Error::new(std::io::ErrorKind::NotFound, message))
+    } else {
+        Ok(uv)
+    }
+}
+
 fn run() -> std::io::Result<ExitStatus> {
     let current_exe = std::env::current_exe()?;
     let Some(bin) = current_exe.parent() else {
@@ -30,22 +82,15 @@ fn run() -> std::io::Result<ExitStatus> {
             "Could not determine the location of the `uvx` binary",
         ));
     };
-    let uv = bin.join(format!("uv{}", std::env::consts::EXE_SUFFIX));
+    let uvx_suffix = get_uvx_suffix(&current_exe);
+    let uv = get_uv_path(bin, uvx_suffix)?;
     let args = ["tool", "uvx"]
         .iter()
         .map(OsString::from)
         // Skip the `uvx` name
         .chain(std::env::args_os().skip(1))
         .collect::<Vec<_>>();
 
-    // If we are sure the uv binary does not exist, display a clearer error message
-    if matches!(uv.try_exists(), Ok(false)) {
-        return Err(std::io::Error::new(
-            std::io::ErrorKind::NotFound,
-            format!("Could not find the `uv` binary at: {}", uv.display()),
-        ));
-    }
-
     let mut cmd = Command::new(uv);
     cmd.args(&args);
     match exec_spawn(&mut cmd)? {}