From bd9fe3ddff3642ae4d29ce17be553c49834e66c7 Mon Sep 17 00:00:00 2001
From: Zanie Blue <contact@zanie.dev>
Date: Thu, 26 Sep 2024 15:26:26 -0500
Subject: [PATCH] Add a `uv-python` shim executable

---
 .github/workflows/build-binaries.yml     |   4 +
 .github/workflows/build-docker.yml       |   2 +-
 Dockerfile                               |   7 +-
 crates/uv-cli/src/lib.rs                 |   8 ++
 crates/uv-python/src/discovery.rs        |   1 +
 crates/uv-python/src/interpreter.rs      |   1 +
 crates/uv-python/src/lib.rs              |   2 +-
 crates/uv/src/bin/uv-python.rs           | 168 +++++++++++++++++++++++
 crates/uv/src/commands/python/install.rs |  64 ++++++++-
 crates/uv/src/lib.rs                     |   8 ++
 crates/uv/src/settings.rs                |  14 +-
 crates/uv/tests/help.rs                  |   8 ++
 docs/reference/cli.md                    |   4 +
 13 files changed, 279 insertions(+), 12 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 74809604419a..291f2b7d8e03 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 c4ce6a9cdf27..d73fda0e7cc4 100644
--- a/.github/workflows/build-docker.yml
+++ b/.github/workflows/build-docker.yml
@@ -196,7 +196,7 @@ jobs:
           # Generate Dockerfile content
           cat <<EOF > Dockerfile
           FROM ${BASE_IMAGE}
-          COPY --from=${{ env.UV_BASE_IMG }}:latest /uv /uvx /usr/local/bin/
+          COPY --from=${{ env.UV_BASE_IMG }}:latest /uv /uvx /uv-python /usr/local/bin/
           ENTRYPOINT []
           CMD ["/usr/local/bin/uv"]
           EOF
diff --git a/Dockerfile b/Dockerfile
index d6aa983854e1..3b091d484d87 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -34,13 +34,14 @@ 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 /uvx /
+COPY --from=build /uv /uvx /uv-python /
 WORKDIR /io
 ENTRYPOINT ["/uv"]
diff --git a/crates/uv-cli/src/lib.rs b/crates/uv-cli/src/lib.rs
index 68b4dba9aada..514d1cdbcb6a 100644
--- a/crates/uv-cli/src/lib.rs
+++ b/crates/uv-cli/src/lib.rs
@@ -3546,6 +3546,14 @@ pub struct PythonInstallArgs {
     /// installed.
     #[arg(long, short, alias = "force")]
     pub reinstall: bool,
+
+    /// Install a `python` shim.
+    #[arg(long, overrides_with("no_shim"))]
+    pub shim: bool,
+
+    /// Do not install a `python` shim.
+    #[arg(long, overrides_with("shim"))]
+    pub no_shim: bool,
 }
 
 #[derive(Args)]
diff --git a/crates/uv-python/src/discovery.rs b/crates/uv-python/src/discovery.rs
index a0ff76405ccd..e1f6dba155df 100644
--- a/crates/uv-python/src/discovery.rs
+++ b/crates/uv-python/src/discovery.rs
@@ -1219,6 +1219,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-python/src/interpreter.rs b/crates/uv-python/src/interpreter.rs
index 4457f14cc6ad..a1349a6d1c2b 100644
--- a/crates/uv-python/src/interpreter.rs
+++ b/crates/uv-python/src/interpreter.rs
@@ -620,6 +620,7 @@ impl InterpreterInfo {
             .arg("-B") // Don't write bytecode.
             .arg("-c")
             .arg(script)
+            .env("UV_INTERNAL__PYTHON_QUERY", "1")
             .output()
             .map_err(|err| Error::SpawnFailed {
                 path: interpreter.to_path_buf(),
diff --git a/crates/uv-python/src/lib.rs b/crates/uv-python/src/lib.rs
index cb8e99f67f62..895430648dd1 100644
--- a/crates/uv-python/src/lib.rs
+++ b/crates/uv-python/src/lib.rs
@@ -6,7 +6,7 @@ pub use crate::discovery::{
     PythonNotFound, PythonPreference, PythonRequest, PythonSource, VersionRequest,
 };
 pub use crate::environment::{InvalidEnvironment, InvalidEnvironmentKind, PythonEnvironment};
-pub use crate::implementation::ImplementationName;
+pub use crate::implementation::{ImplementationName, LenientImplementationName};
 pub use crate::installation::{PythonInstallation, PythonInstallationKey};
 pub use crate::interpreter::{Error as InterpreterError, Interpreter};
 pub use crate::pointer_size::PointerSize;
diff --git a/crates/uv/src/bin/uv-python.rs b/crates/uv/src/bin/uv-python.rs
new file mode 100644
index 000000000000..21b814a78f50
--- /dev/null
+++ b/crates/uv/src/bin/uv-python.rs
@@ -0,0 +1,168 @@
+use std::convert::Infallible;
+use std::io::Write;
+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),
+    RecursiveQuery,
+}
+
+#[derive(Debug, Default)]
+struct Options {
+    request: Option<String>,
+    system: bool,
+    managed: bool,
+    verbose: 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());
+        } else {
+            // By default, we should never select an alternative implementation with the shim
+            args.push("cpython");
+        }
+        if self.system {
+            args.push("--system");
+        }
+        if self.verbose {
+            args.push("--verbose");
+        }
+        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}"),
+            Self::RecursiveQuery => write!(f, "Ignoring recursive query from uv"),
+        }
+    }
+}
+
+/// 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
+/// - `+v`: Enable verbose mode.
+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,
+            "v" => options.verbose = 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> {
+    if std::env::var_os("UV_INTERNAL__PYTHON_QUERY").is_some() {
+        return Err(Error::RecursiveQuery);
+    }
+
+    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(),
+        ));
+    }
+
+    // If verbose is enabled, print the output of the `uv python find` command
+    if options.verbose {
+        std::io::stderr()
+            .write_all(&output.stderr)
+            .map_err(Error::Io)?;
+    }
+
+    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)
+        }
+    }
+}
diff --git a/crates/uv/src/commands/python/install.rs b/crates/uv/src/commands/python/install.rs
index ac373eafd7c1..08317856a8ff 100644
--- a/crates/uv/src/commands/python/install.rs
+++ b/crates/uv/src/commands/python/install.rs
@@ -1,17 +1,24 @@
-use anyhow::Result;
+use std::collections::BTreeSet;
+use std::fmt::Write;
+use std::path::{Path, PathBuf};
+
+use anyhow::{bail, Result};
 use fs_err as fs;
 use futures::stream::FuturesUnordered;
 use futures::StreamExt;
 use itertools::Itertools;
 use owo_colors::OwoColorize;
-use std::collections::BTreeSet;
-use std::fmt::Write;
-use std::path::Path;
 
+use tracing::debug;
 use uv_client::Connectivity;
+use uv_configuration::PreviewMode;
+use uv_fs::Simplified;
 use uv_python::downloads::{DownloadResult, ManagedPythonDownload, PythonDownloadRequest};
 use uv_python::managed::{ManagedPythonInstallation, ManagedPythonInstallations};
-use uv_python::{PythonDownloads, PythonRequest, PythonVersionFile};
+use uv_python::{
+    ImplementationName, LenientImplementationName, PythonDownloads, PythonRequest,
+    PythonVersionFile,
+};
 
 use crate::commands::python::{ChangeEvent, ChangeEventKind};
 use crate::commands::reporters::PythonDownloadReporter;
@@ -23,6 +30,8 @@ pub(crate) async fn install(
     project_dir: &Path,
     targets: Vec<String>,
     reinstall: bool,
+    shim: Option<bool>,
+    preview: PreviewMode,
     python_downloads: PythonDownloads,
     native_tls: bool,
     connectivity: Connectivity,
@@ -170,6 +179,13 @@ pub(crate) async fn install(
         }
     }
 
+    let any_cpython = installed.iter().any(|key| {
+        matches!(
+            key.implementation(),
+            LenientImplementationName::Known(ImplementationName::CPython)
+        )
+    });
+
     if !installed.is_empty() {
         if let [installed] = installed.as_slice() {
             // Ex) "Installed Python 3.9.7 in 1.68s"
@@ -236,5 +252,43 @@ pub(crate) async fn install(
         return Ok(ExitStatus::Failure);
     }
 
+    // Install a shim if explicitly requested or if we installed a CPython version
+    if shim.unwrap_or(preview.is_enabled() && any_cpython) {
+        let shim_src = find_shim()?;
+        let executable_dir = uv_tool::find_executable_directory()?;
+        let shim_dst = executable_dir.join("python");
+        if shim_dst.try_exists()? {
+            writeln!(
+                printer.stderr(),
+                "Python executable already exists at `{}`",
+                shim_dst.user_display().cyan()
+            )?;
+        } else {
+            debug!(
+                "Linking {} -> {}",
+                shim_src.user_display(),
+                shim_dst.user_display()
+            );
+            uv_fs::replace_symlink(&shim_src, &shim_dst)?;
+            writeln!(
+                printer.stderr(),
+                "Installed Python shim to `{}`",
+                shim_dst.user_display().cyan()
+            )?;
+        }
+    }
+
     Ok(ExitStatus::Success)
 }
+
+fn find_shim() -> Result<PathBuf> {
+    let current_exe = std::env::current_exe()?;
+    let Some(bin) = current_exe.parent() else {
+        bail!("Could not find the directory for the `uv-python` binary");
+    };
+    let uv_python = bin.join("uv-python");
+    if !uv_python.try_exists()? {
+        bail!("Could not find the `uv-python` binary");
+    }
+    Ok(uv_python)
+}
diff --git a/crates/uv/src/lib.rs b/crates/uv/src/lib.rs
index 983bc3c2ba3b..0f59ada12a5a 100644
--- a/crates/uv/src/lib.rs
+++ b/crates/uv/src/lib.rs
@@ -1006,10 +1006,18 @@ async fn run(cli: Cli) -> Result<ExitStatus> {
             let args = settings::PythonInstallSettings::resolve(args, filesystem);
             show_settings!(args);
 
+            if matches!(args.shim, Some(true)) && globals.preview.is_disabled() {
+                warn_user_once!(
+                    "The uv Python shim is experimental and may change without warning"
+                );
+            }
+
             commands::python_install(
                 &project_dir,
                 args.targets,
                 args.reinstall,
+                args.shim,
+                globals.preview,
                 globals.python_downloads,
                 globals.native_tls,
                 globals.connectivity,
diff --git a/crates/uv/src/settings.rs b/crates/uv/src/settings.rs
index fb924949e161..b6e0a78d124d 100644
--- a/crates/uv/src/settings.rs
+++ b/crates/uv/src/settings.rs
@@ -585,15 +585,25 @@ impl PythonListSettings {
 pub(crate) struct PythonInstallSettings {
     pub(crate) targets: Vec<String>,
     pub(crate) reinstall: bool,
+    pub(crate) shim: Option<bool>,
 }
 
 impl PythonInstallSettings {
     /// Resolve the [`PythonInstallSettings`] from the CLI and filesystem configuration.
     #[allow(clippy::needless_pass_by_value)]
     pub(crate) fn resolve(args: PythonInstallArgs, _filesystem: Option<FilesystemOptions>) -> Self {
-        let PythonInstallArgs { targets, reinstall } = args;
+        let PythonInstallArgs {
+            targets,
+            reinstall,
+            shim,
+            no_shim,
+        } = args;
 
-        Self { targets, reinstall }
+        Self {
+            targets,
+            reinstall,
+            shim: flag(shim, no_shim),
+        }
     }
 }
 
diff --git a/crates/uv/tests/help.rs b/crates/uv/tests/help.rs
index e5d229dd5a49..9c10c4f77546 100644
--- a/crates/uv/tests/help.rs
+++ b/crates/uv/tests/help.rs
@@ -432,6 +432,12 @@ fn help_subsubcommand() {
               
               By default, uv will exit successfully if the version is already installed.
 
+          --shim
+              Install a `python` shim
+
+          --no-shim
+              Do not install a `python` shim
+
     Cache options:
       -n, --no-cache
               Avoid reading from or writing to the cache, instead using a temporary directory for the
@@ -633,6 +639,8 @@ fn help_flag_subsubcommand() {
 
     Options:
       -r, --reinstall  Reinstall the requested Python version, if it's already installed
+          --shim       Install a `python` shim
+          --no-shim    Do not install a `python` shim
 
     Cache options:
       -n, --no-cache               Avoid reading from or writing to the cache, instead using a temporary
diff --git a/docs/reference/cli.md b/docs/reference/cli.md
index 221f35f6599b..b026bf3145cf 100644
--- a/docs/reference/cli.md
+++ b/docs/reference/cli.md
@@ -4026,6 +4026,8 @@ uv python install [OPTIONS] [TARGETS]...
 
 </dd><dt><code>--no-python-downloads</code></dt><dd><p>Disable automatic downloads of Python.</p>
 
+</dd><dt><code>--no-shim</code></dt><dd><p>Do not install a <code>python</code> shim</p>
+
 </dd><dt><code>--offline</code></dt><dd><p>Disable network access.</p>
 
 <p>When disabled, uv will only use locally cached data and locally available files.</p>
@@ -4062,6 +4064,8 @@ uv python install [OPTIONS] [TARGETS]...
 
 <p>By default, uv will exit successfully if the version is already installed.</p>
 
+</dd><dt><code>--shim</code></dt><dd><p>Install a <code>python</code> shim</p>
+
 </dd><dt><code>--verbose</code>, <code>-v</code></dt><dd><p>Use verbose output.</p>
 
 <p>You can configure fine-grained logging using the <code>RUST_LOG</code> environment variable. (&lt;https://docs.rs/tracing-subscriber/latest/tracing_subscriber/filter/struct.EnvFilter.html#directives&gt;)</p>