Support install uv from pypi

Cargo.toml

@@ -40,6 +40,7 @@ miette = { version = "7.2.0", features = ["owo-colors", "textwrap"] }
 owo-colors = "4.1.0"
 rand = "0.8.5"
 rayon = "1.10.0"
+reqwest = { version = "0.12.9", default-features = false }
 rusqlite = { version = "0.32.1", features = ["bundled"] }
 same-file = "1.0.6"
 serde = { version = "1.0.210", features = ["derive"] }
@@ -49,7 +50,7 @@ shlex = "1.3.0"
 tempfile = "3.13.0"
 textwrap = "0.16.1"
 thiserror = "1.0.64"
-tokio = { version = "1.40.0", features = ["fs", "process", "rt", "sync"] }
+tokio = { version = "1.40.0", features = ["fs", "process", "rt", "sync", "macros"] }
 tracing = "0.1.40"
 tracing-subscriber = { version = "0.3.18", features = ["env-filter"] }
 unicode-width = "0.2.0"

src/languages/python/impl.rs

@@ -4,7 +4,7 @@ use std::sync::Arc;
 
 use crate::env_vars::EnvVars;
 use crate::hook::Hook;
-use crate::languages::python::uv::ensure_uv;
+use crate::languages::python::uv::UvInstaller;
 use crate::languages::LanguageImpl;
 use crate::process::Cmd;
 use crate::run::run_by_batch;
@@ -26,7 +26,7 @@ impl LanguageImpl for Python {
     async fn install(&self, hook: &Hook) -> anyhow::Result<()> {
         let venv = hook.environment_dir().expect("No environment dir found");
 
-        let uv = ensure_uv().await?;
+        let uv = UvInstaller::install().await?;
 
         let uv_cmd = |summary| {
             #[allow(unused_mut)]

src/languages/python/uv.rs

@@ -1,78 +1,184 @@
 use std::env;
-use std::path::PathBuf;
+use std::path::{Path, PathBuf};
+use std::time::Duration;
 
 use anyhow::Result;
 use axoupdater::{AxoUpdater, ReleaseSource, ReleaseSourceType, UpdateRequest};
+use tokio::task::JoinSet;
 use tracing::{debug, enabled, trace, warn};
 
 use crate::fs::LockedFile;
 use crate::store::Store;
 
 // The version of `uv` to install. Should update periodically.
-const UV_VERSION: &str = "0.5.2";
+const UV_VERSION: &str = "0.5.8";
 
+#[derive(Debug)]
+enum PyPiMirror {
+    Pypi,
+    Tuna,
+    Aliyun,
+    Tencent,
+    Custom(String),
+}
+
+// TODO: support reading pypi source user config, or allow user to set mirror
 // TODO: allow opt-out uv
-// TODO: allow install uv using pip
-
-/// Ensure that the `uv` binary is available.
-pub(crate) async fn ensure_uv() -> Result<PathBuf> {
-    // 1) Check if `uv` is installed already.
-    if let Ok(uv) = which::which("uv") {
-        trace!(uv = %uv.display(), "Found uv from PATH");
-        return Ok(uv);
+
+impl PyPiMirror {
+    fn url(&self) -> &str {
+        match self {
+            Self::Pypi => "https://pypi.org/simple/",
+            Self::Tuna => "https://pypi.tuna.tsinghua.edu.cn/simple/",
+            Self::Aliyun => "https://mirrors.aliyun.com/pypi/simple/",
+            Self::Tencent => "https://mirrors.cloud.tencent.com/pypi/simple/",
+            Self::Custom(url) => url,
+        }
     }
 
-    // 2) Check if `uv` is installed by `prefligit`
-    let store = Store::from_settings()?;
+    fn iter() -> impl Iterator<Item = Self> {
+        vec![Self::Pypi, Self::Tuna, Self::Aliyun, Self::Tencent].into_iter()
+    }
+}
+
+#[derive(Debug)]
+enum InstallSource {
+    GitHub,
+    PyPi(PyPiMirror),
+}
 
-    let uv_dir = store.uv_path();
-    let uv = uv_dir.join("uv").with_extension(env::consts::EXE_EXTENSION);
-    if uv.is_file() {
-        trace!(uv = %uv.display(), "Found managed uv");
-        return Ok(uv);
+impl InstallSource {
+    async fn install(&self, target: &Path) -> Result<()> {
+        match self {
+            Self::GitHub => self.install_from_github(target).await,
+            Self::PyPi(source) => self.install_from_pypi(target, source).await,
+        }
     }
 
-    fs_err::create_dir_all(&uv_dir)?;
-    let _lock = LockedFile::acquire(uv_dir.join(".lock"), "uv").await?;
+    async fn install_from_github(&self, target: &Path) -> Result<()> {
+        let mut installer = AxoUpdater::new_for("uv");
+        installer.configure_version_specifier(UpdateRequest::SpecificTag(UV_VERSION.to_string()));
+        installer.always_update(true);
+        installer.set_install_dir(&target.to_string_lossy());
+        installer.set_release_source(ReleaseSource {
+            release_type: ReleaseSourceType::GitHub,
+            owner: "astral-sh".to_string(),
+            name: "uv".to_string(),
+            app_name: "uv".to_string(),
+        });
+        if enabled!(tracing::Level::DEBUG) {
+            installer.enable_installer_output();
+            env::set_var("INSTALLER_PRINT_VERBOSE", "1");
+        } else {
+            installer.disable_installer_output();
+        }
+        // We don't want the installer to modify the PATH, and don't need the receipt.
+        env::set_var("UV_UNMANAGED_INSTALL", "1");
 
-    if uv.is_file() {
-        trace!(uv = %uv.display(), "Found managed uv");
-        return Ok(uv);
+        match installer.run().await {
+            Ok(Some(result)) => {
+                debug!(
+                    uv = %target.display(),
+                    version = result.new_version_tag,
+                    "Successfully installed uv"
+                );
+                Ok(())
+            }
+            Ok(None) => Ok(()),
+            Err(err) => {
+                warn!(?err, "Failed to install uv");
+                Err(err.into())
+            }
+        }
     }
 
-    // 3) Download and install `uv`
-    let mut installer = AxoUpdater::new_for("uv");
-    installer.configure_version_specifier(UpdateRequest::SpecificTag(UV_VERSION.to_string()));
-    installer.always_update(true);
-    installer.set_install_dir(&uv_dir.to_string_lossy());
-    installer.set_release_source(ReleaseSource {
-        release_type: ReleaseSourceType::GitHub,
-        owner: "astral-sh".to_string(),
-        name: "uv".to_string(),
-        app_name: "uv".to_string(),
-    });
-    if enabled!(tracing::Level::DEBUG) {
-        installer.enable_installer_output();
-        env::set_var("INSTALLER_PRINT_VERBOSE", "1");
-    } else {
-        installer.disable_installer_output();
+    async fn install_from_pypi(&self, _target: &Path, _source: &PyPiMirror) -> Result<()> {
+        unimplemented!()
+    }
+}
+
+pub struct UvInstaller;
+
+impl UvInstaller {
+    async fn select_source() -> Result<InstallSource> {
+        async fn check_github() -> Result<bool> {
+            let url = format!("https://github.com/astral-sh/uv/releases/download/{UV_VERSION}/uv-x86_64-unknown-linux-gnu.tar.gz");
+            let response = reqwest::Client::new()
+                .head(url)
+                .timeout(Duration::from_secs(3))
+                .send()
+                .await?;
+            Ok(response.status().is_success())
+        }
+
+        async fn best_pypi() -> Result<PyPiMirror> {
+            let mut best = PyPiMirror::Pypi;
+            let mut tasks = PyPiMirror::iter()
+                .map(|source| async move {
+                    let url = format!("{}uv/", source.url());
+                    let response = reqwest::Client::new()
+                        .head(&url)
+                        .timeout(Duration::from_secs(2))
+                        .send()
+                        .await;
+                    (source, response)
+                })
+                .collect::<JoinSet<_>>();
+
+            while let Some(result) = tasks.join_next().await {
+                if let Ok((source, response)) = result {
+                    trace!(?source, ?response, "Checked source");
+                    if response.is_ok_and(|resp|resp.status().is_success()) {
+                        best = source;
+                        break;
+                    }
+                }
+            }
+
+            Ok(best)
+        }
+
+        let source = tokio::select! {
+            Ok(true) = check_github() => InstallSource::GitHub,
+            Ok(source) = best_pypi() => InstallSource::PyPi(source),
+            else => {
+                warn!("Failed to check uv source availability, defaulting to GitHub");
+                InstallSource::GitHub
+            }
+        };
+
+        trace!(?source, "Selected uv source");
+        Ok(source)
     }
-    // We don't want the installer to modify the PATH, and don't need the receipt.
-    env::set_var("UV_UNMANAGED_INSTALL", "1");
-
-    match installer.run().await {
-        Ok(Some(result)) => {
-            debug!(
-                uv = %uv.display(),
-                version = result.new_version_tag,
-                "Successfully installed uv"
-            );
-            Ok(uv)
+
+    pub async fn install() -> Result<PathBuf> {
+        // 1) Check if `uv` is installed already.
+        if let Ok(uv) = which::which("uv") {
+            trace!(uv = %uv.display(), "Found uv from PATH");
+            return Ok(uv);
         }
-        Ok(None) => Ok(uv),
-        Err(err) => {
-            warn!(?err, "Failed to install uv");
-            Err(err.into())
+
+        // 2) Check if `uv` is installed by `prefligit`
+        let store = Store::from_settings()?;
+
+        let uv_dir = store.uv_path();
+        let uv = uv_dir.join("uv").with_extension(env::consts::EXE_EXTENSION);
+        if uv.is_file() {
+            trace!(uv = %uv.display(), "Found managed uv");
+            return Ok(uv);
         }
+
+        fs_err::create_dir_all(&uv_dir)?;
+        let _lock = LockedFile::acquire(uv_dir.join(".lock"), "uv").await?;
+
+        if uv.is_file() {
+            trace!(uv = %uv.display(), "Found managed uv");
+            return Ok(uv);
+        }
+
+        let source = Self::select_source().await?;
+        source.install(&uv_dir).await?;
+
+        Ok(uv)
     }
 }