Allow custom sources for UV (#153)

Co-authored-by: Jp <coruscant11@proton.me>
ofek · Aug 3, 2024 · 0b8d25b · 0b8d25b
1 parent 2980970
commit 0b8d25b
Show file tree

Hide file tree

Showing 6 changed files with 72 additions and 32 deletions.
diff --git a/Cargo.toml b/Cargo.toml
@@ -82,5 +82,6 @@ passthrough = [
   "PYAPP_UPGRADE_VIRTUALENV",
   "PYAPP_UV_ENABLED",
   "PYAPP_UV_ONLY_BOOTSTRAP",
+  "PYAPP_UV_SOURCE",
   "PYAPP_UV_VERSION",
 ]
diff --git a/build.rs b/build.rs
@@ -239,6 +239,19 @@ fn check_environment_variable(name: &str) -> String {
     value
 }
 
+fn filename_from_url(url: &str) -> String {
+    let parsed =
+        reqwest::Url::parse(url).unwrap_or_else(|_| panic!("unable to parse URL: {}", &url));
+
+    if let Some(segments) = parsed.path_segments() {
+        if let Some(segment) = segments.last() {
+            return segment.into();
+        }
+    }
+
+    panic!("unable to determine artifact name from URL: {}", &url);
+}
+
 fn is_enabled(name: &str) -> bool {
     ["true", "1"].contains(&env::var(name).unwrap_or_default().as_str())
 }
@@ -922,27 +935,51 @@ fn set_uv_only_bootstrap() {
     }
 }
 
-fn set_uv_version() {
-    let variable = "PYAPP_UV_VERSION";
-    let version = env::var(variable).unwrap_or("any".to_string());
-    set_runtime_variable(variable, version);
+fn set_uv_source() {
+    let source_variable = "PYAPP_UV_SOURCE";
+    let mut source = env::var(source_variable).unwrap_or_default();
 
-    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" {
+    if !source.is_empty() {
+        set_runtime_variable("PYAPP__UV_ARTIFACT_NAME", filename_from_url(&source));
+
+        let mut hasher = PortableHash::default();
+        source.hash(&mut hasher);
+        set_runtime_variable("PYAPP__UV_VERSION", hasher.finish());
+    } else {
+        let version_variable = "PYAPP_UV_VERSION";
+        let version = env::var(version_variable).unwrap_or("any".to_string());
+
+        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())
+        };
+
+        source = if version == "any" {
             format!(
-                "uv-{}-pc-windows-msvc.zip",
-                env::var("CARGO_CFG_TARGET_ARCH").unwrap()
+                "https://github.com/astral-sh/uv/releases/latest/download/{}",
+                &artifact_name,
             )
         } 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);
+            format!(
+                "https://github.com/astral-sh/uv/releases/download/{}/{}",
+                &version, &artifact_name,
+            )
+        };
+        set_runtime_variable("PYAPP__UV_ARTIFACT_NAME", artifact_name);
+        set_runtime_variable("PYAPP__UV_VERSION", &version);
+    }
+    set_runtime_variable(source_variable, &source);
 }
 
 fn set_skip_install() {
@@ -1070,7 +1107,7 @@ fn main() {
     set_pip_allow_config();
     set_uv_enabled();
     set_uv_only_bootstrap();
-    set_uv_version();
+    set_uv_source();
     set_allow_updates();
     set_indicator();
     set_self_command();

diff --git a/docs/changelog.md b/docs/changelog.md
@@ -8,6 +8,10 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 
 ## Unreleased
 
+***Added:***
+
+- Add `PYAPP_UV_SOURCE` option
+
 ## 0.22.0 - 2024-05-26
 
 ***Added:***

diff --git a/docs/config/installation.md b/docs/config/installation.md
@@ -18,6 +18,12 @@ By default, a version of UV that has already been downloaded by a PyApp applicat
 
 You may set the `PYAPP_UV_ONLY_BOOTSTRAP` option to `true` or `1` to only use UV for virtual environment creation and continue using pip for project installation.
 
+### Source
+
+You may explicitly set the `PYAPP_UV_SOURCE` option in order to download your own UV release archive. This takes precedence over any defined [version](#uv-version).
+
+The value must end with the archive's real file extension, which is used to determine the extraction method.
+
 ## pip
 
 These options have no effect when UV is [enabled](#uv).

diff --git a/src/app.rs b/src/app.rs
@@ -194,8 +194,12 @@ pub fn uv_only_bootstrap() -> bool {
     env!("PYAPP_UV_ONLY_BOOTSTRAP") == "1"
 }
 
+pub fn uv_source() -> String {
+    env!("PYAPP_UV_SOURCE").into()
+}
+
 pub fn uv_version() -> String {
-    env!("PYAPP_UV_VERSION").into()
+    env!("PYAPP__UV_VERSION").into()
 }
 
 pub fn uv_artifact_name() -> String {

diff --git a/src/distribution.rs b/src/distribution.rs
@@ -494,19 +494,7 @@ fn ensure_uv_available() -> Result<()> {
     let mut f = fs::File::create(&temp_path)
         .with_context(|| format!("unable to create temporary file: {}", &temp_path.display()))?;
 
-    let url = if uv_version == "any" {
-        format!(
-            "https://github.com/astral-sh/uv/releases/latest/download/{}",
-            &artifact_name,
-        )
-    } else {
-        format!(
-            "https://github.com/astral-sh/uv/releases/download/{}/{}",
-            &uv_version, &artifact_name,
-        )
-    };
-
-    network::download(&url, &mut f, "UV")?;
+    network::download(&app::uv_source(), &mut f, "UV")?;
 
     if artifact_name.ends_with(".zip") {
         compression::unpack_zip(temp_path, dir.path(), "Unpacking UV".to_string())