From b4ea48955ba798ecd8bdcb9fcf2c6d3a04ca157c Mon Sep 17 00:00:00 2001
From: Charlie Marsh <charlie.r.marsh@gmail.com>
Date: Fri, 16 Feb 2024 16:30:54 -0500
Subject: [PATCH] Use comparable representation for `PackageId` (#1543)

## Summary

By using the display representation of `Version` to form a `PackageId`,
we run the risk (as seen in the linked issue) of thinking that versions
like `2021.1` and `2021.1.0` are not equivalent.

Closes https://github.com/astral-sh/uv/issues/1536
---
 crates/distribution-types/src/id.rs     | 25 ++++++++++++++++++----
 crates/distribution-types/src/traits.rs | 10 ++++-----
 crates/uv/tests/pip_compile.rs          | 28 +++++++++++++++++++++++++
 3 files changed, 53 insertions(+), 10 deletions(-)

diff --git a/crates/distribution-types/src/id.rs b/crates/distribution-types/src/id.rs
index f61b8378e..1fb83ab7e 100644
--- a/crates/distribution-types/src/id.rs
+++ b/crates/distribution-types/src/id.rs
@@ -1,18 +1,35 @@
 use std::fmt::{Display, Formatter};
 
+use url::Url;
+
+use pep440_rs::Version;
+use uv_normalize::PackageName;
+
 /// A unique identifier for a package (e.g., `black==23.10.0`).
 #[derive(Debug, Clone, PartialEq, Eq, Hash, PartialOrd, Ord)]
-pub struct PackageId(String);
+pub enum PackageId {
+    NameVersion(PackageName, Version),
+    Url(String),
+}
 
 impl PackageId {
-    pub fn new(id: impl Into<String>) -> Self {
-        Self(id.into())
+    /// Create a new [`PackageId`] from a package name and version.
+    pub fn from_registry(name: PackageName, version: Version) -> Self {
+        Self::NameVersion(name, version)
+    }
+
+    /// Create a new [`PackageId`] from a URL.
+    pub fn from_url(url: &Url) -> Self {
+        Self::Url(cache_key::digest(&cache_key::CanonicalUrl::new(url)))
     }
 }
 
 impl Display for PackageId {
     fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
-        f.write_str(&self.0)
+        match self {
+            PackageId::NameVersion(name, version) => write!(f, "{name}-{version}"),
+            PackageId::Url(url) => write!(f, "{url}"),
+        }
     }
 }
 
diff --git a/crates/distribution-types/src/traits.rs b/crates/distribution-types/src/traits.rs
index 9f963a1dd..e18d58d6a 100644
--- a/crates/distribution-types/src/traits.rs
+++ b/crates/distribution-types/src/traits.rs
@@ -31,14 +31,12 @@ pub trait DistributionMetadata: Name {
     /// registry-based distributions (e.g., different wheels for the same package and version)
     /// will return the same package ID, but different distribution IDs.
     fn package_id(&self) -> PackageId {
-        PackageId::new(match self.version_or_url() {
+        match self.version_or_url() {
             VersionOrUrl::Version(version) => {
-                // https://packaging.python.org/en/latest/specifications/recording-installed-packages/#the-dist-info-directory
-                // `version` is normalized by its `ToString` impl
-                format!("{}-{}", self.name(), version)
+                PackageId::from_registry(self.name().clone(), version.clone())
             }
-            VersionOrUrl::Url(url) => cache_key::digest(&cache_key::CanonicalUrl::new(url)),
-        })
+            VersionOrUrl::Url(url) => PackageId::from_url(url),
+        }
     }
 }
 
diff --git a/crates/uv/tests/pip_compile.rs b/crates/uv/tests/pip_compile.rs
index 5fd10af4f..cbd6585d1 100644
--- a/crates/uv/tests/pip_compile.rs
+++ b/crates/uv/tests/pip_compile.rs
@@ -3376,3 +3376,31 @@ fn compile_none_extra() -> Result<()> {
 
     Ok(())
 }
+
+/// Resolve a package (`pytz`) with a preference that omits a trailing zero.
+///
+/// See: <https://github.com/astral-sh/uv/issues/1536>
+#[test]
+fn compile_types_pytz() -> Result<()> {
+    let context = TestContext::new("3.12");
+    let requirements_in = context.temp_dir.child("requirements.in");
+    requirements_in.write_str("types-pytz")?;
+
+    let requirements_txt = context.temp_dir.child("requirements.txt");
+    requirements_txt.write_str("types-pytz==2021.1")?;
+
+    uv_snapshot!(context
+        .compile()
+        .arg("requirements.in")
+        .arg("-o")
+        .arg("requirements.txt"), @r###"
+    success: true
+    exit_code: 0
+    ----- stdout -----
+
+    ----- stderr -----
+    Resolved 1 package in [TIME]
+    "###);
+
+    Ok(())
+}