From b8ff32f6bedb29f81f9f5cf9589c2b62351c41f1 Mon Sep 17 00:00:00 2001
From: Zanie Blue <contact@zanie.dev>
Date: Wed, 1 Nov 2023 20:20:32 -0500
Subject: [PATCH] Respect markers on constraints (#282)

Closes #252
---
 crates/puffin-cli/tests/pip_compile.rs        | 50 +++++++++++++++++++
 ..._compile__compile_constraints_markers.snap | 28 +++++++++++
 crates/puffin-resolver/src/resolver.rs        | 15 ++----
 3 files changed, 82 insertions(+), 11 deletions(-)
 create mode 100644 crates/puffin-cli/tests/snapshots/pip_compile__compile_constraints_markers.snap

diff --git a/crates/puffin-cli/tests/pip_compile.rs b/crates/puffin-cli/tests/pip_compile.rs
index a130c21ac..5417e11c3 100644
--- a/crates/puffin-cli/tests/pip_compile.rs
+++ b/crates/puffin-cli/tests/pip_compile.rs
@@ -228,6 +228,56 @@ fn compile_constraints_inline() -> Result<()> {
     Ok(())
 }
 
+/// Resolve a package from a `requirements.in` file, with a `constraints.txt` file that
+/// uses markers.
+#[test]
+fn compile_constraints_markers() -> Result<()> {
+    let temp_dir = assert_fs::TempDir::new()?;
+    let cache_dir = assert_fs::TempDir::new()?;
+    let venv = temp_dir.child(".venv");
+
+    Command::new(get_cargo_bin(BIN_NAME))
+        .arg("venv")
+        .arg(venv.as_os_str())
+        .arg("--cache-dir")
+        .arg(cache_dir.path())
+        .current_dir(&temp_dir)
+        .assert()
+        .success();
+    venv.assert(predicates::path::is_dir());
+
+    let requirements_in = temp_dir.child("requirements.in");
+    requirements_in.touch()?;
+    requirements_in.write_str("anyio")?;
+
+    // Constrain a transitive dependency based on the Python version
+    let constraints_txt = temp_dir.child("constraints.txt");
+    constraints_txt.touch()?;
+    // If constraints are ignored, these will conflict
+    constraints_txt.write_str("sniffio==1.2.0;python_version<='3.7'")?;
+    constraints_txt.write_str("sniffio==1.3.0;python_version>'3.7'")?;
+
+    insta::with_settings!({
+        filters => vec![
+            (r"(\d|\.)+(ms|s)", "[TIME]"),
+            (r"#    .* pip-compile", "#    [BIN_PATH] pip-compile"),
+            (r"--cache-dir .*", "--cache-dir [CACHE_DIR]"),
+        ]
+    }, {
+        assert_cmd_snapshot!(Command::new(get_cargo_bin(BIN_NAME))
+            .arg("pip-compile")
+            .arg("requirements.in")
+            .arg("--constraint")
+            .arg("constraints.txt")
+            .arg("--cache-dir")
+            .arg(cache_dir.path())
+            .env("VIRTUAL_ENV", venv.as_os_str())
+            .current_dir(&temp_dir));
+    });
+
+    Ok(())
+}
+
 /// Resolve a package from an optional dependency group in a `pyproject.toml` file.
 #[test]
 fn compile_pyproject_toml_extra() -> Result<()> {
diff --git a/crates/puffin-cli/tests/snapshots/pip_compile__compile_constraints_markers.snap b/crates/puffin-cli/tests/snapshots/pip_compile__compile_constraints_markers.snap
new file mode 100644
index 000000000..070d5bf32
--- /dev/null
+++ b/crates/puffin-cli/tests/snapshots/pip_compile__compile_constraints_markers.snap
@@ -0,0 +1,28 @@
+---
+source: crates/puffin-cli/tests/pip_compile.rs
+info:
+  program: puffin
+  args:
+    - pip-compile
+    - requirements.in
+    - "--constraint"
+    - constraints.txt
+    - "--cache-dir"
+    - /var/folders/bc/qlsk3t6x7c9fhhbvvcg68k9c0000gp/T/.tmp1Ts53o
+  env:
+    VIRTUAL_ENV: /var/folders/bc/qlsk3t6x7c9fhhbvvcg68k9c0000gp/T/.tmp1vzBQa/.venv
+---
+success: true
+exit_code: 0
+----- stdout -----
+# This file was autogenerated by Puffin v0.0.1 via the following command:
+#    [BIN_PATH] pip-compile requirements.in --constraint constraints.txt --cache-dir [CACHE_DIR]
+anyio==4.0.0
+idna==3.4
+    # via anyio
+sniffio==1.3.0
+    # via anyio
+
+----- stderr -----
+Resolved 3 packages in [TIME]
+
diff --git a/crates/puffin-resolver/src/resolver.rs b/crates/puffin-resolver/src/resolver.rs
index da91ba7e5..64b15385f 100644
--- a/crates/puffin-resolver/src/resolver.rs
+++ b/crates/puffin-resolver/src/resolver.rs
@@ -484,19 +484,12 @@ impl<'a, Context: BuildContext + Sync> Resolver<'a, Context> {
                     }
                 }
 
-                debug!("Got constraints: {:#?}", constraints);
-
                 // If any requirements were further constrained by the user, add those constraints.
-                for constraint in &self.constraints {
-                    let package = PubGrubPackage::Package(
-                        PackageName::normalize(&constraint.name),
-                        None,
-                        None,
-                    );
+                for (package, version) in
+                    iter_requirements(self.constraints.iter(), None, None, self.markers)
+                {
                     if let Some(range) = constraints.get_mut(&package) {
-                        *range = range.intersection(
-                            &version_range(constraint.version_or_url.as_ref()).unwrap(),
-                        );
+                        *range = range.intersection(&version);
                     }
                 }