Support toolchain requests with platform-tag style Python implementations and version (#4407)

Closes #4399

---------

Co-authored-by: Andrew Gallant <[email protected]>

Author: zanieb

crates/uv-toolchain/src/discovery.rs

@@ -44,9 +44,9 @@ pub enum ToolchainRequest {
     File(PathBuf),
     /// The name of a Python executable (i.e. for lookup in the PATH) e.g. `foopython3`
     ExecutableName(String),
-    /// A Python implementation without a version e.g. `pypy`
+    /// A Python implementation without a version e.g. `pypy` or `pp`
     Implementation(ImplementationName),
-    /// A Python implementation name and version e.g. `pypy3.8` or `[email protected]`
+    /// A Python implementation name and version e.g. `pypy3.8` or `[email protected]` or `pp38`
     ImplementationVersion(ImplementationName, VersionRequest),
     /// A request for a specific toolchain key e.g. `cpython-3.12-x86_64-linux-gnu`
     /// Generally these refer to uv-managed toolchain downloads.
@@ -919,7 +919,7 @@ impl ToolchainRequest {
     ///
     /// This cannot fail, which means weird inputs will be parsed as [`ToolchainRequest::File`] or [`ToolchainRequest::ExecutableName`].
     pub fn parse(value: &str) -> Self {
-        // e.g. `3.12.1`
+        // e.g. `3.12.1`, `312`, or `>=3.12`
         if let Ok(version) = VersionRequest::from_str(value) {
             return Self::Version(version);
         }
@@ -937,18 +937,25 @@ impl ToolchainRequest {
                 }
             }
         }
-        for implementation in ImplementationName::iter() {
+        for implementation in ImplementationName::possible_names() {
             if let Some(remainder) = value
                 .to_ascii_lowercase()
                 .strip_prefix(Into::<&str>::into(implementation))
             {
                 // e.g. `pypy`
                 if remainder.is_empty() {
-                    return Self::Implementation(*implementation);
+                    return Self::Implementation(
+                        // Safety: The name matched the possible names above
+                        ImplementationName::from_str(implementation).unwrap(),
+                    );
                 }
-                // e.g. `pypy3.12`
+                // e.g. `pypy3.12` or `pp312`
                 if let Ok(version) = VersionRequest::from_str(remainder) {
-                    return Self::ImplementationVersion(*implementation, version);
+                    return Self::ImplementationVersion(
+                        // Safety: The name matched the possible names above
+                        ImplementationName::from_str(implementation).unwrap(),
+                        version,
+                    );
                 }
             }
         }
@@ -1267,8 +1274,22 @@ impl FromStr for VersionRequest {
     type Err = Error;
 
     fn from_str(s: &str) -> Result<Self, Self::Err> {
+        fn parse_nosep(s: &str) -> Option<VersionRequest> {
+            let mut chars = s.chars();
+            let major = chars.next()?.to_digit(10)?.try_into().ok()?;
+            if chars.as_str().is_empty() {
+                return Some(VersionRequest::Major(major));
+            }
+            let minor = chars.as_str().parse::<u8>().ok()?;
+            Some(VersionRequest::MajorMinor(major, minor))
+        }
+
+        // e.g. `3`, `38`, `312`
+        if let Some(request) = parse_nosep(s) {
+            Ok(request)
+        }
         // e.g. `3.12.1`
-        if let Ok(versions) = s
+        else if let Ok(versions) = s
             .splitn(3, '.')
             .map(str::parse::<u8>)
             .collect::<Result<Vec<_>, _>>()
@@ -1284,6 +1305,7 @@ impl FromStr for VersionRequest {
             };
 
             Ok(selector)
+
         // e.g. `>=3.12.1,<3.12`
         } else if let Ok(specifiers) = VersionSpecifiers::from_str(s) {
             if specifiers.is_empty() {
@@ -1549,6 +1571,7 @@ mod tests {
     use assert_fs::{prelude::*, TempDir};
     use test_log::test;
 
+    use super::Error;
     use crate::{
         discovery::{ToolchainRequest, VersionRequest},
         implementation::ImplementationName,
@@ -1587,13 +1610,35 @@ mod tests {
             ToolchainRequest::parse("pypy"),
             ToolchainRequest::Implementation(ImplementationName::PyPy)
         );
+        assert_eq!(
+            ToolchainRequest::parse("pp"),
+            ToolchainRequest::Implementation(ImplementationName::PyPy)
+        );
+        assert_eq!(
+            ToolchainRequest::parse("cp"),
+            ToolchainRequest::Implementation(ImplementationName::CPython)
+        );
         assert_eq!(
             ToolchainRequest::parse("pypy3.10"),
             ToolchainRequest::ImplementationVersion(
                 ImplementationName::PyPy,
                 VersionRequest::from_str("3.10").unwrap()
             )
         );
+        assert_eq!(
+            ToolchainRequest::parse("pp310"),
+            ToolchainRequest::ImplementationVersion(
+                ImplementationName::PyPy,
+                VersionRequest::from_str("3.10").unwrap()
+            )
+        );
+        assert_eq!(
+            ToolchainRequest::parse("cp38"),
+            ToolchainRequest::ImplementationVersion(
+                ImplementationName::CPython,
+                VersionRequest::from_str("3.8").unwrap()
+            )
+        );
         assert_eq!(
             ToolchainRequest::parse("[email protected]"),
             ToolchainRequest::ImplementationVersion(
@@ -1603,7 +1648,10 @@ mod tests {
         );
         assert_eq!(
             ToolchainRequest::parse("pypy310"),
-            ToolchainRequest::ExecutableName("pypy310".to_string())
+            ToolchainRequest::ImplementationVersion(
+                ImplementationName::PyPy,
+                VersionRequest::from_str("3.10").unwrap()
+            )
         );
 
         let tempdir = TempDir::new().unwrap();
@@ -1645,5 +1693,28 @@ mod tests {
             VersionRequest::MajorMinorPatch(3, 12, 1)
         );
         assert!(VersionRequest::from_str("1.foo.1").is_err());
+        assert_eq!(
+            VersionRequest::from_str("3").unwrap(),
+            VersionRequest::Major(3)
+        );
+        assert_eq!(
+            VersionRequest::from_str("38").unwrap(),
+            VersionRequest::MajorMinor(3, 8)
+        );
+        assert_eq!(
+            VersionRequest::from_str("312").unwrap(),
+            VersionRequest::MajorMinor(3, 12)
+        );
+        assert_eq!(
+            VersionRequest::from_str("3100").unwrap(),
+            VersionRequest::MajorMinor(3, 100)
+        );
+        assert!(
+            // Test for overflow
+            matches!(
+                VersionRequest::from_str("31000"),
+                Err(Error::InvalidVersionRequest(_))
+            )
+        );
     }
 }

crates/uv-toolchain/src/implementation.rs

@@ -24,10 +24,8 @@ pub enum LenientImplementationName {
 }
 
 impl ImplementationName {
-    pub(crate) fn iter() -> impl Iterator<Item = &'static ImplementationName> {
-        static NAMES: &[ImplementationName] =
-            &[ImplementationName::CPython, ImplementationName::PyPy];
-        NAMES.iter()
+    pub(crate) fn possible_names() -> impl Iterator<Item = &'static str> {
+        ["cpython", "pypy", "cp", "pp"].into_iter()
     }
 
     pub fn pretty(self) -> &'static str {
@@ -68,10 +66,13 @@ impl<'a> From<&'a LenientImplementationName> for &'a str {
 impl FromStr for ImplementationName {
     type Err = Error;
 
+    /// Parse a Python implementation name from a string.
+    ///
+    /// Supports the full name and the platform compatibility tag style name.
     fn from_str(s: &str) -> Result<Self, Self::Err> {
         match s.to_ascii_lowercase().as_str() {
-            "cpython" => Ok(Self::CPython),
-            "pypy" => Ok(Self::PyPy),
+            "cpython" | "cp" => Ok(Self::CPython),
+            "pypy" | "pp" => Ok(Self::PyPy),
             _ => Err(Error::UnknownImplementation(s.to_string())),
         }
     }