Use iterators rather than eagerly collecting tags

astral-sh · Jan 18, 2025 · f3a4192 · f3a4192
1 parent 4f80747
commit f3a4192
Show file tree

Hide file tree

Showing 2 changed files with 23 additions and 23 deletions.
diff --git a/crates/uv-distribution-types/src/prioritized_distribution.rs b/crates/uv-distribution-types/src/prioritized_distribution.rs
@@ -1,4 +1,3 @@
-use std::collections::BTreeSet;
 use std::fmt::{Display, Formatter};
 
 use arcstr::ArcStr;
@@ -522,38 +521,41 @@ impl PrioritizedDist {
         self.0.best_wheel_index.map(|i| &self.0.wheels[i])
     }
 
-    /// Returns the set of all Python tags for the distribution.
-    pub fn python_tags(&self) -> BTreeSet<LanguageTag> {
+    /// Returns an iterator over all Python tags for the distribution.
+    pub fn python_tags(&self) -> impl Iterator<Item = LanguageTag> + '_ {
         self.0
             .wheels
             .iter()
             .flat_map(|(wheel, _)| wheel.filename.python_tags().iter().copied())
-            .collect()
     }
 
-    /// Returns the set of all ABI tags for the distribution.
-    pub fn abi_tags(&self) -> BTreeSet<AbiTag> {
+    /// Returns an iterator over all ABI tags for the distribution.
+    pub fn abi_tags(&self) -> impl Iterator<Item = AbiTag> + '_ {
         self.0
             .wheels
             .iter()
             .flat_map(|(wheel, _)| wheel.filename.abi_tags().iter().copied())
-            .collect()
     }
 
     /// Returns the set of platform tags for the distribution that are ABI-compatible with the given
     /// tags.
-    pub fn platform_tags<'a>(&'a self, tags: &'a Tags) -> BTreeSet<&'a PlatformTag> {
-        let mut candidates = BTreeSet::new();
-        for (wheel, _) in &self.0.wheels {
-            for wheel_py in wheel.filename.python_tags() {
-                for wheel_abi in wheel.filename.abi_tags() {
-                    if tags.is_compatible_abi(*wheel_py, *wheel_abi) {
-                        candidates.extend(wheel.filename.platform_tags().iter());
-                    }
-                }
+    pub fn platform_tags<'a>(
+        &'a self,
+        tags: &'a Tags,
+    ) -> impl Iterator<Item = &'a PlatformTag> + 'a {
+        self.0.wheels.iter().flat_map(move |(wheel, _)| {
+            if wheel.filename.python_tags().iter().any(|wheel_py| {
+                wheel
+                    .filename
+                    .abi_tags()
+                    .iter()
+                    .any(|wheel_abi| tags.is_compatible_abi(*wheel_py, *wheel_abi))
+            }) {
+                wheel.filename.platform_tags().iter()
+            } else {
+                [].iter()
             }
-        }
-        candidates
+        })
     }
 }
 

diff --git a/crates/uv-resolver/src/pubgrub/report.rs b/crates/uv-resolver/src/pubgrub/report.rs
@@ -745,7 +745,7 @@ impl PubGrubReportFormatter<'_> {
             IncompatibleTag::Invalid => None,
             IncompatibleTag::Python => {
                 let best = tags.and_then(Tags::python_tag);
-                let tags = prioritized.python_tags();
+                let tags = prioritized.python_tags().collect::<BTreeSet<_>>();
                 if tags.is_empty() {
                     None
                 } else {
@@ -761,7 +761,6 @@ impl PubGrubReportFormatter<'_> {
                 let best = tags.and_then(Tags::abi_tag);
                 let tags = prioritized
                     .abi_tags()
-                    .into_iter()
                     // Ignore `none`, which is universally compatible.
                     //
                     // As an example, `none` can appear here if we're solving for Python 3.13, and
@@ -796,9 +795,8 @@ impl PubGrubReportFormatter<'_> {
                 // we only show platforms for ABI-compatible wheels.
                 let tags = prioritized
                     .platform_tags(self.tags?)
-                    .into_iter()
                     .cloned()
-                    .collect::<Vec<_>>();
+                    .collect::<BTreeSet<_>>();
                 if tags.is_empty() {
                     None
                 } else {
@@ -1132,7 +1130,7 @@ pub(crate) enum PubGrubHint {
         // excluded from `PartialEq` and `Hash`
         version: Version,
         // excluded from `PartialEq` and `Hash`
-        tags: Vec<PlatformTag>,
+        tags: BTreeSet<PlatformTag>,
     },
 }