Add support for UV

ofek · Apr 21, 2024 · 67227ae · 67227ae
1 parent eb317a4
commit 67227ae
Show file tree

Hide file tree

Showing 5 changed files with 272 additions and 99 deletions.
diff --git a/Cargo.toml b/Cargo.toml
@@ -74,4 +74,7 @@ passthrough = [
   "PYAPP_SELF_COMMAND",
   "PYAPP_SKIP_INSTALL",
   "PYAPP_UPGRADE_VIRTUALENV",
+  "PYAPP_UV_ENABLED",
+  "PYAPP_UV_ONLY_BOOTSTRAP",
+  "PYAPP_UV_VERSION",
 ]
diff --git a/build.rs b/build.rs
@@ -309,13 +309,12 @@ fn get_distribution_source() -> String {
     let selected_arch = env::var("CARGO_CFG_TARGET_ARCH").unwrap();
     let selected_variant = {
         let mut variant = env::var("PYAPP_DISTRIBUTION_VARIANT").unwrap_or_default();
-        if variant.is_empty() {
-            if selected_platform == "linux"
-                && selected_arch == "x86_64"
-                && selected_python_version != "3.7"
-            {
-                variant = "v3".to_string();
-            }
+        if variant.is_empty()
+            && selected_platform == "linux"
+            && selected_arch == "x86_64"
+            && selected_python_version != "3.7"
+        {
+            variant = "v3".to_string();
         };
         variant
     };
@@ -852,6 +851,47 @@ fn set_pip_allow_config() {
     }
 }
 
+fn set_uv_enabled() {
+    let variable = "PYAPP_UV_ENABLED";
+    if is_enabled(variable) {
+        set_runtime_variable(variable, "1");
+    } else {
+        set_runtime_variable(variable, "0");
+    }
+}
+
+fn set_uv_only_bootstrap() {
+    let variable = "PYAPP_UV_ONLY_BOOTSTRAP";
+    if is_enabled(variable) {
+        set_runtime_variable(variable, "1");
+    } else {
+        set_runtime_variable(variable, "0");
+    }
+}
+
+fn set_uv_version() {
+    let variable = "PYAPP_UV_VERSION";
+    let version = env::var(variable).unwrap_or("any".to_string());
+    set_runtime_variable(variable, version);
+
+    let artifact_name = if !is_enabled("PYAPP_UV_ENABLED") {
+        "".to_string()
+    } else if env::var("CARGO_CFG_TARGET_OS").unwrap() == "windows" {
+        // Force MinGW-w64 to use msvc
+        if env::var("CARGO_CFG_TARGET_ENV").unwrap_or_default() == "gnu" {
+            format!(
+                "uv-{}-pc-windows-msvc.zip",
+                env::var("CARGO_CFG_TARGET_ARCH").unwrap()
+            )
+        } else {
+            format!("uv-{}.zip", env::var("TARGET").unwrap())
+        }
+    } else {
+        format!("uv-{}.tar.gz", env::var("TARGET").unwrap())
+    };
+    set_runtime_variable("PYAPP__UV_ARTIFACT_NAME", artifact_name);
+}
+
 fn set_skip_install() {
     let variable = "PYAPP_SKIP_INSTALL";
     if is_enabled(variable) {
@@ -942,6 +982,9 @@ fn main() {
     set_pip_project_features();
     set_pip_extra_args();
     set_pip_allow_config();
+    set_uv_enabled();
+    set_uv_only_bootstrap();
+    set_uv_version();
     set_skip_install();
     set_indicator();
     set_self_command();

diff --git a/src/app.rs b/src/app.rs
@@ -186,6 +186,26 @@ pub fn pip_external() -> bool {
     env!("PYAPP_PIP_EXTERNAL") == "1"
 }
 
+pub fn uv_enabled() -> bool {
+    env!("PYAPP_UV_ENABLED") == "1"
+}
+
+pub fn uv_only_bootstrap() -> bool {
+    env!("PYAPP_UV_ONLY_BOOTSTRAP") == "1"
+}
+
+pub fn uv_version() -> String {
+    env!("PYAPP_UV_VERSION").into()
+}
+
+pub fn uv_artifact_name() -> String {
+    env!("PYAPP__UV_ARTIFACT_NAME").into()
+}
+
+pub fn uv_as_installer() -> bool {
+    uv_enabled() && !uv_only_bootstrap()
+}
+
 pub fn is_gui() -> bool {
     env!("PYAPP_IS_GUI") == "1"
 }
@@ -235,6 +255,10 @@ pub fn external_pip_cache() -> PathBuf {
     cache_dir().join("pip")
 }
 
+pub fn managed_uv_cache() -> PathBuf {
+    cache_dir().join("uv").join(uv_version())
+}
+
 pub fn external_pip_zipapp() -> PathBuf {
     let pip_version = pip_version();
     let filename = if pip_version == "latest" {
@@ -244,3 +268,13 @@ pub fn external_pip_zipapp() -> PathBuf {
     };
     external_pip_cache().join(filename)
 }
+
+pub fn managed_uv() -> PathBuf {
+    let uv_artifact_name = uv_artifact_name();
+    let filename = if uv_artifact_name.ends_with(".zip") {
+        "uv.exe".to_string()
+    } else {
+        "uv".to_string()
+    };
+    managed_uv_cache().join(filename)
+}
diff --git a/src/compression.rs b/src/compression.rs
@@ -6,57 +6,58 @@ use anyhow::{bail, Result};
 use crate::terminal;
 
 pub fn unpack(format: String, archive: &PathBuf, destination: &PathBuf) -> Result<()> {
+    let wait_message = format!("Unpacking distribution ({})", format);
     match format.as_ref() {
-        "tar|bzip2" => unpack_tar_bzip2(archive, destination)?,
-        "tar|gzip" => unpack_tar_gzip(archive, destination)?,
-        "tar|zstd" => unpack_tar_zstd(archive, destination)?,
-        "zip" => unpack_zip(archive, destination)?,
+        "tar|bzip2" => unpack_tar_bzip2(archive, destination, wait_message)?,
+        "tar|gzip" => unpack_tar_gzip(archive, destination, wait_message)?,
+        "tar|zstd" => unpack_tar_zstd(archive, destination, wait_message)?,
+        "zip" => unpack_zip(archive, destination, wait_message)?,
         _ => bail!("unsupported distribution format: {}", format),
     }
 
     Ok(())
 }
 
-fn unpack_tar_bzip2(path: &PathBuf, destination: &PathBuf) -> Result<()> {
+fn unpack_tar_bzip2(path: &PathBuf, destination: &PathBuf, wait_message: String) -> Result<()> {
     let bz = bzip2::read::BzDecoder::new(File::open(path)?);
     let mut archive = tar::Archive::new(bz);
 
-    let spinner = terminal::spinner("Unpacking distribution (tar|bzip2)".to_string());
+    let spinner = terminal::spinner(wait_message);
     let result = archive.unpack(destination);
     spinner.finish_and_clear();
     result?;
 
     Ok(())
 }
 
-fn unpack_tar_gzip(path: &PathBuf, destination: &PathBuf) -> Result<()> {
+pub fn unpack_tar_gzip(path: &PathBuf, destination: &PathBuf, wait_message: String) -> Result<()> {
     let gz = flate2::read::GzDecoder::new(File::open(path)?);
     let mut archive = tar::Archive::new(gz);
 
-    let spinner = terminal::spinner("Unpacking distribution (tar|gzip)".to_string());
+    let spinner = terminal::spinner(wait_message);
     let result = archive.unpack(destination);
     spinner.finish_and_clear();
     result?;
 
     Ok(())
 }
 
-fn unpack_tar_zstd(path: &PathBuf, destination: &PathBuf) -> Result<()> {
+fn unpack_tar_zstd(path: &PathBuf, destination: &PathBuf, wait_message: String) -> Result<()> {
     let zst = zstd::stream::read::Decoder::new(File::open(path)?)?;
     let mut archive = tar::Archive::new(zst);
 
-    let spinner = terminal::spinner("Unpacking distribution (tar|zstd)".to_string());
+    let spinner = terminal::spinner(wait_message);
     let result = archive.unpack(destination);
     spinner.finish_and_clear();
     result?;
 
     Ok(())
 }
 
-fn unpack_zip(path: &PathBuf, destination: &PathBuf) -> Result<()> {
+pub fn unpack_zip(path: &PathBuf, destination: &PathBuf, wait_message: String) -> Result<()> {
     let mut archive = zip::ZipArchive::new(File::open(path)?)?;
 
-    let spinner = terminal::spinner("Unpacking distribution (zip)".to_string());
+    let spinner = terminal::spinner(wait_message);
     let result = archive.extract(destination);
     spinner.finish_and_clear();
     result?;