Backtrack to non-local versions when wheels are missing platform support (#10046)

## Summary

This is yet another variation on
#9928, with a few minor changes:

1. It only applies to local versions (e.g., `2.5.1+cpu`).
2. It only _considers_ the non-local version as an alternative (e.g.,
`2.5.1`).
3. It only _considers_ the non-local alternative if it _does_ support
the unsupported platform.
4. Instead of failing, it falls back to using the local version.

So, this is far less strict, and is effectively designed to solve
PyTorch but nothing else. It's also not user-configurable, except by way
of using `environments` to exclude platforms.

Author: charliermarsh

crates/uv-distribution-types/src/prioritized_distribution.rs

@@ -1,7 +1,8 @@
 use std::fmt::{Display, Formatter};
-use uv_distribution_filename::BuildTag;
+use uv_distribution_filename::{BuildTag, WheelFilename};
 
 use uv_pep440::VersionSpecifiers;
+use uv_pep508::{MarkerExpression, MarkerOperator, MarkerTree, MarkerValueString};
 use uv_platform_tags::{IncompatibleTag, TagPriority};
 use uv_pypi_types::{HashDigest, Yanked};
 
@@ -14,7 +15,7 @@ use crate::{
 pub struct PrioritizedDist(Box<PrioritizedDistInner>);
 
 /// [`PrioritizedDist`] is boxed because [`Dist`] is large.
-#[derive(Debug, Default, Clone)]
+#[derive(Debug, Clone)]
 struct PrioritizedDistInner {
     /// The highest-priority source distribution. Between compatible source distributions this priority is arbitrary.
     source: Option<(RegistrySourceDist, SourceDistCompatibility)>,
@@ -25,6 +26,20 @@ struct PrioritizedDistInner {
     wheels: Vec<(RegistryBuiltWheel, WheelCompatibility)>,
     /// The hashes for each distribution.
     hashes: Vec<HashDigest>,
+    /// The set of supported platforms for the distribution, described in terms of their markers.
+    markers: MarkerTree,
+}
+
+impl Default for PrioritizedDistInner {
+    fn default() -> Self {
+        Self {
+            source: None,
+            best_wheel_index: None,
+            wheels: Vec::new(),
+            hashes: Vec::new(),
+            markers: MarkerTree::FALSE,
+        }
+    }
 }
 
 /// A distribution that can be used for both resolution and installation.
@@ -70,6 +85,16 @@ impl CompatibleDist<'_> {
             CompatibleDist::IncompatibleWheel { sdist, .. } => sdist.file.requires_python.as_ref(),
         }
     }
+
+    /// Return the set of supported platform the distribution, in terms of their markers.
+    pub fn implied_markers(&self) -> MarkerTree {
+        match self {
+            CompatibleDist::InstalledDist(_) => MarkerTree::TRUE,
+            CompatibleDist::SourceDist { prioritized, .. } => prioritized.0.markers,
+            CompatibleDist::CompatibleWheel { prioritized, .. } => prioritized.0.markers,
+            CompatibleDist::IncompatibleWheel { prioritized, .. } => prioritized.0.markers,
+        }
+    }
 }
 
 #[derive(Debug, PartialEq, Eq, Clone)]
@@ -257,6 +282,7 @@ impl PrioritizedDist {
         compatibility: WheelCompatibility,
     ) -> Self {
         Self(Box::new(PrioritizedDistInner {
+            markers: implied_markers(&dist.filename),
             best_wheel_index: Some(0),
             wheels: vec![(dist, compatibility)],
             source: None,
@@ -271,6 +297,7 @@ impl PrioritizedDist {
         compatibility: SourceDistCompatibility,
     ) -> Self {
         Self(Box::new(PrioritizedDistInner {
+            markers: MarkerTree::TRUE,
             best_wheel_index: None,
             wheels: vec![],
             source: Some((dist, compatibility)),
@@ -293,8 +320,11 @@ impl PrioritizedDist {
         } else {
             self.0.best_wheel_index = Some(self.0.wheels.len());
         }
-        self.0.wheels.push((dist, compatibility));
         self.0.hashes.extend(hashes);
+        if !self.0.markers.is_true() {
+            self.0.markers.or(implied_markers(&dist.filename));
+        }
+        self.0.wheels.push((dist, compatibility));
     }
 
     /// Insert the given source distribution into the [`PrioritizedDist`].
@@ -312,7 +342,9 @@ impl PrioritizedDist {
         } else {
             self.0.source = Some((dist, compatibility));
         }
-
+        if !self.0.markers.is_true() {
+            self.0.markers.or(MarkerTree::TRUE);
+        }
         self.0.hashes.extend(hashes);
     }
 
@@ -563,6 +595,7 @@ impl IncompatibleSource {
 }
 
 impl IncompatibleWheel {
+    #[allow(clippy::match_like_matches_macro)]
     fn is_more_compatible(&self, other: &Self) -> bool {
         match self {
             Self::ExcludeNewer(timestamp_self) => match other {
@@ -599,3 +632,39 @@ impl IncompatibleWheel {
         }
     }
 }
+
+/// Given a wheel filename, determine the set of supported platforms, in terms of their markers.
+pub fn implied_markers(filename: &WheelFilename) -> MarkerTree {
+    let mut marker = MarkerTree::FALSE;
+    for platform_tag in &filename.platform_tag {
+        match platform_tag.as_str() {
+            "any" => marker.or(MarkerTree::TRUE),
+            tag if tag.starts_with("win") => {
+                marker.or(MarkerTree::expression(MarkerExpression::String {
+                    key: MarkerValueString::SysPlatform,
+                    operator: MarkerOperator::Equal,
+                    value: "win32".to_string(),
+                }));
+            }
+            tag if tag.starts_with("macosx") => {
+                marker.or(MarkerTree::expression(MarkerExpression::String {
+                    key: MarkerValueString::SysPlatform,
+                    operator: MarkerOperator::Equal,
+                    value: "darwin".to_string(),
+                }));
+            }
+            tag if tag.starts_with("manylinux")
+                || tag.starts_with("musllinux")
+                || tag.starts_with("linux") =>
+            {
+                marker.or(MarkerTree::expression(MarkerExpression::String {
+                    key: MarkerValueString::SysPlatform,
+                    operator: MarkerOperator::Equal,
+                    value: "linux".to_string(),
+                }));
+            }
+            _ => {}
+        }
+    }
+    marker
+}

crates/uv-resolver/src/resolver/availability.rs

@@ -3,8 +3,7 @@ use std::fmt::{Display, Formatter};
 use uv_distribution_types::IncompatibleDist;
 use uv_pep440::{Version, VersionSpecifiers};
 
-use crate::resolver::MetadataUnavailable;
-use crate::ResolverEnvironment;
+use crate::resolver::{MetadataUnavailable, VersionFork};
 
 /// The reason why a package or a version cannot be used.
 #[derive(Debug, Clone, Eq, PartialEq)]
@@ -170,6 +169,6 @@ pub(crate) enum ResolverVersion {
     Unavailable(Version, UnavailableVersion),
     /// A usable version
     Unforked(Version),
-    /// A set of forks.
-    Forked(Vec<ResolverEnvironment>),
+    /// A set of forks, optionally with resolved versions
+    Forked(Vec<VersionFork>),
 }

crates/uv-resolver/src/resolver/environment.rs

@@ -512,7 +512,7 @@ impl<'d> Forker<'d> {
 }
 
 /// Fork the resolver based on a `Requires-Python` specifier.
-pub(crate) fn fork_python_requirement(
+pub(crate) fn fork_version_by_python_requirement(
     requires_python: &VersionSpecifiers,
     python_requirement: &PythonRequirement,
     env: &ResolverEnvironment,
@@ -555,6 +555,49 @@ pub(crate) fn fork_python_requirement(
     envs
 }
 
+/// Fork the resolver based on a marker.
+pub(crate) fn fork_version_by_marker(
+    env: &ResolverEnvironment,
+    marker: MarkerTree,
+) -> Option<(ResolverEnvironment, ResolverEnvironment)> {
+    let Kind::Universal {
+        markers: ref env_marker,
+        ..
+    } = env.kind
+    else {
+        panic!("resolver must be in universal mode for forking")
+    };
+
+    // Attempt to split based on the marker.
+    //
+    // For example, given `python_version >= '3.10'` and the split marker `sys_platform == 'linux'`,
+    // the result will be:
+    //
+    //   `python_version >= '3.10' and sys_platform == 'linux'`
+    //   `python_version >= '3.10' and sys_platform != 'linux'`
+    //
+    // If the marker is disjoint with the current environment, then we should return an empty list.
+    // If the marker complement is disjoint with the current environment, then we should also return
+    // an empty list.
+    //
+    // For example, given `python_version >= '3.10' and sys_platform == 'linux'` and the split marker
+    // `sys_platform == 'win32'`, return an empty list, since the following isn't satisfiable:
+    //
+    //   python_version >= '3.10' and sys_platform == 'linux' and sys_platform == 'win32'
+    if env_marker.is_disjoint(marker) {
+        return None;
+    }
+    let with_marker = env.narrow_environment(marker);
+
+    let complement = marker.negate();
+    if env_marker.is_disjoint(complement) {
+        return None;
+    }
+    let without_marker = env.narrow_environment(complement);
+
+    Some((with_marker, without_marker))
+}
+
 #[cfg(test)]
 mod tests {
     use std::ops::Bound;