Add support for PyPy wheels (#1028)

## Summary

This PR adds support for PyPy wheels by changing the compatible tags
based on the implementation name and version of the current interpreter.

For now, we only support CPython and PyPy, and explicitly error out when
given other interpreters. (Is this right? Should we just fallback to
CPython tags...? Or skip the ABI-specific tags for unknown
interpreters?)

The logic is based on
4d85340613/src/packaging/tags.py (L247).
Note, however, that `packaging` uses the `EXT_SUFFIX` variable from
`sysconfig`... Instead, I looked at the way that PyPy formats the tags,
and recreated them based on the Python and implementation version. For
example, PyPy wheels look like
`cchardet-2.1.7-pp37-pypy37_pp73-win_amd64.whl` -- so that's `pp37` for
PyPy with Python version 3.7, and then `pypy37_pp73` for PyPy with
Python version 3.7 and PyPy version 7.3.

Closes https://github.com/astral-sh/puffin/issues/1013.

## Test Plan

I tested this manually, but I couldn't find macOS universal PyPy
wheels... So instead I added `cchardet` to a `requirements.in`, ran
`cargo run pip sync requirements.in --index-url
https://pypy.kmtea.eu/simple --verbose`, and added logging to verify
that the platform tags matched (even if the architecture didn't).

Cargo.lock

@@ -2156,9 +2156,9 @@ dependencies = [
 name = "platform-tags"
 version = "0.0.1"
 dependencies = [
- "anyhow",
  "platform-host",
  "rustc-hash",
+ "thiserror",
 ]
 
 [[package]]

crates/gourgeist/src/bare.rs

@@ -114,14 +114,14 @@ pub fn create_bare_venv(location: &Utf8Path, interpreter: &Interpreter) -> io::R
         symlink(&base_python, &venv_python)?;
         symlink(
             "python",
-            bin_dir.join(format!("python{}", interpreter.simple_version().0)),
+            bin_dir.join(format!("python{}", interpreter.python_major())),
         )?;
         symlink(
             "python",
             bin_dir.join(format!(
                 "python{}.{}",
-                interpreter.simple_version().0,
+                interpreter.python_major(),
-                interpreter.simple_version().1
+                interpreter.python_minor(),
             )),
         )?;
     }
@@ -134,8 +134,8 @@ pub fn create_bare_venv(location: &Utf8Path, interpreter: &Interpreter) -> io::R
                 "{{ RELATIVE_SITE_PACKAGES }}",
                 &format!(
                     "../lib/python{}.{}/site-packages",
-                    interpreter.simple_version().0,
+                    interpreter.python_major(),
-                    interpreter.simple_version().1
+                    interpreter.python_minor(),
                 ),
             );
         fs::write(bin_dir.join(name), activator)?;
@@ -180,8 +180,8 @@ pub fn create_bare_venv(location: &Utf8Path, interpreter: &Interpreter) -> io::R
         .join("lib")
         .join(format!(
             "python{}.{}",
-            interpreter.simple_version().0,
+            interpreter.python_major(),
-            interpreter.simple_version().1
+            interpreter.python_minor(),
         ))
         .join("site-packages");
     fs::create_dir_all(&site_packages)?;

crates/platform-tags/Cargo.toml

@@ -15,5 +15,5 @@ workspace = true
 [dependencies]
 platform-host = { path = "../platform-host" }
 
-anyhow = { workspace = true }
 rustc-hash = { workspace = true }
+thiserror = { workspace = true }

crates/platform-tags/src/lib.rs

@@ -1,10 +1,22 @@
 use std::num::NonZeroU32;
+use std::str::FromStr;
 
-use anyhow::{Error, Result};
 use rustc_hash::FxHashMap;
 
 use platform_host::{Arch, Os, Platform, PlatformError};
 
+#[derive(Debug, thiserror::Error)]
+pub enum TagsError {
+    #[error(transparent)]
+    PlatformError(#[from] PlatformError),
+    #[error("Unsupported implementation: {0}")]
+    UnsupportedImplementation(String),
+    #[error("Unknown implementation: {0}")]
+    UnknownImplementation(String),
+    #[error("Invalid priority: {0}")]
+    InvalidPriority(usize, #[source] std::num::TryFromIntError),
+}
+
 /// A set of compatible tags for a given Python version and platform.
 ///
 /// Its principle function is to determine whether the tags for a particular
@@ -33,8 +45,15 @@ impl Tags {
         Self { map }
     }
 
-    /// Returns the compatible tags for the given Python version and platform.
+    /// Returns the compatible tags for the given Python implementation (e.g., `cpython`), version,
-    pub fn from_env(platform: &Platform, python_version: (u8, u8)) -> Result<Self, PlatformError> {
+    /// and platform.
+    pub fn from_env(
+        platform: &Platform,
+        python_version: (u8, u8),
+        implementation_name: &str,
+        implementation_version: (u8, u8),
+    ) -> Result<Self, TagsError> {
+        let implementation = Implementation::from_str(implementation_name)?;
         let platform_tags = compatible_tags(platform)?;
 
         let mut tags = Vec::with_capacity(5 * platform_tags.len());
@@ -42,31 +61,27 @@ impl Tags {
         // 1. This exact c api version
         for platform_tag in &platform_tags {
             tags.push((
-                format!("cp{}{}", python_version.0, python_version.1),
+                implementation.language_tag(python_version),
-                format!(
+                implementation.abi_tag(python_version, implementation_version),
-                    "cp{}{}{}",
-                    python_version.0,
-                    python_version.1,
-                    // hacky but that's legacy anyways
-                    if python_version.1 <= 7 { "m" } else { "" }
-                ),
                 platform_tag.clone(),
             ));
             tags.push((
-                format!("cp{}{}", python_version.0, python_version.1),
+                implementation.language_tag(python_version),
                 "none".to_string(),
                 platform_tag.clone(),
             ));
         }
         // 2. abi3 and no abi (e.g. executable binary)
-        // For some reason 3.2 is the minimum python for the cp abi
+        if matches!(implementation, Implementation::CPython) {
-        for minor in 2..=python_version.1 {
+            // For some reason 3.2 is the minimum python for the cp abi
-            for platform_tag in &platform_tags {
+            for minor in 2..=python_version.1 {
-                tags.push((
+                for platform_tag in &platform_tags {
-                    format!("cp{}{}", python_version.0, minor),
+                    tags.push((
-                    "abi3".to_string(),
+                        implementation.language_tag((python_version.0, minor)),
-                    platform_tag.clone(),
+                        "abi3".to_string(),
-                ));
+                        platform_tag.clone(),
+                    ));
+                }
             }
         }
         // 3. no abi (e.g. executable binary)
@@ -174,12 +189,73 @@ impl Tags {
 pub struct TagPriority(NonZeroU32);
 
 impl TryFrom<usize> for TagPriority {
-    type Error = Error;
+    type Error = TagsError;
 
     /// Create a [`TagPriority`] from a `usize`, where higher `usize` values are given higher
     /// priority.
-    fn try_from(priority: usize) -> Result<Self> {
+    fn try_from(priority: usize) -> Result<Self, TagsError> {
-        Ok(Self(NonZeroU32::try_from(1 + u32::try_from(priority)?)?))
+        match u32::try_from(priority).and_then(|priority| NonZeroU32::try_from(1 + priority)) {
+            Ok(priority) => Ok(Self(priority)),
+            Err(err) => Err(TagsError::InvalidPriority(priority, err)),
+        }
+    }
+}
+
+#[derive(Debug, Clone, Copy)]
+pub enum Implementation {
+    CPython,
+    PyPy,
+}
+
+impl Implementation {
+    /// Returns the "language implementation and version tag" for the current implementation and
+    /// Python version (e.g., `cp39` or `pp37`).
+    pub fn language_tag(&self, python_version: (u8, u8)) -> String {
+        match self {
+            // Ex) `cp39`
+            Implementation::CPython => format!("cp{}{}", python_version.0, python_version.1),
+            // Ex) `pp39`
+            Implementation::PyPy => format!("pp{}{}", python_version.0, python_version.1),
+        }
+    }
+
+    pub fn abi_tag(&self, python_version: (u8, u8), implementation_version: (u8, u8)) -> String {
+        match self {
+            // Ex) `cp39`
+            Implementation::CPython => {
+                if python_version.1 <= 7 {
+                    format!("cp{}{}m", python_version.0, python_version.1)
+                } else {
+                    format!("cp{}{}", python_version.0, python_version.1)
+                }
+            }
+            // Ex) `pypy39_pp73`
+            Implementation::PyPy => format!(
+                "pypy{}{}_pp{}{}",
+                python_version.0,
+                python_version.1,
+                implementation_version.0,
+                implementation_version.1
+            ),
+        }
+    }
+}
+
+impl FromStr for Implementation {
+    type Err = TagsError;
+
+    fn from_str(s: &str) -> Result<Self, TagsError> {
+        match s {
+            // Known and supported implementations.
+            "cpython" => Ok(Self::CPython),
+            "pypy" => Ok(Self::PyPy),
+            // Known but unsupported implementations.
+            "python" => Err(TagsError::UnsupportedImplementation(s.to_string())),
+            "ironpython" => Err(TagsError::UnsupportedImplementation(s.to_string())),
+            "jython" => Err(TagsError::UnsupportedImplementation(s.to_string())),
+            // Unknown implementations.
+            _ => Err(TagsError::UnknownImplementation(s.to_string())),
+        }
     }
 }
 

crates/puffin-build/src/lib.rs

@@ -460,7 +460,7 @@ impl SourceBuild {
         let span = info_span!(
             "run_python_script",
             script="prepare_metadata_for_build_wheel",
-            python_version = %self.venv.interpreter().version()
+            python_version = %self.venv.interpreter().python_version()
         );
         let output = run_python_script(&self.venv, &script, &self.source_tree)
             .instrument(span)
@@ -525,7 +525,7 @@ impl SourceBuild {
             let span = info_span!(
                 "run_python_script",
                 script="setup.py bdist_wheel",
-                python_version = %self.venv.interpreter().version()
+                python_version = %self.venv.interpreter().python_version()
             );
             let output = Command::new(&python_interpreter)
                 .args(["setup.py", "bdist_wheel"])
@@ -585,7 +585,7 @@ impl SourceBuild {
         let span = info_span!(
             "run_python_script",
             script=format!("build_{}", self.build_kind),
-            python_version = %self.venv.interpreter().version()
+            python_version = %self.venv.interpreter().python_version()
         );
         let output = run_python_script(&self.venv, &script, &self.source_tree)
             .instrument(span)
@@ -663,7 +663,7 @@ async fn create_pep517_build_environment(
     let span = info_span!(
         "run_python_script",
         script=format!("get_requires_for_build_{}", build_kind),
-        python_version = %venv.interpreter().version()
+        python_version = %venv.interpreter().python_version()
     );
     let output = run_python_script(venv, &script, source_tree)
         .instrument(span)

crates/puffin-installer/src/installer.rs

@@ -43,7 +43,7 @@ impl<'a> Installer<'a> {
             wheels.par_iter().try_for_each(|wheel| {
                 let location = install_wheel_rs::InstallLocation::new(
                     self.venv.root(),
-                    self.venv.interpreter().simple_version(),
+                    self.venv.interpreter().python_tuple(),
                 );
 
                 install_wheel_rs::linker::install_wheel(

crates/puffin-installer/src/site_packages.rs

@@ -180,10 +180,10 @@ impl<'a> SitePackages<'a> {
 
             // Verify that the package is compatible with the current Python version.
             if let Some(requires_python) = metadata.requires_python.as_ref() {
-                if !requires_python.contains(self.venv.interpreter().version()) {
+                if !requires_python.contains(self.venv.interpreter().python_version()) {
                     diagnostics.push(Diagnostic::IncompatiblePythonVersion {
                         package: package.clone(),
-                        version: self.venv.interpreter().version().clone(),
+                        version: self.venv.interpreter().python_version().clone(),
                         requires_python: requires_python.clone(),
                     });
                 }

crates/puffin-interpreter/src/interpreter.rs

@@ -9,8 +9,8 @@ use tracing::{debug, warn};
 use cache_key::digest;
 use pep440_rs::Version;
 use pep508_rs::MarkerEnvironment;
-use platform_host::{Platform, PlatformError};
+use platform_host::Platform;
-use platform_tags::Tags;
+use platform_tags::{Tags, TagsError};
 use puffin_cache::{Cache, CacheBucket, CachedByTimestamp};
 use puffin_fs::write_atomic_sync;
 
@@ -156,34 +156,60 @@ impl Interpreter {
     }
 
     /// Returns the [`Tags`] for this Python executable.
-    pub fn tags(&self) -> Result<&Tags, PlatformError> {
+    pub fn tags(&self) -> Result<&Tags, TagsError> {
-        self.tags
+        self.tags.get_or_try_init(|| {
-            .get_or_try_init(|| Tags::from_env(self.platform(), self.simple_version()))
+            Tags::from_env(
+                self.platform(),
+                self.python_tuple(),
+                self.implementation_name(),
+                self.implementation_tuple(),
+            )
+        })
     }
 
     /// Returns the Python version.
     #[inline]
-    pub const fn version(&self) -> &Version {
+    pub const fn python_version(&self) -> &Version {
         &self.markers.python_full_version.version
     }
 
     /// Return the major version of this Python version.
-    pub fn major(&self) -> u8 {
+    pub fn python_major(&self) -> u8 {
-        let major = self.version().release()[0];
+        let major = self.markers.python_full_version.version.release()[0];
         u8::try_from(major).expect("invalid major version")
     }
 
     /// Return the minor version of this Python version.
-    pub fn minor(&self) -> u8 {
+    pub fn python_minor(&self) -> u8 {
-        let minor = self.version().release()[1];
+        let minor = self.markers.python_full_version.version.release()[1];
         u8::try_from(minor).expect("invalid minor version")
     }
 
     /// Returns the Python version as a simple tuple.
-    pub fn simple_version(&self) -> (u8, u8) {
+    pub fn python_tuple(&self) -> (u8, u8) {
-        (self.major(), self.minor())
+        (self.python_major(), self.python_minor())
     }
 
+    /// Return the major version of the implementation (e.g., `CPython` or `PyPy`).
+    pub fn implementation_major(&self) -> u8 {
+        let major = self.markers.implementation_version.version.release()[0];
+        u8::try_from(major).expect("invalid major version")
+    }
+
+    /// Return the minor version of the implementation (e.g., `CPython` or `PyPy`).
+    pub fn implementation_minor(&self) -> u8 {
+        let minor = self.markers.implementation_version.version.release()[1];
+        u8::try_from(minor).expect("invalid minor version")
+    }
+
+    /// Returns the implementation version as a simple tuple.
+    pub fn implementation_tuple(&self) -> (u8, u8) {
+        (self.implementation_major(), self.implementation_minor())
+    }
+
+    pub fn implementation_name(&self) -> &str {
+        &self.markers.implementation_name
+    }
     pub fn base_exec_prefix(&self) -> &Path {
         &self.base_exec_prefix
     }

crates/puffin-interpreter/src/virtual_env.rs

@@ -78,7 +78,7 @@ impl Virtualenv {
     pub fn site_packages(&self) -> PathBuf {
         self.interpreter
             .platform
-            .venv_site_packages(&self.root, self.interpreter().simple_version())
+            .venv_site_packages(&self.root, self.interpreter().python_tuple())
     }
 
     pub fn bin_dir(&self) -> PathBuf {

crates/puffin-resolver/src/finder.rs

@@ -183,7 +183,7 @@ impl<'a> DistFinder<'a> {
                         .requires_python
                         .as_ref()
                         .map_or(true, |requires_python| {
-                            requires_python.contains(self.interpreter.version())
+                            requires_python.contains(self.interpreter.python_version())
                         })
                     {
                         continue;
@@ -219,7 +219,7 @@ impl<'a> DistFinder<'a> {
                         .requires_python
                         .as_ref()
                         .map_or(true, |requires_python| {
-                            requires_python.contains(self.interpreter.version())
+                            requires_python.contains(self.interpreter.python_version())
                         })
                     {
                         continue;

crates/puffin-resolver/src/python_requirement.rs

@@ -15,7 +15,7 @@ pub struct PythonRequirement {
 impl PythonRequirement {
     pub fn new(interpreter: &Interpreter, markers: &MarkerEnvironment) -> Self {
         Self {
-            installed: interpreter.version().clone(),
+            installed: interpreter.python_version().clone(),
             target: markers.python_full_version.version.clone(),
         }
     }

crates/puffin-resolver/tests/resolver.rs

@@ -694,6 +694,8 @@ static TAGS_311: Lazy<Tags> = Lazy::new(|| {
             Arch::Aarch64,
         ),
         (3, 11),
+        "cpython",
+        (3, 11),
     )
     .unwrap()
 });
@@ -724,6 +726,8 @@ static TAGS_310: Lazy<Tags> = Lazy::new(|| {
             Arch::Aarch64,
         ),
         (3, 10),
+        "cpython",
+        (3, 10),
     )
     .unwrap()
 });

crates/puffin/src/commands/pip_compile.rs

@@ -149,6 +149,8 @@ pub(crate) async fn pip_compile(
         Cow::Owned(Tags::from_env(
             interpreter.platform(),
             python_version.simple_version(),
+            interpreter.implementation_name(),
+            interpreter.implementation_tuple(),
         )?)
     } else {
         Cow::Borrowed(interpreter.tags()?)

crates/puffin/src/commands/venv.rs

@@ -68,7 +68,7 @@ enum VenvError {
 
     #[error("Failed to extract interpreter tags")]
     #[diagnostic(code(puffin::venv::tags))]
-    TagsError(#[source] platform_host::PlatformError),
+    TagsError(#[source] platform_tags::TagsError),
 
     #[error("Failed to resolve `--find-links` entry")]
     #[diagnostic(code(puffin::venv::flat_index))]
@@ -107,7 +107,7 @@ async fn venv_impl(
     writeln!(
         printer,
         "Using Python {} at {}",
-        interpreter.version(),
+        interpreter.python_version(),
         interpreter.sys_executable().display().cyan()
     )
     .into_diagnostic()?;