Refactor storage of distribution metadata needed in resolver (#1291)

Follows #1290 and https://github.com/astral-sh/puffin/pull/912 with some
minor clean-up.

clippy.toml

@@ -0,0 +1,4 @@
+doc-valid-idents = [
+  "PyPI",
+  ".." # Include the defaults
+] 

crates/distribution-types/src/prioritized_distribution.rs

@@ -4,12 +4,16 @@ use pypi_types::{Hashes, Yanked};
 
 use crate::Dist;
 
-/// Attach its requires-python to a [`Dist`], since downstream needs this information to filter
+/// A [`Dist`] and metadata about it required for downstream filtering.
-/// [`PrioritizedDistribution`].
 #[derive(Debug, Clone)]
-pub struct DistRequiresPython {
+pub struct DistMetadata {
     pub dist: Dist,
+
+    /// The version of Python required by the distribution
     pub requires_python: Option<VersionSpecifiers>,
+
+    /// Is the distribution file yanked?
+    pub yanked: Yanked,
 }
 
 // Boxed because `Dist` is large.
@@ -19,13 +23,11 @@ pub struct PrioritizedDistribution(Box<PrioritizedDistributionInner>);
 #[derive(Debug, Default, Clone)]
 struct PrioritizedDistributionInner {
     /// An arbitrary source distribution for the package version.
-    source: Option<DistRequiresPython>,
+    source: Option<DistMetadata>,
     /// The highest-priority, platform-compatible wheel for the package version.
-    compatible_wheel: Option<(DistRequiresPython, TagPriority)>,
+    compatible_wheel: Option<(DistMetadata, TagPriority)>,
     /// An arbitrary, platform-incompatible wheel for the package version.
-    incompatible_wheel: Option<DistRequiresPython>,
+    incompatible_wheel: Option<DistMetadata>,
-    /// Is the distribution yanked
-    yanked: Yanked,
     /// The hashes for each distribution.
     hashes: Vec<Hashes>,
 }
@@ -35,34 +37,34 @@ impl PrioritizedDistribution {
     pub fn from_built(
         dist: Dist,
         requires_python: Option<VersionSpecifiers>,
+        yanked: Yanked,
         hash: Option<Hashes>,
         priority: Option<TagPriority>,
-        yanked: Yanked,
     ) -> Self {
         if let Some(priority) = priority {
             Self(Box::new(PrioritizedDistributionInner {
                 source: None,
                 compatible_wheel: Some((
-                    DistRequiresPython {
+                    DistMetadata {
                         dist,
                         requires_python,
+                        yanked,
                     },
                     priority,
                 )),
                 incompatible_wheel: None,
                 hashes: hash.map(|hash| vec![hash]).unwrap_or_default(),
-                yanked,
             }))
         } else {
             Self(Box::new(PrioritizedDistributionInner {
                 source: None,
                 compatible_wheel: None,
-                incompatible_wheel: Some(DistRequiresPython {
+                incompatible_wheel: Some(DistMetadata {
                     dist,
                     requires_python,
+                    yanked,
                 }),
                 hashes: hash.map(|hash| vec![hash]).unwrap_or_default(),
-                yanked,
             }))
         }
     }
@@ -71,18 +73,18 @@ impl PrioritizedDistribution {
     pub fn from_source(
         dist: Dist,
         requires_python: Option<VersionSpecifiers>,
-        hash: Option<Hashes>,
         yanked: Yanked,
+        hash: Option<Hashes>,
     ) -> Self {
         Self(Box::new(PrioritizedDistributionInner {
-            source: Some(DistRequiresPython {
+            source: Some(DistMetadata {
                 dist,
                 requires_python,
+                yanked,
             }),
             compatible_wheel: None,
             incompatible_wheel: None,
             hashes: hash.map(|hash| vec![hash]).unwrap_or_default(),
-            yanked,
         }))
     }
 
@@ -91,39 +93,38 @@ impl PrioritizedDistribution {
         &mut self,
         dist: Dist,
         requires_python: Option<VersionSpecifiers>,
+        yanked: Yanked,
         hash: Option<Hashes>,
         priority: Option<TagPriority>,
-        yanked: Yanked,
     ) {
-        if yanked.is_yanked() {
-            self.0.yanked = yanked;
-        }
-
         // Prefer the highest-priority, platform-compatible wheel.
         if let Some(priority) = priority {
             if let Some((.., existing_priority)) = &self.0.compatible_wheel {
                 if priority > *existing_priority {
                     self.0.compatible_wheel = Some((
-                        DistRequiresPython {
+                        DistMetadata {
                             dist,
                             requires_python,
+                            yanked,
                         },
                         priority,
                     ));
                 }
             } else {
                 self.0.compatible_wheel = Some((
-                    DistRequiresPython {
+                    DistMetadata {
                         dist,
                         requires_python,
+                        yanked,
                     },
                     priority,
                 ));
             }
         } else if self.0.incompatible_wheel.is_none() {
-            self.0.incompatible_wheel = Some(DistRequiresPython {
+            self.0.incompatible_wheel = Some(DistMetadata {
                 dist,
                 requires_python,
+                yanked,
             });
         }
 
@@ -137,17 +138,14 @@ impl PrioritizedDistribution {
         &mut self,
         dist: Dist,
         requires_python: Option<VersionSpecifiers>,
-        hash: Option<Hashes>,
         yanked: Yanked,
+        hash: Option<Hashes>,
     ) {
-        if yanked.is_yanked() {
-            self.0.yanked = yanked;
-        }
-
         if self.0.source.is_none() {
-            self.0.source = Some(DistRequiresPython {
+            self.0.source = Some(DistMetadata {
                 dist,
                 requires_python,
+                yanked,
             });
         }
 
@@ -164,35 +162,29 @@ impl PrioritizedDistribution {
             &self.0.incompatible_wheel,
         ) {
             // Prefer the highest-priority, platform-compatible wheel.
-            (Some((wheel, tag_priority)), _, _) => Some(ResolvableDist::CompatibleWheel(
+            (Some((wheel, tag_priority)), _, _) => {
-                wheel,
+                Some(ResolvableDist::CompatibleWheel(wheel, *tag_priority))
-                &self.0.yanked,
+            }
-                *tag_priority,
-            )),
             // If we have a compatible source distribution and an incompatible wheel, return the
             // wheel. We assume that all distributions have the same metadata for a given package
             // version. If a compatible source distribution exists, we assume we can build it, but
             // using the wheel is faster.
-            (_, Some(source_dist), Some(wheel)) => Some(ResolvableDist::IncompatibleWheel {
+            (_, Some(source_dist), Some(wheel)) => {
-                source_dist,
+                Some(ResolvableDist::IncompatibleWheel { source_dist, wheel })
-                yanked: &self.0.yanked,
-                wheel,
-            }),
-            // Otherwise, if we have a source distribution, return it.
-            (_, Some(source_dist), _) => {
-                Some(ResolvableDist::SourceDist(source_dist, &self.0.yanked))
             }
+            // Otherwise, if we have a source distribution, return it.
+            (_, Some(source_dist), _) => Some(ResolvableDist::SourceDist(source_dist)),
             _ => None,
         }
     }
 
     /// Return the source distribution, if any.
-    pub fn source(&self) -> Option<&DistRequiresPython> {
+    pub fn source(&self) -> Option<&DistMetadata> {
         self.0.source.as_ref()
     }
 
     /// Return the compatible built distribution, if any.
-    pub fn compatible_wheel(&self) -> Option<&(DistRequiresPython, TagPriority)> {
+    pub fn compatible_wheel(&self) -> Option<&(DistMetadata, TagPriority)> {
         self.0.compatible_wheel.as_ref()
     }
 
@@ -210,57 +202,56 @@ impl PrioritizedDistribution {
     }
 }
 
+/// A collection of distributions ([`Dist`]) that can be used for resolution and installation.
 #[derive(Debug, Clone)]
 pub enum ResolvableDist<'a> {
     /// The distribution should be resolved and installed using a source distribution.
-    SourceDist(&'a DistRequiresPython, &'a Yanked),
+    SourceDist(&'a DistMetadata),
     /// The distribution should be resolved and installed using a wheel distribution.
-    CompatibleWheel(&'a DistRequiresPython, &'a Yanked, TagPriority),
+    CompatibleWheel(&'a DistMetadata, TagPriority),
     /// The distribution should be resolved using an incompatible wheel distribution, but
     /// installed using a source distribution.
     IncompatibleWheel {
-        source_dist: &'a DistRequiresPython,
+        source_dist: &'a DistMetadata,
-        yanked: &'a Yanked,
+        wheel: &'a DistMetadata,
-        wheel: &'a DistRequiresPython,
     },
 }
 
 impl<'a> ResolvableDist<'a> {
-    /// Return the [`DistRequiresPython`] to use during resolution.
+    /// Return the [`DistMetadata`] to use during resolution.
-    pub fn for_resolution(&self) -> &DistRequiresPython {
+    pub fn for_resolution(&self) -> &DistMetadata {
         match *self {
-            ResolvableDist::SourceDist(sdist, _) => sdist,
+            ResolvableDist::SourceDist(sdist) => sdist,
-            ResolvableDist::CompatibleWheel(wheel, _, _) => wheel,
+            ResolvableDist::CompatibleWheel(wheel, _) => wheel,
             ResolvableDist::IncompatibleWheel {
                 source_dist: _,
-                yanked: _,
                 wheel,
             } => wheel,
         }
     }
 
-    /// Return the [`DistRequiresPython`] to use during installation.
+    /// Return the [`DistMetadata`] to use during installation.
-    pub fn for_installation(&self) -> &DistRequiresPython {
+    pub fn for_installation(&self) -> &DistMetadata {
         match *self {
-            ResolvableDist::SourceDist(sdist, _) => sdist,
+            ResolvableDist::SourceDist(sdist) => sdist,
-            ResolvableDist::CompatibleWheel(wheel, _, _) => wheel,
+            ResolvableDist::CompatibleWheel(wheel, _) => wheel,
             ResolvableDist::IncompatibleWheel {
                 source_dist,
-                yanked: _,
                 wheel: _,
             } => source_dist,
         }
     }
 
+    /// Return the [`Yanked`] status of the distribution.
+    ///
+    /// It is possible for files to have a different yank status per PEP 592 but in the official
+    /// PyPI warehouse this cannot happen.
+    ///
+    /// Here, we will treat the distribution is yanked if the file we will install with
+    /// is yanked.
+    ///
+    /// PEP 592: <https://peps.python.org/pep-0592/#warehouse-pypi-implementation-notes>
     pub fn yanked(&self) -> &Yanked {
-        match *self {
+        &self.for_installation().yanked
-            ResolvableDist::SourceDist(_, yanked) => yanked,
-            ResolvableDist::CompatibleWheel(_, yanked, _) => yanked,
-            ResolvableDist::IncompatibleWheel {
-                source_dist: _,
-                yanked,
-                wheel: _,
-            } => yanked,
-        }
     }
 }

crates/puffin-client/src/flat_index.rs

@@ -300,15 +300,15 @@ impl FlatIndex {
                     Entry::Occupied(mut entry) => {
                         entry
                             .get_mut()
-                            .insert_built(dist, None, None, priority, Yanked::default());
+                            .insert_built(dist, None, Yanked::default(), None, priority);
                     }
                     Entry::Vacant(entry) => {
                         entry.insert(PrioritizedDistribution::from_built(
                             dist,
                             None,
+                            Yanked::default(),
                             None,
                             priority,
-                            Yanked::default(),
                         ));
                     }
                 }
@@ -323,14 +323,14 @@ impl FlatIndex {
                     Entry::Occupied(mut entry) => {
                         entry
                             .get_mut()
-                            .insert_source(dist, None, None, Yanked::default());
+                            .insert_source(dist, None, Yanked::default(), None);
                     }
                     Entry::Vacant(entry) => {
                         entry.insert(PrioritizedDistribution::from_source(
                             dist,
                             None,
-                            None,
                             Yanked::default(),
+                            None,
                         ));
                     }
                 }

crates/puffin-resolver/src/candidate_selector.rs

@@ -3,7 +3,7 @@ use pypi_types::Yanked;
 use rustc_hash::FxHashMap;
 
 use distribution_types::{Dist, DistributionMetadata, Name};
-use distribution_types::{DistRequiresPython, ResolvableDist};
+use distribution_types::{DistMetadata, ResolvableDist};
 use pep440_rs::{Version, VersionSpecifiers};
 use pep508_rs::{Requirement, VersionOrUrl};
 use puffin_normalize::PackageName;
@@ -248,12 +248,12 @@ impl<'a> Candidate<'a> {
     }
 
     /// Return the [`DistFile`] to use when resolving the package.
-    pub(crate) fn resolution_dist(&self) -> &DistRequiresPython {
+    pub(crate) fn resolution_dist(&self) -> &DistMetadata {
         self.dist.for_resolution()
     }
 
     /// Return the [`DistFile`] to use when installing the package.
-    pub(crate) fn installation_dist(&self) -> &DistRequiresPython {
+    pub(crate) fn installation_dist(&self) -> &DistMetadata {
         self.dist.for_installation()
     }
 
@@ -295,8 +295,9 @@ impl<'a> Candidate<'a> {
         None
     }
 
+    /// If the distribution that would be installed is yanked.
     pub(crate) fn yanked(&self) -> &Yanked {
-        return self.dist.yanked();
+        self.dist.yanked()
     }
 }
 

crates/puffin-resolver/src/version_map.rs

@@ -78,10 +78,6 @@ impl VersionMap {
                     }
                 }
 
-                // It is possible for files to have a different yank status per PEP 592 but in the official
-                // PyPI warehouse this cannot happen.
-                // If any file is yanked, the version will be marked as yanked.
-                // <https://peps.python.org/pep-0592/#warehouse-pypi-implementation-notes>
                 let yanked = if let Some(ref yanked) = file.yanked {
                     yanked.clone()
                 } else {
@@ -115,9 +111,9 @@ impl VersionMap {
                         priority_dist.insert_built(
                             dist,
                             requires_python,
+                            yanked,
                             Some(hash),
                             priority,
-                            yanked,
                         );
                     }
                     DistFilename::SourceDistFilename(filename) => {
@@ -126,7 +122,7 @@ impl VersionMap {
                             file,
                             index.clone(),
                         );
-                        priority_dist.insert_source(dist, requires_python, Some(hash), yanked);
+                        priority_dist.insert_source(dist, requires_python, yanked, Some(hash));
                     }
                 }
             }