Allow multiple disjoint URLs in overrides

charliermarsh · charliermarsh · commit 19b25f50116f · 2024-12-13T21:00:44.000-05:00
diff --git a/crates/uv-resolver/src/resolver/mod.rs b/crates/uv-resolver/src/resolver/mod.rs
@@ -2245,10 +2245,10 @@ impl ForkState {
                 // requirement was a URL requirement. `Urls` applies canonicalization to this and
                 // override URLs to both URL and registry requirements, which we then check for
                 // conflicts using [`ForkUrl`].
-                if let Some(url) = urls.get_url(&self.env, name, url.as_ref(), git)? {
+                for url in urls.get_url(&self.env, name, url.as_ref(), git)? {
                     self.fork_urls.insert(name, url, &self.env)?;
                     has_url = true;
-                };
+                }
 
                 // If the package is pinned to an exact index, add it to the fork.
                 for index in indexes.get(name, &self.env) {
diff --git a/crates/uv-resolver/src/resolver/urls.rs b/crates/uv-resolver/src/resolver/urls.rs
@@ -1,5 +1,4 @@
-use std::iter;
-
+use either::Either;
 use rustc_hash::FxHashMap;
 use same_file::is_same_file;
 use tracing::debug;
@@ -8,7 +7,7 @@ use uv_cache_key::CanonicalUrl;
 use uv_distribution_types::Verbatim;
 use uv_git::GitResolver;
 use uv_normalize::PackageName;
-use uv_pep508::VerbatimUrl;
+use uv_pep508::{MarkerTree, VerbatimUrl};
 use uv_pypi_types::{ParsedDirectoryUrl, ParsedUrl, VerbatimParsedUrl};
 
 use crate::{DependencyMode, Manifest, ResolveError, ResolverEnvironment};
@@ -24,10 +23,10 @@ use crate::{DependencyMode, Manifest, ResolveError, ResolverEnvironment};
 /// [`crate::fork_urls::ForkUrls`].
 #[derive(Debug, Default)]
 pub(crate) struct Urls {
-    /// URL requirements in overrides. There can only be a single URL per package in overrides
-    /// (since it replaces all other URLs), and an override URL replaces all requirements and
-    /// constraints URLs.
-    overrides: FxHashMap<PackageName, VerbatimParsedUrl>,
+    /// URL requirements in overrides. An override URL replaces all requirements and constraints
+    /// URLs. There can be multiple URLs for the same package as long as they are in different
+    /// forks.
+    overrides: FxHashMap<PackageName, Vec<(MarkerTree, VerbatimParsedUrl)>>,
     /// URLs from regular requirements or from constraints. There can be multiple URLs for the same
     /// package as long as they are in different forks.
     regular: FxHashMap<PackageName, Vec<VerbatimParsedUrl>>,
@@ -41,7 +40,8 @@ impl Urls {
         dependencies: DependencyMode,
     ) -> Result<Self, ResolveError> {
         let mut urls: FxHashMap<PackageName, Vec<VerbatimParsedUrl>> = FxHashMap::default();
-        let mut overrides: FxHashMap<PackageName, VerbatimParsedUrl> = FxHashMap::default();
+        let mut overrides: FxHashMap<PackageName, Vec<(MarkerTree, VerbatimParsedUrl)>> =
+            FxHashMap::default();
 
         // Add all direct regular requirements and constraints URL.
         for requirement in manifest.requirements_no_overrides(env, dependencies) {
@@ -86,16 +86,10 @@ impl Urls {
             // a requirements.txt entry `./anyio`, we still use the URL. See
             // `allow_recursive_url_local_path_override_constraint`.
             urls.remove(&requirement.name);
-            let previous = overrides.insert(requirement.name.clone(), url.clone());
-            if let Some(previous) = previous {
-                if !same_resource(&previous.parsed_url, &url.parsed_url, git) {
-                    return Err(ResolveError::ConflictingOverrideUrls(
-                        requirement.name.clone(),
-                        previous.verbatim.verbatim().to_string(),
-                        url.verbatim.verbatim().to_string(),
-                    ));
-                }
-            }
+            overrides
+                .entry(requirement.name.clone())
+                .or_default()
+                .push((requirement.marker, url));
         }
 
         Ok(Self {
@@ -104,41 +98,47 @@ impl Urls {
         })
     }
 
-    /// Check and canonicalize the URL of a requirement.
+    /// Return an iterator over the allowed URLs for the given package.
     ///
     /// If we have a URL override, apply it unconditionally for registry and URL requirements.
-    /// Otherwise, there are two case: For a URL requirement (`url` isn't `None`), check that the
-    /// URL is allowed and return its canonical form. For registry requirements, we return `None`
-    /// if there is no override.
+    /// Otherwise, there are two case: for a URL requirement (`url` isn't `None`), check that the
+    /// URL is allowed and return its canonical form.
+    ///
+    /// For registry requirements, we return an empty iterator.
     pub(crate) fn get_url<'a>(
         &'a self,
-        env: &ResolverEnvironment,
+        env: &'a ResolverEnvironment,
         name: &'a PackageName,
         url: Option<&'a VerbatimParsedUrl>,
         git: &'a GitResolver,
-    ) -> Result<Option<&'a VerbatimParsedUrl>, ResolveError> {
-        if let Some(override_url) = self.get_override(name) {
-            Ok(Some(override_url))
+    ) -> Result<impl Iterator<Item = &'a VerbatimParsedUrl>, ResolveError> {
+        if let Some(override_urls) = self.get_overrides(name) {
+            Ok(Either::Left(Either::Left(
+                override_urls.into_iter().filter_map(|(marker, url)| {
+                    if env.included_by_marker(*marker) {
+                        Some(url)
+                    } else {
+                        None
+                    }
+                }),
+            )))
         } else if let Some(url) = url {
-            Ok(Some(self.canonicalize_allowed_url(
-                env,
-                name,
-                git,
-                &url.verbatim,
-                &url.parsed_url,
-            )?))
+            let url =
+                self.canonicalize_allowed_url(env, name, git, &url.verbatim, &url.parsed_url)?;
+            Ok(Either::Left(Either::Right(std::iter::once(url))))
         } else {
-            Ok(None)
+            Ok(Either::Right(std::iter::empty()))
         }
     }
 
+    /// Return `true` if the package has any URL (from overrides or regular requirements).
     pub(crate) fn any_url(&self, name: &PackageName) -> bool {
-        self.get_override(name).is_some() || self.get_regular(name).is_some()
+        self.get_overrides(name).is_some() || self.get_regular(name).is_some()
     }
 
     /// Return the [`VerbatimUrl`] override for the given package, if any.
-    fn get_override(&self, package: &PackageName) -> Option<&VerbatimParsedUrl> {
-        self.overrides.get(package)
+    fn get_overrides(&self, package: &PackageName) -> Option<&[(MarkerTree, VerbatimParsedUrl)]> {
+        self.overrides.get(package).map(Vec::as_slice)
     }
 
     /// Return the allowed [`VerbatimUrl`]s for given package from regular requirements and
@@ -174,7 +174,7 @@ impl Urls {
             let mut conflicting_urls: Vec<_> = matching_urls
                 .into_iter()
                 .map(|parsed_url| parsed_url.verbatim.verbatim().to_string())
-                .chain(iter::once(verbatim_url.verbatim().to_string()))
+                .chain(std::iter::once(verbatim_url.verbatim().to_string()))
                 .collect();
             conflicting_urls.sort();
             return Err(ResolveError::ConflictingUrls {
diff --git a/crates/uv/tests/it/pip_compile.rs b/crates/uv/tests/it/pip_compile.rs
@@ -13858,6 +13858,84 @@ fn universal_disjoint_deprecated_markers() -> Result<()> {
     Ok(())
 }
 
+#[test]
+fn universal_disjoint_override_urls() -> Result<()> {
+    let context = TestContext::new("3.12");
+    let requirements_in = context.temp_dir.child("requirements.in");
+    requirements_in.write_str(indoc::indoc! {r"
+        anyio
+    "})?;
+
+    let overrides_txt = context.temp_dir.child("overrides.txt");
+    overrides_txt.write_str(indoc::indoc! {r"
+        sniffio @ https://files.pythonhosted.org/packages/e9/44/75a9c9421471a6c4805dbf2356f7c181a29c1879239abab1ea2cc8f38b40/sniffio-1.3.1-py3-none-any.whl ; sys_platform == 'win32'
+        sniffio @ https://files.pythonhosted.org/packages/c3/a0/5dba8ed157b0136607c7f2151db695885606968d1fae123dc3391e0cfdbf/sniffio-1.3.0-py3-none-any.whl ; sys_platform == 'darwin'
+    "})?;
+
+    uv_snapshot!(context.filters(), context.pip_compile()
+            .arg("requirements.in")
+            .arg("--overrides")
+            .arg("overrides.txt")
+            .arg("--universal"), @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 --overrides overrides.txt --universal
+    anyio==4.3.0
+        # via -r requirements.in
+    idna==3.6
+        # via anyio
+    sniffio @ https://files.pythonhosted.org/packages/c3/a0/5dba8ed157b0136607c7f2151db695885606968d1fae123dc3391e0cfdbf/sniffio-1.3.0-py3-none-any.whl ; sys_platform == 'darwin'
+        # via
+        #   --override overrides.txt
+        #   anyio
+    sniffio @ https://files.pythonhosted.org/packages/e9/44/75a9c9421471a6c4805dbf2356f7c181a29c1879239abab1ea2cc8f38b40/sniffio-1.3.1-py3-none-any.whl ; sys_platform == 'win32'
+        # via
+        #   --override overrides.txt
+        #   anyio
+
+    ----- stderr -----
+    Resolved 4 packages in [TIME]
+    "###
+    );
+
+    Ok(())
+}
+
+#[test]
+fn universal_conflicting_override_urls() -> Result<()> {
+    let context = TestContext::new("3.12");
+    let requirements_in = context.temp_dir.child("requirements.in");
+    requirements_in.write_str(indoc::indoc! {r"
+        anyio
+    "})?;
+
+    let overrides_txt = context.temp_dir.child("overrides.txt");
+    overrides_txt.write_str(indoc::indoc! {r"
+        sniffio @ https://files.pythonhosted.org/packages/e9/44/75a9c9421471a6c4805dbf2356f7c181a29c1879239abab1ea2cc8f38b40/sniffio-1.3.1-py3-none-any.whl ; sys_platform == 'win32'
+        sniffio @ https://files.pythonhosted.org/packages/c3/a0/5dba8ed157b0136607c7f2151db695885606968d1fae123dc3391e0cfdbf/sniffio-1.3.0-py3-none-any.whl ; sys_platform == 'darwin' or sys_platform == 'win32'
+    "})?;
+
+    uv_snapshot!(context.filters(), context.pip_compile()
+            .arg("requirements.in")
+            .arg("--overrides")
+            .arg("overrides.txt")
+            .arg("--universal"), @r###"
+    success: false
+    exit_code: 2
+    ----- stdout -----
+
+    ----- stderr -----
+    error: Requirements contain conflicting URLs for package `sniffio` in split `sys_platform == 'win32'`:
+    - https://files.pythonhosted.org/packages/c3/a0/5dba8ed157b0136607c7f2151db695885606968d1fae123dc3391e0cfdbf/sniffio-1.3.0-py3-none-any.whl
+    - https://files.pythonhosted.org/packages/e9/44/75a9c9421471a6c4805dbf2356f7c181a29c1879239abab1ea2cc8f38b40/sniffio-1.3.1-py3-none-any.whl
+    "###
+    );
+
+    Ok(())
+}
+
 #[test]
 fn compile_lowest_extra_unpinned_warning() -> Result<()> {
     let context = TestContext::new("3.12");
