From 5b9b3ffefba06dfea22d74b89d4bec9052b13421 Mon Sep 17 00:00:00 2001
From: Zanie Blue <contact@zanie.dev>
Date: Tue, 24 Sep 2024 12:52:46 -0500
Subject: [PATCH] Add a `uv-python` shim executable

---
 .github/workflows/build-binaries.yml |   4 +
 .github/workflows/build-docker.yml   |   1 +
 Dockerfile                           |   6 +-
 crates/uv-python/src/discovery.rs    |   1 +
 crates/uv/src/bin/uv-python.rs       | 144 +++++++++++++++++++++++++++
 5 files changed, 154 insertions(+), 2 deletions(-)
 create mode 100644 crates/uv/src/bin/uv-python.rs

diff --git a/.github/workflows/build-binaries.yml b/.github/workflows/build-binaries.yml
index eab2348049ee8..471b0dfa7ca9e 100644
--- a/.github/workflows/build-binaries.yml
+++ b/.github/workflows/build-binaries.yml
@@ -56,6 +56,7 @@ jobs:
           ${{ env.MODULE_NAME }} --help
           python -m ${{ env.MODULE_NAME }} --help
           uvx --help
+          uv-python +3.11 --help
       - name: "Upload sdist"
         uses: actions/upload-artifact@v4
         with:
@@ -128,6 +129,7 @@ jobs:
           ${{ env.MODULE_NAME }} --help
           python -m ${{ env.MODULE_NAME }} --help
           uvx --help
+          uv-python +3.11 --help
       - name: "Upload wheels"
         uses: actions/upload-artifact@v4
         with:
@@ -185,6 +187,7 @@ jobs:
           ${{ env.MODULE_NAME }} --help
           python -m ${{ env.MODULE_NAME }} --help
           uvx --help
+          uv-python +3.11 --help
       - name: "Upload wheels"
         uses: actions/upload-artifact@v4
         with:
@@ -251,6 +254,7 @@ jobs:
           ${{ env.MODULE_NAME }} --help
           python -m ${{ env.MODULE_NAME }} --help
           uvx --help
+          uv-python +3.11 --help
       - name: "Upload wheels"
         uses: actions/upload-artifact@v4
         with:
diff --git a/.github/workflows/build-docker.yml b/.github/workflows/build-docker.yml
index 14936abd927af..ff819cf1d8e66 100644
--- a/.github/workflows/build-docker.yml
+++ b/.github/workflows/build-docker.yml
@@ -198,6 +198,7 @@ jobs:
           FROM ${BASE_IMAGE}
           COPY --from=${{ env.UV_BASE_IMG }}:latest /uv /usr/local/bin/uv
           COPY --from=${{ env.UV_BASE_IMG }}:latest /uvx /usr/local/bin/uvx
+          COPY --from=${{ env.UV_BASE_IMG }}:latest /uv-python /usr/local/bin/uv-python
           ENTRYPOINT []
           CMD ["/usr/local/bin/uv"]
           EOF
diff --git a/Dockerfile b/Dockerfile
index 2a46222961157..5e037edd6d194 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -34,14 +34,16 @@ RUN rustup target add $(cat rust_target.txt)
 COPY crates crates
 COPY ./Cargo.toml Cargo.toml
 COPY ./Cargo.lock Cargo.lock
-RUN cargo zigbuild --bin uv --bin uvx --target $(cat rust_target.txt) --release
+RUN cargo zigbuild --bin uv --bin uvx --bin uv-python --target $(cat rust_target.txt) --release
 RUN cp target/$(cat rust_target.txt)/release/uv /uv \
-  && cp target/$(cat rust_target.txt)/release/uvx /uvx
+  && cp target/$(cat rust_target.txt)/release/uvx /uvx \
+  && cp target/$(cat rust_target.txt)/release/uv-python /uv-python
 # TODO(konsti): Optimize binary size, with a version that also works when cross compiling
 # RUN strip --strip-all /uv
 
 FROM scratch
 COPY --from=build /uv /uv
 COPY --from=build /uvx /uvx
+COPY --from=build /uv-python /uv-python
 WORKDIR /io
 ENTRYPOINT ["/uv"]
diff --git a/crates/uv-python/src/discovery.rs b/crates/uv-python/src/discovery.rs
index 05c21e964eeea..64216aac79c6a 100644
--- a/crates/uv-python/src/discovery.rs
+++ b/crates/uv-python/src/discovery.rs
@@ -1188,6 +1188,7 @@ impl PythonRequest {
         if let Ok(version) = VersionRequest::from_str(value) {
             return Self::Version(version);
         }
+
         // e.g. `python3.12.1`
         if let Some(remainder) = value.strip_prefix("python") {
             if let Ok(version) = VersionRequest::from_str(remainder) {
diff --git a/crates/uv/src/bin/uv-python.rs b/crates/uv/src/bin/uv-python.rs
new file mode 100644
index 0000000000000..7db60f367b500
--- /dev/null
+++ b/crates/uv/src/bin/uv-python.rs
@@ -0,0 +1,144 @@
+use std::convert::Infallible;
+use std::{
+    ffi::OsString,
+    process::{Command, ExitCode, ExitStatus},
+};
+
+/// Spawns a command exec style.
+fn exec_spawn(cmd: &mut Command) -> std::io::Result<Infallible> {
+    #[cfg(unix)]
+    {
+        use std::os::unix::process::CommandExt;
+        let err = cmd.exec();
+        Err(err)
+    }
+    #[cfg(windows)]
+    {
+        cmd.stdin(std::process::Stdio::inherit());
+        let status = cmd.status()?;
+
+        #[allow(clippy::exit)]
+        std::process::exit(status.code().unwrap())
+    }
+}
+
+#[derive(Debug)]
+enum Error {
+    Io(std::io::Error),
+    Which(which::Error),
+    NoInterpreter(String),
+}
+
+#[derive(Debug, Default)]
+struct Options {
+    request: Option<String>,
+    system: bool,
+    managed: bool,
+}
+
+impl Options {
+    fn as_args(&self) -> Vec<&str> {
+        let mut args = Vec::new();
+        if let Some(request) = &self.request {
+            args.push(request.as_str());
+        }
+        if self.system {
+            args.push("--system");
+        }
+        if self.managed {
+            args.push("--python-preference");
+            args.push("only-managed");
+        }
+        args
+    }
+}
+
+impl std::fmt::Display for Error {
+    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
+        match self {
+            Self::Io(err) => write!(f, "{err}"),
+            Self::Which(err) => write!(f, "Failed to find uv binary: {err}"),
+            Self::NoInterpreter(inner) => write!(f, "{inner}"),
+        }
+    }
+}
+
+/// Parse `+<option>` into [`Options`].
+///
+/// Supports the following options:
+///
+/// - `+system`: Use the system Python, ignore virtual environments.
+/// - `+managed`: Use only managed Python installations.
+/// - `+<request>`: Request a Python version
+fn parse_options(mut args: Vec<OsString>) -> (Vec<OsString>, Options) {
+    let mut position = 0;
+    let mut options = Options::default();
+    while position < args.len() {
+        let arg = &args[position].to_string_lossy();
+
+        // If the argument doesn't start with `+`, we're done.
+        let Some(option) = arg.strip_prefix('+') else {
+            break;
+        };
+
+        match option {
+            "system" => options.system = true,
+            "managed" => options.managed = true,
+            _ => options.request = Some(option.to_string()),
+        }
+
+        position += 1;
+    }
+
+    (args.split_off(position), options)
+}
+
+/// Find the `uv` binary to use.
+fn find_uv() -> Result<std::path::PathBuf, Error> {
+    // We prefer one next to the current binary.
+    let current_exe = std::env::current_exe().map_err(Error::Io)?;
+    if let Some(bin) = current_exe.parent() {
+        let uv = bin.join("uv");
+        if uv.exists() {
+            return Ok(uv);
+        }
+    }
+    // Otherwise, we'll search for it on the `PATH`.
+    which::which("uv").map_err(Error::Which)
+}
+
+fn run() -> Result<ExitStatus, Error> {
+    let args = std::env::args_os().skip(1).collect::<Vec<_>>();
+    let (args, options) = parse_options(args);
+    let uv = find_uv()?;
+    let mut cmd = Command::new(uv);
+    let uv_args = ["python", "find"].iter().copied().chain(options.as_args());
+    cmd.args(uv_args);
+    let output = cmd.output().map_err(Error::Io)?;
+    if !output.status.success() {
+        let stderr = String::from_utf8_lossy(output.stderr.as_slice());
+        return Err(Error::NoInterpreter(
+            stderr
+                .strip_prefix("error: ")
+                .unwrap_or(&*stderr)
+                .to_string(),
+        ));
+    }
+    let python = std::path::PathBuf::from(String::from_utf8_lossy(output.stdout.as_slice()).trim());
+    let mut cmd = Command::new(python);
+    cmd.args(&args);
+    match exec_spawn(&mut cmd).map_err(Error::Io)? {}
+}
+
+#[allow(clippy::print_stderr)]
+fn main() -> ExitCode {
+    let result = run();
+    match result {
+        // Fail with 2 if the status cannot be cast to an exit code
+        Ok(status) => u8::try_from(status.code().unwrap_or(2)).unwrap_or(2).into(),
+        Err(err) => {
+            eprintln!("error: {err}");
+            ExitCode::from(2)
+        }
+    }
+}