Prune unreachable wheels from lockfile

When a package is included under a platform-specific marker, we know that wheels that mismatch this marker can never be installed, so we drop them from the lockfile.
astral-sh · Sep 3, 2024 · 148c19a · 148c19a
1 parent 572f670
commit 148c19a
Show file tree

Hide file tree

Showing 8 changed files with 78 additions and 104 deletions.
diff --git a/crates/uv-resolver/src/lock/mod.rs b/crates/uv-resolver/src/lock/mod.rs
@@ -1,15 +1,16 @@
+use either::Either;
+use itertools::Itertools;
+use petgraph::graph::NodeIndex;
+use petgraph::visit::EdgeRef;
+use rustc_hash::{FxHashMap, FxHashSet};
 use std::borrow::Cow;
 use std::collections::{BTreeMap, BTreeSet, VecDeque};
 use std::convert::Infallible;
 use std::fmt::{Debug, Display};
 use std::io;
 use std::path::{Path, PathBuf};
 use std::str::FromStr;
-
-use either::Either;
-use itertools::Itertools;
-use petgraph::visit::EdgeRef;
-use rustc_hash::{FxHashMap, FxHashSet};
+use std::sync::LazyLock;
 use toml_edit::{value, Array, ArrayOfTables, InlineTable, Item, Table, Value};
 use tracing::debug;
 use url::Url;
@@ -49,6 +50,25 @@ mod tree;
 /// The current version of the lockfile format.
 const VERSION: u32 = 1;
 
+static LINUX_MARKERS: LazyLock<MarkerTree> = LazyLock::new(|| {
+    MarkerTree::from_str(
+        "platform_system == 'Linux' and os_name == 'posix' and sys_platform == 'linux'",
+    )
+    .unwrap()
+});
+static WINDOWS_MARKERS: LazyLock<MarkerTree> = LazyLock::new(|| {
+    MarkerTree::from_str(
+        "platform_system == 'Windows' and os_name == 'nt' and sys_platform == 'win32'",
+    )
+    .unwrap()
+});
+static MAC_MARKERS: LazyLock<MarkerTree> = LazyLock::new(|| {
+    MarkerTree::from_str(
+        "platform_system == 'Darwin' and os_name == 'posix' and sys_platform == 'darwin'",
+    )
+    .unwrap()
+});
+
 #[derive(Clone, Debug, serde::Deserialize)]
 #[serde(try_from = "LockWire")]
 pub struct Lock {
@@ -103,8 +123,7 @@ impl Lock {
                 .cloned()
                 .unwrap_or_default();
             let mut package = Package::from_annotated_dist(dist, fork_markers, root)?;
-
-            Self::remove_unreachable_wheels(&requires_python, &mut package);
+            Self::remove_unreachable_wheels(graph, &requires_python, node_index, &mut package);
 
             // Add all dependencies
             for edge in graph.petgraph.edges(node_index) {
@@ -207,12 +226,60 @@ impl Lock {
     }
 
     /// Remove wheels that can't be selected for installation due to environment markers.
-    fn remove_unreachable_wheels(requires_python: &RequiresPython, package: &mut Package) {
-        // Remove wheels that don't match `requires-python` and can't be selected for
-        // installation.
-        package
+    ///
+    /// For example, a package included under `sys_platform == 'win32'` does not need Linux
+    /// wheels.
+    fn remove_unreachable_wheels(
+        graph: &ResolutionGraph,
+        requires_python: &RequiresPython,
+        node_index: NodeIndex,
+        locked_dist: &mut Package,
+    ) {
+        // Remove wheels that don't match `requires-python` and can't be selected for installation.
+        locked_dist
             .wheels
             .retain(|wheel| requires_python.matches_wheel_tag(&wheel.filename));
+
+        // Filter by platform tags.
+
+        // See https://github.com/pypi/warehouse/blob/ccff64920db7965078cf1fdb50f028e640328887/warehouse/forklift/legacy.py#L100-L169
+        // for a list of relevant platforms.
+        let linux_tags = [
+            "manylinux1_",
+            "manylinux2010_",
+            "manylinux2014_",
+            "musllinux_",
+            "manylinux_",
+        ];
+        let windows_tags = ["win32", "win_arm64", "win_amd64", "win_ia64"];
+
+        locked_dist.wheels.retain(|wheel| {
+            // Naively, we'd check whether `platform_system == 'Linux'` is disjoint, or
+            // `os_name == 'posix'` is disjoint, or `sys_platform == 'linux'` is disjoint (each on its
+            // own sufficient to exclude linux wheels), but due to
+            // `(A ∩ (B ∩ C) = ∅) => ((A ∩ B = ∅) or (A ∩ C = ∅))`
+            // a single disjointness check with the intersection is sufficient, so we have one
+            // constant per platform.
+            let reachability_markers = &graph.reachability[&node_index];
+            let platform_tags = &wheel.filename.platform_tag;
+            if platform_tags.iter().all(|tag| {
+                linux_tags.into_iter().any(|linux_tag| {
+                    // These two linux tags are allowed by warehouse.
+                    tag.starts_with(linux_tag) || tag == "linux_armv6l" || tag == "linux_armv7l"
+                })
+            }) {
+                !reachability_markers.is_disjoint(&LINUX_MARKERS)
+            } else if platform_tags
+                .iter()
+                .all(|tag| windows_tags.contains(&&**tag))
+            {
+                !graph.reachability[&node_index].is_disjoint(&WINDOWS_MARKERS)
+            } else if platform_tags.iter().all(|tag| tag.starts_with("macosx_")) {
+                !graph.reachability[&node_index].is_disjoint(&MAC_MARKERS)
+            } else {
+                true
+            }
+        });
     }
 
     /// Initialize a [`Lock`] from a list of [`Package`] entries.

diff --git a/crates/uv/tests/export.rs b/crates/uv/tests/export.rs
@@ -292,36 +292,10 @@ fn dependency_multiple_markers() -> Result<()> {
         --hash=sha256:99b87a485a5820b23b879f04c2305b44b951b502fd64be915879d77a7e8fc6f1
     cffi==1.16.0 ; (implementation_name != 'pypy' and os_name == 'nt' and sys_platform == 'win32') or (python_full_version >= '3.12' and implementation_name != 'pypy' and os_name == 'nt') \
         --hash=sha256:bcb3ef43e58665bbda2fb198698fcae6776483e0c4a631aa5647806c25e02cc0 \
-        --hash=sha256:6b3d6606d369fc1da4fd8c357d026317fbb9c9b75d36dc16e90e84c26854b088 \
-        --hash=sha256:ac0f5edd2360eea2f1daa9e26a41db02dd4b0451b48f7c318e217ee092a213e9 \
-        --hash=sha256:7e61e3e4fa664a8588aa25c883eab612a188c725755afff6289454d6362b9673 \
-        --hash=sha256:a72e8961a86d19bdb45851d8f1f08b041ea37d2bd8d4fd19903bc3083d80c896 \
-        --hash=sha256:5b50bf3f55561dac5438f8e70bfcdfd74543fd60df5fa5f62d94e5867deca684 \
-        --hash=sha256:7651c50c8c5ef7bdb41108b7b8c5a83013bfaa8a935590c5d74627c047a583c7 \
-        --hash=sha256:e4108df7fe9b707191e55f33efbcb2d81928e10cea45527879a4749cbe472614 \
-        --hash=sha256:32c68ef735dbe5857c810328cb2481e24722a59a2003018885514d4c09af9743 \
-        --hash=sha256:673739cb539f8cdaa07d92d02efa93c9ccf87e345b9a0b556e3ecc666718468d \
         --hash=sha256:9f90389693731ff1f659e55c7d1640e2ec43ff725cc61b04b2f9c6d8d017df6a \
         --hash=sha256:e6024675e67af929088fda399b2094574609396b1decb609c55fa58b028a32a1 \
-        --hash=sha256:b84834d0cf97e7d27dd5b7f3aca7b6e9263c56308ab9dc8aae9784abb774d404 \
-        --hash=sha256:1b8ebc27c014c59692bb2664c7d13ce7a6e9a629be20e54e7271fa696ff2b417 \
-        --hash=sha256:ee07e47c12890ef248766a6e55bd38ebfb2bb8edd4142d56db91b21ea68b7627 \
-        --hash=sha256:d8a9d3ebe49f084ad71f9269834ceccbf398253c9fac910c4fd7053ff1386936 \
-        --hash=sha256:e70f54f1796669ef691ca07d046cd81a29cb4deb1e5f942003f401c0c4a2695d \
-        --hash=sha256:5bf44d66cdf9e893637896c7faa22298baebcd18d1ddb6d2626a6e39793a1d56 \
-        --hash=sha256:7b78010e7b97fef4bee1e896df8a4bbb6712b7f05b7ef630f9d1da00f6444d2e \
-        --hash=sha256:c6a164aa47843fb1b01e941d385aab7215563bb8816d80ff3a363a9f8448a8dc \
-        --hash=sha256:e09f3ff613345df5e8c3667da1d918f9149bd623cd9070c983c013792a9a62eb \
         --hash=sha256:2c56b361916f390cd758a57f2e16233eb4f64bcbeee88a4881ea90fca14dc6ab \
         --hash=sha256:db8e577c19c0fda0beb7e0d4e09e0ba74b1e4c092e0e40bfa12fe05b6f6d75ba \
-        --hash=sha256:fa3a0128b152627161ce47201262d3140edb5a5c3da88d73a1b790a959126956 \
-        --hash=sha256:68e7c44931cc171c54ccb702482e9fc723192e88d25a0e133edd7aff8fcd1f6e \
-        --hash=sha256:abd808f9c129ba2beda4cfc53bde801e5bcf9d6e0f22f095e45327c038bfe68e \
-        --hash=sha256:88e2b3c14bdb32e440be531ade29d3c50a1a59cd4e51b1dd8b0865c54ea5d2e2 \
-        --hash=sha256:fcc8eb6d5902bb1cf6dc4f187ee3ea80a1eba0a89aba40a5cb20a5087d961357 \
-        --hash=sha256:b7be2d771cdba2942e13215c4e340bfd76398e9227ad10402a8767ab1865d2e6 \
-        --hash=sha256:e715596e683d2ce000574bae5d07bd522c781a822866c20495e52520564f0969 \
-        --hash=sha256:2d92b25dbf6cae33f65005baf472d2c245c050b1ce709cc4588cdcdd5495b520 \
         --hash=sha256:b2ca4e77f9f47c55c194982e10f058db063937845bb2b7a86c84a6cfe0aefa8b \
         --hash=sha256:68678abf380b42ce21a5f2abde8efee05c114c2fdb2e9eef2efdb0257fba1235
     exceptiongroup==1.2.0 ; python_full_version < '3.11' and sys_platform == 'win32' \
@@ -391,14 +365,6 @@ fn dependency_conflicting_markers() -> Result<()> {
         --hash=sha256:99b87a485a5820b23b879f04c2305b44b951b502fd64be915879d77a7e8fc6f1
     cffi==1.16.0 ; (implementation_name != 'pypy' and os_name == 'nt' and sys_platform == 'darwin') or (os_name == 'nt' and sys_platform == 'win32') \
         --hash=sha256:bcb3ef43e58665bbda2fb198698fcae6776483e0c4a631aa5647806c25e02cc0 \
-        --hash=sha256:fa3a0128b152627161ce47201262d3140edb5a5c3da88d73a1b790a959126956 \
-        --hash=sha256:68e7c44931cc171c54ccb702482e9fc723192e88d25a0e133edd7aff8fcd1f6e \
-        --hash=sha256:abd808f9c129ba2beda4cfc53bde801e5bcf9d6e0f22f095e45327c038bfe68e \
-        --hash=sha256:88e2b3c14bdb32e440be531ade29d3c50a1a59cd4e51b1dd8b0865c54ea5d2e2 \
-        --hash=sha256:fcc8eb6d5902bb1cf6dc4f187ee3ea80a1eba0a89aba40a5cb20a5087d961357 \
-        --hash=sha256:b7be2d771cdba2942e13215c4e340bfd76398e9227ad10402a8767ab1865d2e6 \
-        --hash=sha256:e715596e683d2ce000574bae5d07bd522c781a822866c20495e52520564f0969 \
-        --hash=sha256:2d92b25dbf6cae33f65005baf472d2c245c050b1ce709cc4588cdcdd5495b520 \
         --hash=sha256:b2ca4e77f9f47c55c194982e10f058db063937845bb2b7a86c84a6cfe0aefa8b \
         --hash=sha256:68678abf380b42ce21a5f2abde8efee05c114c2fdb2e9eef2efdb0257fba1235
     idna==3.6 ; sys_platform == 'darwin' or sys_platform == 'win32' \

diff --git a/crates/uv/tests/lock.rs b/crates/uv/tests/lock.rs
@@ -10525,7 +10525,6 @@ fn lock_constrained_environment() -> Result<()> {
             { url = "https://files.pythonhosted.org/packages/b6/c6/1d174efa9ff02b22d0124c73fc5f4d4fb006d0d9a081aadc354d05754a13/black-24.3.0-cp312-cp312-macosx_10_9_x86_64.whl", hash = "sha256:2818cf72dfd5d289e48f37ccfa08b460bf469e67fb7c4abb07edc2e9f16fb63f", size = 1600822 },
             { url = "https://files.pythonhosted.org/packages/d9/ed/704731afffe460b8ff0672623b40fce9fe569f2ee617c15857e4d4440a3a/black-24.3.0-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:4acf672def7eb1725f41f38bf6bf425c8237248bb0804faa3965c036f7672d11", size = 1429987 },
             { url = "https://files.pythonhosted.org/packages/a8/05/8dd038e30caadab7120176d4bc109b7ca2f4457f12eef746b0560a583458/black-24.3.0-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:c7ed6668cbbfcd231fa0dc1b137d3e40c04c7f786e626b405c62bcd5db5857e4", size = 1755319 },
-            { url = "https://files.pythonhosted.org/packages/71/9d/e5fa1ff4ef1940be15a64883c0bb8d2fcf626efec996eab4ae5a8c691d2c/black-24.3.0-cp312-cp312-win_amd64.whl", hash = "sha256:56f52cfbd3dabe2798d76dbdd299faa046a901041faf2cf33288bc4e6dae57b5", size = 1385180 },
             { url = "https://files.pythonhosted.org/packages/4d/ea/31770a7e49f3eedfd8cd7b35e78b3a3aaad860400f8673994bc988318135/black-24.3.0-py3-none-any.whl", hash = "sha256:41622020d7120e01d377f74249e677039d20e6344ff5851de8a10f11f513bf93", size = 201493 },
         ]