Implement Toolchain::find_or_fetch and use in uv venv --preview (#4138)

Extends #4121
Part of #2607 

Adds support for managed toolchain fetching to `uv venv`, e.g.

```
❯ cargo run -q -- venv --python 3.9.18 --preview -v
DEBUG Searching for Python 3.9.18 in search path or managed toolchains
DEBUG Searching for managed toolchains at `/Users/zb/Library/Application Support/uv/toolchains`
DEBUG Found CPython 3.12.3 at `/opt/homebrew/bin/python3` (search path)
DEBUG Found CPython 3.9.6 at `/usr/bin/python3` (search path)
DEBUG Found CPython 3.12.3 at `/opt/homebrew/bin/python3` (search path)
DEBUG Requested Python not found, checking for available download...
DEBUG Using registry request timeout of 30s
INFO Fetching requested toolchain...
DEBUG Downloading https://github.com/indygreg/python-build-standalone/releases/download/20240224/cpython-3.9.18%2B20240224-aarch64-apple-darwin-pgo%2Blto-full.tar.zst to temporary location /Users/zb/Library/Application Support/uv/toolchains/.tmpgohKwp
DEBUG Extracting cpython-3.9.18%2B20240224-aarch64-apple-darwin-pgo%2Blto-full.tar.zst
DEBUG Moving /Users/zb/Library/Application Support/uv/toolchains/.tmpgohKwp/python to /Users/zb/Library/Application Support/uv/toolchains/cpython-3.9.18-macos-aarch64-none
Using Python 3.9.18 interpreter at: /Users/zb/Library/Application Support/uv/toolchains/cpython-3.9.18-macos-aarch64-none/install/bin/python3
Creating virtualenv at: .venv
INFO Removing existing directory
Activate with: source .venv/bin/activate
```

The preview flag is required. The fetch is performed if we can't find an
interpreter that satisfies the request. Once fetched, the toolchain will
be available for later invocations that include the `--preview` flag.
There will be follow-ups to improve toolchain management in general,
there is still outstanding work from the initial implementation.

Author: zanieb

Cargo.lock

@@ -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/Cargo.toml

@@ -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-dev/src/fetch_python.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/discovery.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, Default)]
 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,27 @@ impl PythonDownloadRequest {
         self
     }
 
+    pub fn from_request(request: ToolchainRequest) -> Result<Self, Error> {
+        let result = Self::default();
+        let 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(_)
+            | ToolchainRequest::ExecutableName(_)
+            | 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 +166,34 @@ 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 {
+            parts.push(version.to_string());
+        }
+        if let Some(implementation) = self.implementation {
+            parts.push(implementation.to_string());
+        }
+        if let Some(os) = &self.os {
+            parts.push(os.to_string());
+        }
+        if let Some(arch) = self.arch {
+            parts.push(arch.to_string());
+        }
+        if let Some(libc) = self.libc {
+            parts.push(libc.to_string());
+        }
+        write!(f, "{}", parts.join("-"))
+    }
+}
+
 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 +211,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 +229,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 +283,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/downloads.rs

@@ -56,6 +56,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),
 }