Allow conflicting prerelease strategies when forking (#5150)

## Summary

Similar to #5232, we should also
track prerelease strategies per-fork, instead of globally per package.
The common functionality for tracking locals and prerelease versions
across forks is extracted into the `ForkMap` type.

Resolves #4579. This doesn't quite
solve #4959, as that issue relies
on overlapping markers.

crates/uv-resolver/src/candidate_selector.rs

@@ -12,10 +12,10 @@ use uv_normalize::PackageName;
 use uv_types::InstalledPackagesProvider;
 
 use crate::preferences::Preferences;
-use crate::prerelease_mode::PreReleaseStrategy;
+use crate::prerelease_mode::{AllowPreRelease, PreReleaseStrategy};
 use crate::resolution_mode::ResolutionStrategy;
 use crate::version_map::{VersionMap, VersionMapDistHandle};
-use crate::{Exclusions, Manifest, Options};
+use crate::{Exclusions, Manifest, Options, ResolverMarkers};
 
 #[derive(Debug, Clone)]
 #[allow(clippy::struct_field_names)]
@@ -68,13 +68,6 @@ impl CandidateSelector {
     }
 }
 
-#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
-enum AllowPreRelease {
-    Yes,
-    No,
-    IfNecessary,
-}
-
 impl CandidateSelector {
     /// Select a [`Candidate`] from a set of candidate versions and files.
     ///
@@ -89,35 +82,52 @@ impl CandidateSelector {
         preferences: &'a Preferences,
         installed_packages: &'a InstalledPackages,
         exclusions: &'a Exclusions,
+        markers: &ResolverMarkers,
     ) -> Option<Candidate<'a>> {
-        if let Some(preferred) = Self::get_preferred(
+        if let Some(preferred) = self.get_preferred(
             package_name,
             range,
             version_maps,
             preferences,
             installed_packages,
             exclusions,
+            markers,
         ) {
             return Some(preferred);
         }
 
-        self.select_no_preference(package_name, range, version_maps)
+        self.select_no_preference(package_name, range, version_maps, markers)
     }
 
     /// Get a preferred version if one exists. This is the preference from a lockfile or a locally
     /// installed version.
     fn get_preferred<'a, InstalledPackages: InstalledPackagesProvider>(
+        &self,
         package_name: &'a PackageName,
         range: &Range<Version>,
         version_maps: &'a [VersionMap],
         preferences: &'a Preferences,
         installed_packages: &'a InstalledPackages,
         exclusions: &'a Exclusions,
+        markers: &ResolverMarkers,
     ) -> Option<Candidate<'a>> {
         // If the package has a preference (e.g., an existing version from an existing lockfile),
         // and the preference satisfies the current range, use that.
         if let Some(version) = preferences.version(package_name) {
-            if range.contains(version) {
+            'preference: {
+                // Respect the version range for this requirement.
+                if !range.contains(version) {
+                    break 'preference;
+                }
+
+                // Respect the pre-release strategy for this fork.
+                if version.any_prerelease()
+                    && self.prerelease_strategy.allows(package_name, markers)
+                        != AllowPreRelease::Yes
+                {
+                    break 'preference;
+                }
+
                 // Check for a locally installed distribution that matches the preferred version
                 if !exclusions.contains(package_name) {
                     let installed_dists = installed_packages.get_packages(package_name);
@@ -167,16 +177,27 @@ impl CandidateSelector {
                 [] => {}
                 [dist] => {
                     let version = dist.version();
-                    if range.contains(version) {
-                        debug!("Found installed version of {dist} that satisfies {range}");
-
-                        return Some(Candidate {
-                            name: package_name,
-                            version,
-                            dist: CandidateDist::Compatible(CompatibleDist::InstalledDist(dist)),
-                            choice_kind: VersionChoiceKind::Installed,
-                        });
+
+                    // Respect the version range for this requirement.
+                    if !range.contains(version) {
+                        return None;
                     }
+
+                    // Respect the pre-release strategy for this fork.
+                    if version.any_prerelease()
+                        && self.prerelease_strategy.allows(package_name, markers)
+                            != AllowPreRelease::Yes
+                    {
+                        return None;
+                    }
+
+                    debug!("Found installed version of {dist} that satisfies {range}");
+                    return Some(Candidate {
+                        name: package_name,
+                        version,
+                        dist: CandidateDist::Compatible(CompatibleDist::InstalledDist(dist)),
+                        choice_kind: VersionChoiceKind::Installed,
+                    });
                 }
                 // We do not consider installed distributions with multiple versions because
                 // during installation these must be reinstalled from the remote
@@ -189,43 +210,21 @@ impl CandidateSelector {
         None
     }
 
-    /// Determine the appropriate prerelease strategy for the current package.
-    fn allow_prereleases(&self, package_name: &PackageName) -> AllowPreRelease {
-        match &self.prerelease_strategy {
-            PreReleaseStrategy::Disallow => AllowPreRelease::No,
-            PreReleaseStrategy::Allow => AllowPreRelease::Yes,
-            PreReleaseStrategy::IfNecessary => AllowPreRelease::IfNecessary,
-            PreReleaseStrategy::Explicit(packages) => {
-                if packages.contains(package_name) {
-                    AllowPreRelease::Yes
-                } else {
-                    AllowPreRelease::No
-                }
-            }
-            PreReleaseStrategy::IfNecessaryOrExplicit(packages) => {
-                if packages.contains(package_name) {
-                    AllowPreRelease::Yes
-                } else {
-                    AllowPreRelease::IfNecessary
-                }
-            }
-        }
-    }
-
     /// Select a [`Candidate`] without checking for version preference such as an existing
     /// lockfile.
     pub(crate) fn select_no_preference<'a>(
         &'a self,
         package_name: &'a PackageName,
         range: &Range<Version>,
         version_maps: &'a [VersionMap],
+        markers: &ResolverMarkers,
     ) -> Option<Candidate> {
         tracing::trace!(
             "selecting candidate for package {package_name} with range {range:?} with {} remote versions",
             version_maps.iter().map(VersionMap::len).sum::<usize>(),
         );
         let highest = self.use_highest_version(package_name);
-        let allow_prerelease = self.allow_prereleases(package_name);
+        let allow_prerelease = self.prerelease_strategy.allows(package_name, markers);
 
         if self.index_strategy == IndexStrategy::UnsafeBestMatch {
             if highest {

crates/uv-resolver/src/error.rs

@@ -232,6 +232,7 @@ impl std::fmt::Display for NoSolutionError {
             &self.unavailable_packages,
             &self.incomplete_packages,
             &self.fork_urls,
+            &self.markers,
         ) {
             write!(f, "\n\n{hint}")?;
         }

crates/uv-resolver/src/prerelease_mode.rs

@@ -1,10 +1,10 @@
 use pypi_types::RequirementSource;
-use rustc_hash::FxHashSet;
 
 use pep508_rs::MarkerEnvironment;
 use uv_normalize::PackageName;
 
-use crate::{DependencyMode, Manifest};
+use crate::resolver::ForkSet;
+use crate::{DependencyMode, Manifest, ResolverMarkers};
 
 #[derive(Debug, Default, Clone, Copy, PartialEq, Eq, serde::Deserialize)]
 #[serde(deny_unknown_fields, rename_all = "kebab-case")]
@@ -57,11 +57,11 @@ pub(crate) enum PreReleaseStrategy {
 
     /// Allow pre-release versions for first-party packages with explicit pre-release markers in
     /// their version requirements.
-    Explicit(FxHashSet<PackageName>),
+    Explicit(ForkSet),
 
     /// Allow pre-release versions if all versions of a package are pre-release, or if the package
     /// has an explicit pre-release marker in its version requirements.
-    IfNecessaryOrExplicit(FxHashSet<PackageName>),
+    IfNecessaryOrExplicit(ForkSet),
 }
 
 impl PreReleaseStrategy {
@@ -71,51 +71,72 @@ impl PreReleaseStrategy {
         markers: Option<&MarkerEnvironment>,
         dependencies: DependencyMode,
     ) -> Self {
+        let mut packages = ForkSet::default();
+
         match mode {
             PreReleaseMode::Disallow => Self::Disallow,
             PreReleaseMode::Allow => Self::Allow,
             PreReleaseMode::IfNecessary => Self::IfNecessary,
-            PreReleaseMode::Explicit => Self::Explicit(
-                manifest
-                    .requirements(markers, dependencies)
-                    .filter(|requirement| {
-                        let RequirementSource::Registry { specifier, .. } = &requirement.source
-                        else {
-                            return false;
-                        };
-                        specifier
-                            .iter()
-                            .any(pep440_rs::VersionSpecifier::any_prerelease)
-                    })
-                    .map(|requirement| requirement.name.clone())
-                    .collect(),
-            ),
-            PreReleaseMode::IfNecessaryOrExplicit => Self::IfNecessaryOrExplicit(
-                manifest
-                    .requirements(markers, dependencies)
-                    .filter(|requirement| {
-                        let RequirementSource::Registry { specifier, .. } = &requirement.source
-                        else {
-                            return false;
-                        };
-                        specifier
-                            .iter()
-                            .any(pep440_rs::VersionSpecifier::any_prerelease)
-                    })
-                    .map(|requirement| requirement.name.clone())
-                    .collect(),
-            ),
+            _ => {
+                for requirement in manifest.requirements(markers, dependencies) {
+                    let RequirementSource::Registry { specifier, .. } = &requirement.source else {
+                        continue;
+                    };
+
+                    if specifier
+                        .iter()
+                        .any(pep440_rs::VersionSpecifier::any_prerelease)
+                    {
+                        packages.add(&requirement, ());
+                    }
+                }
+
+                match mode {
+                    PreReleaseMode::Explicit => Self::Explicit(packages),
+                    PreReleaseMode::IfNecessaryOrExplicit => Self::IfNecessaryOrExplicit(packages),
+                    _ => unreachable!(),
+                }
+            }
         }
     }
 
     /// Returns `true` if a [`PackageName`] is allowed to have pre-release versions.
-    pub(crate) fn allows(&self, package: &PackageName) -> bool {
+    pub(crate) fn allows(
+        &self,
+        package_name: &PackageName,
+        markers: &ResolverMarkers,
+    ) -> AllowPreRelease {
         match self {
-            Self::Disallow => false,
-            Self::Allow => true,
-            Self::IfNecessary => false,
-            Self::Explicit(packages) => packages.contains(package),
-            Self::IfNecessaryOrExplicit(packages) => packages.contains(package),
+            PreReleaseStrategy::Disallow => AllowPreRelease::No,
+            PreReleaseStrategy::Allow => AllowPreRelease::Yes,
+            PreReleaseStrategy::IfNecessary => AllowPreRelease::IfNecessary,
+            PreReleaseStrategy::Explicit(packages) => {
+                if packages.contains(package_name, markers) {
+                    AllowPreRelease::Yes
+                } else {
+                    AllowPreRelease::No
+                }
+            }
+            PreReleaseStrategy::IfNecessaryOrExplicit(packages) => {
+                if packages.contains(package_name, markers) {
+                    AllowPreRelease::Yes
+                } else {
+                    AllowPreRelease::IfNecessary
+                }
+            }
         }
     }
 }
+
+/// The pre-release strategy for a given package.
+#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
+pub(crate) enum AllowPreRelease {
+    /// Allow all pre-release versions.
+    Yes,
+
+    /// Disallow all pre-release versions.
+    No,
+
+    /// Allow pre-release versions if all versions of this package are pre-release.
+    IfNecessary,
+}

crates/uv-resolver/src/pubgrub/report.rs

@@ -18,9 +18,10 @@ use uv_normalize::PackageName;
 
 use crate::candidate_selector::CandidateSelector;
 use crate::fork_urls::ForkUrls;
+use crate::prerelease_mode::AllowPreRelease;
 use crate::python_requirement::{PythonRequirement, PythonTarget};
 use crate::resolver::{IncompletePackage, UnavailablePackage, UnavailableReason};
-use crate::RequiresPython;
+use crate::{RequiresPython, ResolverMarkers};
 
 use super::{PubGrubPackage, PubGrubPackageInner, PubGrubPython};
 
@@ -407,6 +408,7 @@ impl PubGrubReportFormatter<'_> {
         unavailable_packages: &FxHashMap<PackageName, UnavailablePackage>,
         incomplete_packages: &FxHashMap<PackageName, BTreeMap<Version, IncompletePackage>>,
         fork_urls: &ForkUrls,
+        markers: &ResolverMarkers,
     ) -> IndexSet<PubGrubHint> {
         let mut hints = IndexSet::default();
         match derivation_tree {
@@ -416,7 +418,9 @@ impl PubGrubReportFormatter<'_> {
                 if let PubGrubPackageInner::Package { name, .. } = &**package {
                     // Check for no versions due to pre-release options.
                     if !fork_urls.contains_key(name) {
-                        self.prerelease_available_hint(package, name, set, selector, &mut hints);
+                        self.prerelease_available_hint(
+                            package, name, set, selector, markers, &mut hints,
+                        );
                     }
                 }
 
@@ -465,6 +469,7 @@ impl PubGrubReportFormatter<'_> {
                     unavailable_packages,
                     incomplete_packages,
                     fork_urls,
+                    markers,
                 ));
                 hints.extend(self.hints(
                     &derived.cause2,
@@ -473,6 +478,7 @@ impl PubGrubReportFormatter<'_> {
                     unavailable_packages,
                     incomplete_packages,
                     fork_urls,
+                    markers,
                 ));
             }
         }
@@ -569,6 +575,7 @@ impl PubGrubReportFormatter<'_> {
         name: &PackageName,
         set: &Range<Version>,
         selector: &CandidateSelector,
+        markers: &ResolverMarkers,
         hints: &mut IndexSet<PubGrubHint>,
     ) {
         let any_prerelease = set.iter().any(|(start, end)| {
@@ -587,7 +594,7 @@ impl PubGrubReportFormatter<'_> {
 
         if any_prerelease {
             // A pre-release marker appeared in the version requirements.
-            if !selector.prerelease_strategy().allows(name) {
+            if selector.prerelease_strategy().allows(name, markers) != AllowPreRelease::Yes {
                 hints.insert(PubGrubHint::PreReleaseRequested {
                     package: package.clone(),
                     range: self.simplify_set(set, package).into_owned(),
@@ -601,7 +608,7 @@ impl PubGrubReportFormatter<'_> {
                 .find(|version| set.contains(version))
         }) {
             // There are pre-release versions available for the package.
-            if !selector.prerelease_strategy().allows(name) {
+            if selector.prerelease_strategy().allows(name, markers) != AllowPreRelease::Yes {
                 hints.insert(PubGrubHint::PreReleaseAvailable {
                     package: package.clone(),
                     version: version.clone(),