Allow version specifiers to be used in Python version requests

crates/uv-toolchain/src/discovery.rs

@@ -1,11 +1,11 @@
 use std::borrow::Cow;
 use std::collections::HashSet;
 use std::fmt::{self, Formatter};
-use std::num::ParseIntError;
 use std::{env, io};
 use std::{path::Path, path::PathBuf, str::FromStr};
 
 use itertools::Itertools;
+use pep440_rs::{Version, VersionSpecifiers};
 use same_file::is_same_file;
 use thiserror::Error;
 use tracing::{debug, instrument, trace};
@@ -35,7 +35,7 @@ pub enum ToolchainRequest {
     /// Use any discovered Python toolchain
     #[default]
     Any,
-    /// A Python version without an implementation name e.g. `3.10`
+    /// A Python version without an implementation name e.g. `3.10` or `>=3.12,<3.13`
     Version(VersionRequest),
     /// A path to a directory containing a Python installation, e.g. `.venv`
     Directory(PathBuf),
@@ -63,13 +63,14 @@ pub enum ToolchainSources {
 }
 
 /// A Python toolchain version request.
-#[derive(Copy, Clone, Debug, Default, PartialEq, Eq)]
+#[derive(Clone, Debug, Default, PartialEq, Eq)]
 pub enum VersionRequest {
     #[default]
     Any,
     Major(u8),
     MajorMinor(u8, u8),
     MajorMinorPatch(u8, u8, u8),
+    Range(VersionSpecifiers),
 }
 
 /// The policy for discovery of "system" Python interpreters.
@@ -159,6 +160,10 @@ pub enum Error {
     #[error(transparent)]
     PyLauncher(#[from] crate::py_launcher::Error),
 
+    /// An invalid version request was given
+    #[error("Invalid version request: {0}")]
+    InvalidVersionRequest(String),
+
     #[error("Interpreter discovery for `{0}` requires `{1}` but it is not selected; the following are selected: {2}")]
     SourceNotSelected(ToolchainRequest, ToolchainSource, ToolchainSources),
 }
@@ -563,7 +568,7 @@ pub(crate) fn find_toolchain(
                     ToolchainNotFound::NoMatchingImplementationVersion(
                         sources.clone(),
                         *implementation,
-                        *version,
+                        version.clone(),
                     ),
                 ));
             };
@@ -608,9 +613,9 @@ pub(crate) fn find_toolchain(
                     .transpose()?
             else {
                 let err = if matches!(version, VersionRequest::Any) {
-                    ToolchainNotFound::NoPythonInstallation(sources.clone(), Some(*version))
+                    ToolchainNotFound::NoPythonInstallation(sources.clone(), Some(version.clone()))
                 } else {
-                    ToolchainNotFound::NoMatchingVersion(sources.clone(), *version)
+                    ToolchainNotFound::NoMatchingVersion(sources.clone(), version.clone())
                 };
                 return Ok(ToolchainResult::Err(err));
             };
@@ -679,14 +684,17 @@ pub fn find_best_toolchain(
     if let Some(request) = match request {
         ToolchainRequest::Version(version) => {
             if version.has_patch() {
-                Some(ToolchainRequest::Version((*version).without_patch()))
+                Some(ToolchainRequest::Version(version.clone().without_patch()))
             } else {
                 None
             }
         }
-        ToolchainRequest::ImplementationVersion(implementation, version) => Some(
-            ToolchainRequest::ImplementationVersion(*implementation, (*version).without_patch()),
-        ),
+        ToolchainRequest::ImplementationVersion(implementation, version) => {
+            Some(ToolchainRequest::ImplementationVersion(
+                *implementation,
+                version.clone().without_patch(),
+            ))
+        }
         _ => None,
     } {
         debug!("Looking for relaxed patch version {request}");
@@ -1026,7 +1034,7 @@ impl ToolchainRequest {
 }
 
 impl VersionRequest {
-    pub(crate) fn default_names(self) -> [Option<Cow<'static, str>>; 4] {
+    pub(crate) fn default_names(&self) -> [Option<Cow<'static, str>>; 4] {
         let (python, python3, extension) = if cfg!(windows) {
             (
                 Cow::Borrowed("python.exe"),
@@ -1038,7 +1046,7 @@ impl VersionRequest {
         };
 
         match self {
-            Self::Any => [Some(python3), Some(python), None, None],
+            Self::Any | Self::Range(_) => [Some(python3), Some(python), None, None],
             Self::Major(major) => [
                 Some(Cow::Owned(format!("python{major}{extension}"))),
                 Some(python),
@@ -1081,7 +1089,7 @@ impl VersionRequest {
                 };
 
                 match self {
-                    Self::Any => [Some(python3), Some(python), None, None],
+                    Self::Any | Self::Range(_) => [Some(python3), Some(python), None, None],
                     Self::Major(major) => [
                         Some(Cow::Owned(format!("{name}{major}{extension}"))),
                         Some(python),
@@ -1109,98 +1117,126 @@ impl VersionRequest {
     }
 
     /// Check if a interpreter matches the requested Python version.
-    fn matches_interpreter(self, interpreter: &Interpreter) -> bool {
+    fn matches_interpreter(&self, interpreter: &Interpreter) -> bool {
         match self {
             Self::Any => true,
-            Self::Major(major) => interpreter.python_major() == major,
+            Self::Major(major) => interpreter.python_major() == *major,
             Self::MajorMinor(major, minor) => {
-                (interpreter.python_major(), interpreter.python_minor()) == (major, minor)
+                (interpreter.python_major(), interpreter.python_minor()) == (*major, *minor)
             }
             Self::MajorMinorPatch(major, minor, patch) => {
                 (
                     interpreter.python_major(),
                     interpreter.python_minor(),
                     interpreter.python_patch(),
-                ) == (major, minor, patch)
+                ) == (*major, *minor, *patch)
             }
+            Self::Range(specifiers) => specifiers.contains(interpreter.python_version()),
         }
     }
 
-    pub(crate) fn matches_version(self, version: &PythonVersion) -> bool {
+    pub(crate) fn matches_version(&self, version: &PythonVersion) -> bool {
         match self {
             Self::Any => true,
-            Self::Major(major) => version.major() == major,
-            Self::MajorMinor(major, minor) => (version.major(), version.minor()) == (major, minor),
+            Self::Major(major) => version.major() == *major,
+            Self::MajorMinor(major, minor) => {
+                (version.major(), version.minor()) == (*major, *minor)
+            }
             Self::MajorMinorPatch(major, minor, patch) => {
-                (version.major(), version.minor(), version.patch()) == (major, minor, Some(patch))
+                (version.major(), version.minor(), version.patch())
+                    == (*major, *minor, Some(*patch))
             }
+            Self::Range(specifiers) => specifiers.contains(&version.version),
         }
     }
 
-    fn matches_major_minor(self, major: u8, minor: u8) -> bool {
+    fn matches_major_minor(&self, major: u8, minor: u8) -> bool {
         match self {
             Self::Any => true,
-            Self::Major(self_major) => self_major == major,
-            Self::MajorMinor(self_major, self_minor) => (self_major, self_minor) == (major, minor),
+            Self::Major(self_major) => *self_major == major,
+            Self::MajorMinor(self_major, self_minor) => {
+                (*self_major, *self_minor) == (major, minor)
+            }
             Self::MajorMinorPatch(self_major, self_minor, _) => {
-                (self_major, self_minor) == (major, minor)
+                (*self_major, *self_minor) == (major, minor)
+            }
+            Self::Range(specifiers) => {
+                specifiers.contains(&Version::new([u64::from(major), u64::from(minor)]))
             }
         }
     }
 
-    pub(crate) fn matches_major_minor_patch(self, major: u8, minor: u8, patch: u8) -> bool {
+    pub(crate) fn matches_major_minor_patch(&self, major: u8, minor: u8, patch: u8) -> bool {
         match self {
             Self::Any => true,
-            Self::Major(self_major) => self_major == major,
-            Self::MajorMinor(self_major, self_minor) => (self_major, self_minor) == (major, minor),
+            Self::Major(self_major) => *self_major == major,
+            Self::MajorMinor(self_major, self_minor) => {
+                (*self_major, *self_minor) == (major, minor)
+            }
             Self::MajorMinorPatch(self_major, self_minor, self_patch) => {
-                (self_major, self_minor, self_patch) == (major, minor, patch)
+                (*self_major, *self_minor, *self_patch) == (major, minor, patch)
             }
+            Self::Range(specifiers) => specifiers.contains(&Version::new([
+                u64::from(major),
+                u64::from(minor),
+                u64::from(patch),
+            ])),
         }
     }
 
     /// Return true if a patch version is present in the request.
-    fn has_patch(self) -> bool {
+    fn has_patch(&self) -> bool {
         match self {
             Self::Any => false,
             Self::Major(..) => false,
             Self::MajorMinor(..) => false,
             Self::MajorMinorPatch(..) => true,
+            Self::Range(_) => false,
         }
     }
 
-    /// Return a new `VersionRequest` without the patch version.
+    /// Return a new [`VersionRequest`] without the patch version if possible.
+    ///
+    /// If the patch version is not present, it is returned unchanged.
     #[must_use]
     fn without_patch(self) -> Self {
         match self {
             Self::Any => Self::Any,
             Self::Major(major) => Self::Major(major),
             Self::MajorMinor(major, minor) => Self::MajorMinor(major, minor),
             Self::MajorMinorPatch(major, minor, _) => Self::MajorMinor(major, minor),
+            Self::Range(_) => self,
         }
     }
 }
 
 impl FromStr for VersionRequest {
-    type Err = ParseIntError;
+    type Err = Error;
 
     fn from_str(s: &str) -> Result<Self, Self::Err> {
-        let versions = s
+        // e.g. `3.12.1`
+        if let Ok(versions) = s
             .splitn(3, '.')
             .map(str::parse::<u8>)
-            .collect::<Result<Vec<_>, _>>()?;
-
-        let selector = match versions.as_slice() {
-            // e.g. `3`
-            [major] => VersionRequest::Major(*major),
-            // e.g. `3.10`
-            [major, minor] => VersionRequest::MajorMinor(*major, *minor),
-            // e.g. `3.10.4`
-            [major, minor, patch] => VersionRequest::MajorMinorPatch(*major, *minor, *patch),
-            _ => unreachable!(),
-        };
+            .collect::<Result<Vec<_>, _>>()
+        {
+            let selector = match versions.as_slice() {
+                // e.g. `3`
+                [major] => VersionRequest::Major(*major),
+                // e.g. `3.10`
+                [major, minor] => VersionRequest::MajorMinor(*major, *minor),
+                // e.g. `3.10.4`
+                [major, minor, patch] => VersionRequest::MajorMinorPatch(*major, *minor, *patch),
+                _ => unreachable!(),
+            };
 
-        Ok(selector)
+            Ok(selector)
+        // e.g. `>=3.12.1,<3.12`
+        } else if let Ok(specifiers) = VersionSpecifiers::from_str(s) {
+            Ok(Self::Range(specifiers))
+        } else {
+            Err(Error::InvalidVersionRequest(s.to_string()))
+        }
     }
 }
 
@@ -1220,6 +1256,7 @@ impl fmt::Display for VersionRequest {
             Self::MajorMinorPatch(major, minor, patch) => {
                 write!(f, "{major}.{minor}.{patch}")
             }
+            Self::Range(specifiers) => write!(f, "{specifiers}"),
         }
     }
 }
@@ -1458,6 +1495,14 @@ mod tests {
             ToolchainRequest::parse("3.12"),
             ToolchainRequest::Version(VersionRequest::from_str("3.12").unwrap())
         );
+        assert_eq!(
+            ToolchainRequest::parse(">=3.12"),
+            ToolchainRequest::Version(VersionRequest::from_str(">=3.12").unwrap())
+        );
+        assert_eq!(
+            ToolchainRequest::parse(">=3.12"),
+            ToolchainRequest::Version(VersionRequest::from_str(">=3.12,<3.13").unwrap())
+        );
         assert_eq!(
             ToolchainRequest::parse("foo"),
             ToolchainRequest::ExecutableName("foo".to_string())
@@ -1522,14 +1567,17 @@ mod tests {
 
     #[test]
     fn version_request_from_str() {
-        assert_eq!(VersionRequest::from_str("3"), Ok(VersionRequest::Major(3)));
         assert_eq!(
-            VersionRequest::from_str("3.12"),
-            Ok(VersionRequest::MajorMinor(3, 12))
+            VersionRequest::from_str("3").unwrap(),
+            VersionRequest::Major(3)
+        );
+        assert_eq!(
+            VersionRequest::from_str("3.12").unwrap(),
+            VersionRequest::MajorMinor(3, 12)
         );
         assert_eq!(
-            VersionRequest::from_str("3.12.1"),
-            Ok(VersionRequest::MajorMinorPatch(3, 12, 1))
+            VersionRequest::from_str("3.12.1").unwrap(),
+            VersionRequest::MajorMinorPatch(3, 12, 1)
         );
         assert!(VersionRequest::from_str("1.foo.1").is_err());
     }

crates/uv-toolchain/src/downloads.rs

@@ -1,6 +1,5 @@
 use std::fmt::Display;
 use std::io;
-use std::num::ParseIntError;
 use std::path::{Path, PathBuf};
 use std::str::FromStr;
 
@@ -26,7 +25,7 @@ pub enum Error {
     #[error(transparent)]
     ImplementationError(#[from] ImplementationError),
     #[error("Invalid python version: {0}")]
-    InvalidPythonVersion(ParseIntError),
+    InvalidPythonVersion(String),
     #[error("Download failed")]
     NetworkError(#[from] BetterReqwestError),
     #[error("Download failed")]
@@ -215,7 +214,7 @@ impl PythonDownloadRequest {
 impl Display for PythonDownloadRequest {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         let mut parts = Vec::new();
-        if let Some(version) = self.version {
+        if let Some(version) = &self.version {
             parts.push(version.to_string());
         }
         if let Some(implementation) = self.implementation {
@@ -239,7 +238,8 @@ impl FromStr for PythonDownloadRequest {
 
     fn from_str(s: &str) -> Result<Self, Self::Err> {
         // TODO(zanieb): Implement parsing of additional request parts
-        let version = VersionRequest::from_str(s).map_err(Error::InvalidPythonVersion)?;
+        let version =
+            VersionRequest::from_str(s).map_err(|_| Error::InvalidPythonVersion(s.to_string()))?;
         Ok(Self::new(Some(version), None, None, None, None))
     }
 }