diff --git a/Cargo.lock b/Cargo.lock index 31ea355..ec6cd5a 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1569,6 +1569,7 @@ dependencies = [ "rand", "rayon", "regex", + "reqwest", "rusqlite", "same-file", "serde", @@ -2284,9 +2285,21 @@ dependencies = [ "pin-project-lite", "signal-hook-registry", "socket2", + "tokio-macros", "windows-sys 0.52.0", ] +[[package]] +name = "tokio-macros" +version = "2.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "693d596312e88961bc67d7f1f97af8a70227d9f90c31bba5806eec004978d752" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "tokio-rustls" version = "0.26.1" diff --git a/Cargo.toml b/Cargo.toml index 4dbc7b6..587a9b2 100644 --- a/Cargo.toml +++ b/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" diff --git a/src/languages/python/impl.rs b/src/languages/python/impl.rs index 6f23cf7..ce2a431 100644 --- a/src/languages/python/impl.rs +++ b/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)] diff --git a/src/languages/python/uv.rs b/src/languages/python/uv.rs index 4b3525a..3ec84f9 100644 --- a/src/languages/python/uv.rs +++ b/src/languages/python/uv.rs @@ -1,78 +1,230 @@ 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::process::Cmd; 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 { - // 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 { + vec![Self::Pypi, Self::Tuna, Self::Aliyun, Self::Tencent].into_iter() + } +} + +#[derive(Debug)] +enum InstallSource { + /// Download uv from GitHub releases. + GitHub, + /// Download uv from `PyPi`. + PyPi(PyPiMirror), + /// Install uv by running `pip install uv`. + Pip, +} - 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, + Self::Pip => self.install_from_pip(target).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<()> { + // TODO: Implement this, currently just fallback to pip install + // Determine the host system + // Get the html page + // Parse html, get the latest version url + // Download the tarball + // Extract the tarball + self.install_from_pip(target).await } - // 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) + + async fn install_from_pip(&self, target: &Path) -> Result<()> { + Cmd::new("python3", "pip install uv") + .arg("-m") + .arg("pip") + .arg("install") + .arg("--prefix") + .arg(target) + .arg(format!("uv=={UV_VERSION}")) + .check(true) + .output() + .await?; + + let bin_dir = target.join(if cfg!(windows) { "Scripts" } else { "bin" }); + let lib_dir = target.join(if cfg!(windows) { "Lib" } else { "lib" }); + + let uv = target + .join(&bin_dir) + .join("uv") + .with_extension(env::consts::EXE_EXTENSION); + fs_err::rename( + &uv, + target.join("uv").with_extension(env::consts::EXE_EXTENSION), + )?; + fs_err::remove_dir_all(bin_dir)?; + fs_err::remove_dir_all(lib_dir)?; + + Ok(()) + } +} + +pub struct UvInstaller; + +impl UvInstaller { + async fn select_source() -> Result { + async fn check_github(client: &reqwest::Client) -> Result { + let url = format!("https://github.com/astral-sh/uv/releases/download/{UV_VERSION}/uv-x86_64-unknown-linux-gnu.tar.gz"); + let response = client + .head(url) + .timeout(Duration::from_secs(3)) + .send() + .await?; + trace!(?response, "Checked GitHub"); + Ok(response.status().is_success()) + } + + async fn select_best_pypi(client: &reqwest::Client) -> Result { + let mut best = PyPiMirror::Pypi; + let mut tasks = PyPiMirror::iter() + .map(|source| { + let client = client.clone(); + async move { + let url = format!("{}uv/", source.url()); + let response = client + .head(&url) + .timeout(Duration::from_secs(2)) + .send() + .await; + (source, response) + } + }) + .collect::>(); + + 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) } - Ok(None) => Ok(uv), - Err(err) => { - warn!(?err, "Failed to install uv"); - Err(err.into()) + + let client = reqwest::Client::new(); + let source = tokio::select! { + Ok(true) = check_github(&client) => InstallSource::GitHub, + Ok(source) = select_best_pypi(&client) => InstallSource::PyPi(source), + else => { + warn!("Failed to check uv source availability, falling back to pip install"); + InstallSource::Pip + } + }; + + trace!(?source, "Selected uv source"); + Ok(source) + } + + pub async fn install() -> Result { + // 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); } + + // 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) } }