From 1bc5485094342fb859bab6427331f0bc2521f1bc Mon Sep 17 00:00:00 2001
From: Charlie Marsh <charlie.r.marsh@gmail.com>
Date: Wed, 28 Feb 2024 20:16:23 -0500
Subject: [PATCH] Allow pre-releases for requirements in constraints files
 (#2069)

## Summary

If a pre-release marker is present on a requirement in a constraint
file, we should allow pre-releases for that package.

Closes https://github.com/astral-sh/uv/issues/2063.
---
 crates/uv-resolver/src/candidate_selector.rs |  2 +
 crates/uv-resolver/src/prerelease_mode.rs    | 15 +++++--
 crates/uv/tests/pip_compile.rs               | 41 ++++++++++++++++++++
 3 files changed, 55 insertions(+), 3 deletions(-)

diff --git a/crates/uv-resolver/src/candidate_selector.rs b/crates/uv-resolver/src/candidate_selector.rs
index 1a3bd285b..d538f421c 100644
--- a/crates/uv-resolver/src/candidate_selector.rs
+++ b/crates/uv-resolver/src/candidate_selector.rs
@@ -32,6 +32,8 @@ impl CandidateSelector {
             prerelease_strategy: PreReleaseStrategy::from_mode(
                 options.prerelease_mode,
                 manifest.requirements.as_slice(),
+                manifest.constraints.as_slice(),
+                manifest.overrides.as_slice(),
             ),
             preferences: Preferences::from(manifest.preferences.as_slice()),
         }
diff --git a/crates/uv-resolver/src/prerelease_mode.rs b/crates/uv-resolver/src/prerelease_mode.rs
index e0ec0f604..858fd5dc9 100644
--- a/crates/uv-resolver/src/prerelease_mode.rs
+++ b/crates/uv-resolver/src/prerelease_mode.rs
@@ -48,14 +48,21 @@ pub(crate) enum PreReleaseStrategy {
 }
 
 impl PreReleaseStrategy {
-    pub(crate) fn from_mode(mode: PreReleaseMode, direct_dependencies: &[Requirement]) -> Self {
+    pub(crate) fn from_mode(
+        mode: PreReleaseMode,
+        requirements: &[Requirement],
+        constraints: &[Requirement],
+        overrides: &[Requirement],
+    ) -> Self {
         match mode {
             PreReleaseMode::Disallow => Self::Disallow,
             PreReleaseMode::Allow => Self::Allow,
             PreReleaseMode::IfNecessary => Self::IfNecessary,
             PreReleaseMode::Explicit => Self::Explicit(
-                direct_dependencies
+                requirements
                     .iter()
+                    .chain(constraints.iter())
+                    .chain(overrides.iter())
                     .filter(|requirement| {
                         let Some(version_or_url) = &requirement.version_or_url else {
                             return false;
@@ -74,8 +81,10 @@ impl PreReleaseStrategy {
                     .collect(),
             ),
             PreReleaseMode::IfNecessaryOrExplicit => Self::IfNecessaryOrExplicit(
-                direct_dependencies
+                requirements
                     .iter()
+                    .chain(constraints.iter())
+                    .chain(overrides.iter())
                     .filter(|requirement| {
                         let Some(version_or_url) = &requirement.version_or_url else {
                             return false;
diff --git a/crates/uv/tests/pip_compile.rs b/crates/uv/tests/pip_compile.rs
index 4b5538f51..0b5721e28 100644
--- a/crates/uv/tests/pip_compile.rs
+++ b/crates/uv/tests/pip_compile.rs
@@ -4386,6 +4386,47 @@ fn pre_alias() -> Result<()> {
     Ok(())
 }
 
+/// Allow a pre-release for a version specifier in a constraint file.
+#[test]
+fn pre_release_constraint() -> Result<()> {
+    let context = TestContext::new("3.12");
+    let requirements_in = context.temp_dir.child("requirements.in");
+    requirements_in.write_str("flask")?;
+
+    let constraints_txt = context.temp_dir.child("constraints.txt");
+    constraints_txt.write_str("flask<=2.0.0rc2")?;
+
+    uv_snapshot!(context.compile()
+            .arg("requirements.in")
+            .arg("--constraint")
+            .arg("constraints.txt"), @r###"
+    success: true
+    exit_code: 0
+    ----- stdout -----
+    # This file was autogenerated by uv via the following command:
+    #    uv pip compile --cache-dir [CACHE_DIR] --exclude-newer 2023-11-18T12:00:00Z requirements.in --constraint constraints.txt
+    click==8.1.7
+        # via flask
+    flask==2.0.0rc2
+    itsdangerous==2.1.2
+        # via flask
+    jinja2==3.1.2
+        # via flask
+    markupsafe==2.1.3
+        # via
+        #   jinja2
+        #   werkzeug
+    werkzeug==3.0.1
+        # via flask
+
+    ----- stderr -----
+    Resolved 6 packages in [TIME]
+    "###
+    );
+
+    Ok(())
+}
+
 /// Resolve from a `pyproject.toml` file with a recursive extra.
 #[test]
 fn compile_pyproject_toml_recursive_extra() -> Result<()> {