From 1f6255e8509b1f3e2e149178569d025aa4305932 Mon Sep 17 00:00:00 2001
From: Ibraheem Ahmed <ibraheem@ibraheem.ca>
Date: Fri, 19 Jul 2024 17:44:32 -0400
Subject: [PATCH] exclude non-local version from range

---
 crates/uv-resolver/src/resolver/mod.rs |  9 ++-
 crates/uv/tests/pip_compile.rs         | 81 +++++++++++++++++++++++++-
 2 files changed, 88 insertions(+), 2 deletions(-)

diff --git a/crates/uv-resolver/src/resolver/mod.rs b/crates/uv-resolver/src/resolver/mod.rs
index 6000f31dd6758..8bacd64cfc3e8 100644
--- a/crates/uv-resolver/src/resolver/mod.rs
+++ b/crates/uv-resolver/src/resolver/mod.rs
@@ -2049,10 +2049,17 @@ impl ForkState {
                 // If the specifier is an exact version and the user requested a local version for this
                 // fork that's more precise than the specifier, use the local version instead.
                 if let Some(specifier) = specifier {
+                    let locals = locals.get(name, &self.markers);
+
+                    // Prioritize local versions over the original version range.
+                    if locals.len() > 0 {
+                        *version = Range::empty();
+                    }
+
                     // It's possible that there are multiple matching local versions requested with
                     // different marker expressions. All of these are potentially compatible until we
                     // narrow to a specific fork.
-                    for local in locals.get(name, &self.markers) {
+                    for local in locals {
                         let local = specifier
                             .iter()
                             .map(|specifier| {
diff --git a/crates/uv/tests/pip_compile.rs b/crates/uv/tests/pip_compile.rs
index 2ce23fe2510ac..14a5b4f6f1c0d 100644
--- a/crates/uv/tests/pip_compile.rs
+++ b/crates/uv/tests/pip_compile.rs
@@ -6986,7 +6986,7 @@ fn universal_overlapping_local_requirement() -> Result<()> {
     Ok(())
 }
 
-/// If a dependency requests distinct local versions with distinct marker expressions,
+/// If a dependency requests distinct local versions with disjoint marker expressions,
 /// we should fork the root requirement.
 #[test]
 fn universal_disjoint_local_requirement() -> Result<()> {
@@ -7064,6 +7064,85 @@ fn universal_disjoint_local_requirement() -> Result<()> {
     Ok(())
 }
 
+/// If a dependency requests distinct local versions and non-local versions with disjoint marker
+/// expressions, we should fork the root requirement.
+#[test]
+fn universal_disjoint_base_or_local_requirement() -> Result<()> {
+    let context = TestContext::new("3.12");
+
+    let pyproject_toml = context.temp_dir.child("pyproject.toml");
+    pyproject_toml.write_str(indoc! {r#"
+        [project]
+        name = "example"
+        version = "0.0.0"
+        dependencies = [
+            "torch==2.0.0; python_version < '3.10'",
+            "torch==2.0.0+cu118 ; python_version >= '3.10' and python_version <= '3.12'",
+            "torch==2.0.0+cpu ; python_version > '3.12'"
+        ]
+        requires-python = ">=3.11"
+    "#})?;
+
+    let requirements_in = context.temp_dir.child("requirements.in");
+    requirements_in.write_str(indoc! {"
+        torch==2.0.0
+        .
+    "})?;
+
+    // Some marker expressions on the output here are missing due to https://github.com/astral-sh/uv/issues/5086,
+    // but the local versions are still respected correctly.
+    uv_snapshot!(context.pip_compile()
+        .arg("requirements.in")
+        .arg("--universal")
+        .arg("--find-links")
+        .arg("https://download.pytorch.org/whl/torch_stable.html"), @r###"
+    success: true
+    exit_code: 0
+    ----- stdout -----
+    # This file was autogenerated by uv via the following command:
+    #    uv pip compile --cache-dir [CACHE_DIR] requirements.in --universal
+    cmake==3.28.4 ; platform_machine == 'x86_64' and platform_system == 'Linux'
+        # via triton
+    .
+        # via -r requirements.in
+    filelock==3.13.1
+        # via
+        #   torch
+        #   triton
+    jinja2==3.1.3
+        # via torch
+    lit==18.1.2 ; platform_machine == 'x86_64' and platform_system == 'Linux'
+        # via triton
+    markupsafe==2.1.5
+        # via jinja2
+    mpmath==1.3.0
+        # via sympy
+    networkx==3.2.1
+        # via torch
+    sympy==1.12
+        # via torch
+    torch==2.0.0+cpu
+        # via
+        #   -r requirements.in
+        #   example
+    torch==2.0.0+cu118
+        # via
+        #   -r requirements.in
+        #   example
+        #   triton
+    triton==2.0.0 ; platform_machine == 'x86_64' and platform_system == 'Linux'
+        # via torch
+    typing-extensions==4.10.0
+        # via torch
+
+    ----- stderr -----
+    Resolved 13 packages in [TIME]
+    "###
+    );
+
+    Ok(())
+}
+
 /// If a dependency requests a local version with an overlapping marker expression
 /// that form a nested fork, we should prefer the local in both children of the outer
 /// fork.