Implement Toolchain::find_or_fetch and use in uv venv --preview

crates/uv-dev/Cargo.toml

@@ -42,7 +42,6 @@ anyhow = { workspace = true }
 clap = { workspace = true, features = ["derive", "wrap_help"] }
 fs-err = { workspace = true, features = ["tokio"] }
 futures = { workspace = true }
-itertools = { workspace = true }
 owo-colors = { workspace = true }
 poloto = { version = "19.1.2", optional = true }
 pretty_assertions = { version = "1.4.0" }

crates/uv-dev/src/fetch_python.rs

@@ -1,16 +1,10 @@
 use anyhow::Result;
 use clap::Parser;
 use fs_err as fs;
-#[cfg(unix)]
-use fs_err::tokio::symlink;
 use futures::StreamExt;
-#[cfg(unix)]
-use itertools::Itertools;
-use std::str::FromStr;
-#[cfg(unix)]
-use std::{collections::HashMap, path::PathBuf};
 use tokio::time::Instant;
 use tracing::{info, info_span, Instrument};
+use uv_toolchain::ToolchainRequest;
 
 use uv_fs::Simplified;
 use uv_toolchain::downloads::{DownloadResult, Error, PythonDownload, PythonDownloadRequest};
@@ -37,17 +31,16 @@ pub(crate) async fn fetch_python(args: FetchPythonArgs) -> Result<()> {
     let requests = versions
         .iter()
         .map(|version| {
-            PythonDownloadRequest::from_str(version).and_then(PythonDownloadRequest::fill)
+            PythonDownloadRequest::from_request(ToolchainRequest::parse(version))
+                // Populate platform information on the request
+                .and_then(PythonDownloadRequest::fill)
         })
         .collect::<Result<Vec<_>, Error>>()?;
 
     let downloads = requests
         .iter()
-        .map(|request| match PythonDownload::from_request(request) {
-            Some(download) => download,
-            None => panic!("No download found for request {request:?}"),
-        })
-        .collect::<Vec<_>>();
+        .map(PythonDownload::from_request)
+        .collect::<Result<Vec<_>, Error>>()?;
 
     let client = uv_client::BaseClientBuilder::new().build();
 
@@ -91,40 +84,6 @@ pub(crate) async fn fetch_python(args: FetchPythonArgs) -> Result<()> {
         info!("All versions downloaded already.");
     };
 
-    // Order matters here, as we overwrite previous links
-    info!("Installing to `{}`...", toolchain_dir.user_display());
-
-    // On Windows, linking the executable generally results in broken installations
-    // and each toolchain path will need to be added to the PATH separately in the
-    // desired order
-    #[cfg(unix)]
-    {
-        let mut links: HashMap<PathBuf, PathBuf> = HashMap::new();
-        for (version, path) in results {
-            // TODO(zanieb): This path should be a part of the download metadata
-            let executable = path.join("install").join("bin").join("python3");
-            for target in [
-                toolchain_dir.join(format!("python{}", version.python_full_version())),
-                toolchain_dir.join(format!("python{}.{}", version.major(), version.minor())),
-                toolchain_dir.join(format!("python{}", version.major())),
-                toolchain_dir.join("python"),
-            ] {
-                // Attempt to remove it, we'll fail on link if we couldn't remove it for some reason
-                // but if it's missing we don't want to error
-                let _ = fs::remove_file(&target);
-                symlink(&executable, &target).await?;
-                links.insert(target, executable.clone());
-            }
-        }
-        for (target, executable) in links.iter().sorted() {
-            info!(
-                "Linked `{}` to `{}`",
-                target.user_display(),
-                executable.user_display()
-            );
-        }
-    };
-
     info!("Installed {} versions", requests.len());
 
     Ok(())

crates/uv-toolchain/src/discovery.rs

@@ -1059,6 +1059,17 @@ impl VersionRequest {
         }
     }
 
+    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::MajorMinorPatch(self_major, self_minor, self_patch) => {
+                (self_major, self_minor, self_patch) == (major, minor, patch)
+            }
+        }
+    }
+
     /// Return true if a patch version is present in the request.
     fn has_patch(self) -> bool {
         match self {

crates/uv-toolchain/src/downloads.rs

@@ -1,11 +1,12 @@
 use std::fmt::Display;
 use std::io;
+use std::num::ParseIntError;
 use std::path::{Path, PathBuf};
 use std::str::FromStr;
 
 use crate::implementation::{Error as ImplementationError, ImplementationName};
 use crate::platform::{Arch, Error as PlatformError, Libc, Os};
-use crate::PythonVersion;
+use crate::{PythonVersion, ToolchainRequest, VersionRequest};
 use thiserror::Error;
 use uv_client::BetterReqwestError;
 
@@ -25,13 +26,13 @@ pub enum Error {
     #[error(transparent)]
     ImplementationError(#[from] ImplementationError),
     #[error("Invalid python version: {0}")]
-    InvalidPythonVersion(String),
+    InvalidPythonVersion(ParseIntError),
     #[error("Download failed")]
     NetworkError(#[from] BetterReqwestError),
     #[error("Download failed")]
     NetworkMiddlewareError(#[source] anyhow::Error),
-    #[error(transparent)]
-    ExtractError(#[from] uv_extract::Error),
+    #[error("Failed to extract archive: {0}")]
+    ExtractError(String, #[source] uv_extract::Error),
     #[error("Invalid download url")]
     InvalidUrl(#[from] url::ParseError),
     #[error("Failed to create download directory")]
@@ -50,6 +51,11 @@ pub enum Error {
     },
     #[error("Failed to parse toolchain directory name: {0}")]
     NameError(String),
+    #[error("Cannot download toolchain for request: {0}")]
+    InvalidRequestKind(ToolchainRequest),
+    // TODO(zanieb): Implement display for `PythonDownloadRequest`
+    #[error("No download found for request: {0:?}")]
+    NoDownloadFound(PythonDownloadRequest),
 }
 
 #[derive(Debug, PartialEq)]
@@ -66,9 +72,9 @@ pub struct PythonDownload {
     sha256: Option<&'static str>,
 }
 
-#[derive(Debug)]
+#[derive(Debug, Clone)]
 pub struct PythonDownloadRequest {
-    version: Option<PythonVersion>,
+    version: Option<VersionRequest>,
     implementation: Option<ImplementationName>,
     arch: Option<Arch>,
     os: Option<Os>,
@@ -77,7 +83,7 @@ pub struct PythonDownloadRequest {
 
 impl PythonDownloadRequest {
     pub fn new(
-        version: Option<PythonVersion>,
+        version: Option<VersionRequest>,
         implementation: Option<ImplementationName>,
         arch: Option<Arch>,
         os: Option<Os>,
@@ -98,6 +104,12 @@ impl PythonDownloadRequest {
         self
     }
 
+    #[must_use]
+    pub fn with_version(mut self, version: VersionRequest) -> Self {
+        self.version = Some(version);
+        self
+    }
+
     #[must_use]
     pub fn with_arch(mut self, arch: Arch) -> Self {
         self.arch = Some(arch);
@@ -116,6 +128,32 @@ impl PythonDownloadRequest {
         self
     }
 
+    pub fn from_request(request: ToolchainRequest) -> Result<Self, Error> {
+        let mut result = Self::default();
+        result = match request {
+            ToolchainRequest::Version(version) => result.with_version(version),
+
+            ToolchainRequest::Implementation(implementation) => {
+                result.with_implementation(implementation)
+            }
+            ToolchainRequest::ImplementationVersion(implementation, version) => result
+                .with_implementation(implementation)
+                .with_version(version),
+            ToolchainRequest::Any => result,
+            // We can't download a toolchain for these request kinds
+            ToolchainRequest::Directory(_) => {
+                return Err(Error::InvalidRequestKind(request));
+            }
+            ToolchainRequest::ExecutableName(_) => {
+                return Err(Error::InvalidRequestKind(request));
+            }
+            ToolchainRequest::File(_) => {
+                return Err(Error::InvalidRequestKind(request));
+            }
+        };
+        Ok(result)
+    }
+
     pub fn fill(mut self) -> Result<Self, Error> {
         if self.implementation.is_none() {
             self.implementation = Some(ImplementationName::CPython);
@@ -133,12 +171,18 @@ impl PythonDownloadRequest {
     }
 }
 
+impl Default for PythonDownloadRequest {
+    fn default() -> Self {
+        Self::new(None, None, None, None, None)
+    }
+}
+
 impl FromStr for PythonDownloadRequest {
     type Err = Error;
 
     fn from_str(s: &str) -> Result<Self, Self::Err> {
         // TODO(zanieb): Implement parsing of additional request parts
-        let version = PythonVersion::from_str(s).map_err(Error::InvalidPythonVersion)?;
+        let version = VersionRequest::from_str(s).map_err(Error::InvalidPythonVersion)?;
         Ok(Self::new(Some(version), None, None, None, None))
     }
 }
@@ -156,7 +200,7 @@ impl PythonDownload {
         PYTHON_DOWNLOADS.iter().find(|&value| value.key == key)
     }
 
-    pub fn from_request(request: &PythonDownloadRequest) -> Option<&'static PythonDownload> {
+    pub fn from_request(request: &PythonDownloadRequest) -> Result<&'static PythonDownload, Error> {
         for download in PYTHON_DOWNLOADS {
             if let Some(arch) = &request.arch {
                 if download.arch != *arch {
@@ -174,21 +218,17 @@ impl PythonDownload {
                 }
             }
             if let Some(version) = &request.version {
-                if download.major != version.major() {
-                    continue;
-                }
-                if download.minor != version.minor() {
+                if !version.matches_major_minor_patch(
+                    download.major,
+                    download.minor,
+                    download.patch,
+                ) {
                     continue;
                 }
-                if let Some(patch) = version.patch() {
-                    if download.patch != patch {
-                        continue;
-                    }
-                }
             }
-            return Some(download);
+            return Ok(download);
         }
-        None
+        Err(Error::NoDownloadFound(request.clone()))
     }
 
     pub fn url(&self) -> &str {
@@ -232,13 +272,15 @@ impl PythonDownload {
             .into_async_read();
 
         debug!("Extracting {filename}");
-        uv_extract::stream::archive(reader.compat(), filename, temp_dir.path()).await?;
+        uv_extract::stream::archive(reader.compat(), filename, temp_dir.path())
+            .await
+            .map_err(|err| Error::ExtractError(filename.to_string(), err))?;
 
         // Extract the top-level directory.
         let extracted = match uv_extract::strip_component(temp_dir.path()) {
             Ok(top_level) => top_level,
             Err(uv_extract::Error::NonSingularArchive(_)) => temp_dir.into_path(),
-            Err(err) => return Err(err.into()),
+            Err(err) => return Err(Error::ExtractError(filename.to_string(), err)),
         };
 
         // Persist it to the target

crates/uv-toolchain/src/lib.rs

@@ -57,6 +57,12 @@ pub enum Error {
     #[error(transparent)]
     PyLauncher(#[from] py_launcher::Error),
 
+    #[error(transparent)]
+    ManagedToolchain(#[from] managed::Error),
+
+    #[error(transparent)]
+    Download(#[from] downloads::Error),
+
     #[error(transparent)]
     NotFound(#[from] ToolchainNotFound),
 }