Simplfy pre

astral-sh · Sep 8, 2024 · 6940009 · 6940009
1 parent 6ef7cdf
commit 6940009
Showing 1 changed file with 40 additions and 81 deletions.
diff --git a/crates/uv-resolver/src/candidate_selector.rs b/crates/uv-resolver/src/candidate_selector.rs
@@ -325,6 +325,25 @@ impl CandidateSelector {
         );
         let highest = self.use_highest_version(package_name);
         let allow_prerelease = self.prerelease_strategy.allows(package_name, markers);
+        let allow_prerelease = match allow_prerelease {
+            AllowPrerelease::Yes => true,
+            AllowPrerelease::No => false,
+            AllowPrerelease::IfNecessary => {
+                // If the range is a non-pre-release singleton, we won't consider pre-releases anyway.
+                if range
+                    .as_singleton()
+                    .is_some_and(|version| !version.any_prerelease())
+                {
+                    false
+                } else {
+                    version_maps.iter().any(|version_map| {
+                        version_map
+                            .versions()
+                            .all(pep440_rs::Version::any_prerelease)
+                    })
+                }
+            }
+        };
 
         if self.index_strategy == IndexStrategy::UnsafeBestMatch {
             if highest {
@@ -416,83 +435,35 @@ impl CandidateSelector {
         versions: impl Iterator<Item = (&'a Version, VersionMapDistHandle<'a>)>,
         package_name: &'a PackageName,
         range: &Range<Version>,
-        allow_prerelease: AllowPrerelease,
+        allow_prerelease: bool,
     ) -> Option<Candidate<'a>> {
-        #[derive(Debug)]
-        enum PrereleaseCandidate<'a> {
-            NotNecessary,
-            IfNecessary(&'a Version, &'a PrioritizedDist),
-        }
-
-        let mut prerelease = None;
         let mut steps = 0usize;
         for (version, maybe_dist) in versions {
             steps += 1;
             let candidate = if version.any_prerelease() {
-                if range.contains(version) {
-                    match allow_prerelease {
-                        AllowPrerelease::Yes => {
-                            let Some(dist) = maybe_dist.prioritized_dist() else {
-                                continue;
-                            };
-                            tracing::trace!(
-                                "found candidate for package {:?} with range {:?} \
-                                 after {} steps: {:?} version",
-                                package_name,
-                                range,
-                                steps,
-                                version,
-                            );
-                            // If pre-releases are allowed, treat them equivalently
-                            // to stable distributions.
-                            Candidate::new(
-                                package_name,
-                                version,
-                                dist,
-                                VersionChoiceKind::Compatible,
-                            )
-                        }
-                        AllowPrerelease::IfNecessary => {
-                            let Some(dist) = maybe_dist.prioritized_dist() else {
-                                continue;
-                            };
-                            // If pre-releases are allowed as a fallback, store the
-                            // first-matching prerelease.
-                            if prerelease.is_none() {
-                                prerelease = Some(PrereleaseCandidate::IfNecessary(version, dist));
-                            }
-                            continue;
-                        }
-                        AllowPrerelease::No => {
-                            continue;
-                        }
-                    }
-                } else {
+                if !allow_prerelease {
+                    continue;
+                }
+                if !range.contains(version) {
                     continue;
                 }
+                let Some(dist) = maybe_dist.prioritized_dist() else {
+                    continue;
+                };
+                trace!("found candidate for package {package_name:?} with range {range:?} after {steps} steps: {version:?} version");
+                // If pre-releases are allowed, treat them equivalently
+                // to stable distributions.
+                Candidate::new(package_name, version, dist, VersionChoiceKind::Compatible)
             } else {
-                // If we have at least one stable release, we shouldn't allow the "if-necessary"
-                // pre-release strategy, regardless of whether that stable release satisfies the
-                // current range.
-                prerelease = Some(PrereleaseCandidate::NotNecessary);
-
-                // Return the first-matching stable distribution.
-                if range.contains(version) {
-                    let Some(dist) = maybe_dist.prioritized_dist() else {
-                        continue;
-                    };
-                    tracing::trace!(
-                        "found candidate for package {:?} with range {:?} \
-                         after {} steps: {:?} version",
-                        package_name,
-                        range,
-                        steps,
-                        version,
-                    );
-                    Candidate::new(package_name, version, dist, VersionChoiceKind::Compatible)
-                } else {
+                if !range.contains(version) {
                     continue;
                 }
+                let Some(dist) = maybe_dist.prioritized_dist() else {
+                    continue;
+                };
+                trace!("found candidate for package {package_name:?} with range {range:?} after {steps} steps: {version:?} version");
+                // Return the first-matching stable distribution.
+                Candidate::new(package_name, version, dist, VersionChoiceKind::Compatible)
             };
 
             // If candidate is not compatible due to exclude newer, continue searching.
@@ -513,20 +484,8 @@ impl CandidateSelector {
 
             return Some(candidate);
         }
-        trace!(
-            "Exhausted all candidates for package {package_name} with range {range} \
-             after {steps} steps",
-        );
-        match prerelease {
-            None => None,
-            Some(PrereleaseCandidate::NotNecessary) => None,
-            Some(PrereleaseCandidate::IfNecessary(version, dist)) => Some(Candidate::new(
-                package_name,
-                version,
-                dist,
-                VersionChoiceKind::Compatible,
-            )),
-        }
+        trace!("Exhausted all candidates for package {package_name} with range {range} after {steps} steps");
+        None
     }
 }