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

Extends https://github.com/astral-sh/uv/pull/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 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.

Cargo.lock

@@ -4593,7 +4593,6 @@ dependencies = [
  "fs-err",
  "futures",
  "install-wheel-rs",
- "itertools 0.13.0",
  "mimalloc",
  "owo-colors",
  "pep508_rs",

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) {
+        .map(PythonDownload::from_request)
-            Some(download) => download,
+        .collect::<Result<Vec<_>, Error>>()?;
-            None => panic!("No download found for request {request:?}"),
-        })
-        .collect::<Vec<_>>();
 
     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)]
+    #[error("Failed to extract archive: {0}")]
-    ExtractError(#[from] uv_extract::Error),
+    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() {
+                if !version.matches_major_minor_patch(
-                    continue;
+                    download.major,
-                }
+                    download.minor,
-                if download.minor != version.minor() {
+                    download.patch,
-                    continue;
+                ) {
-                }
-                if let Some(patch) = version.patch() {
-                    if download.patch != patch {
                     continue;
                 }
             }
+            return Ok(download);
         }
-            return Some(download);
+        Err(Error::NoDownloadFound(request.clone()))
-        }
-        None
     }
 
     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/lib.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),
 }

crates/uv-toolchain/src/managed.rs

@@ -5,14 +5,46 @@ use std::ffi::OsStr;
 use std::io::{self, Write};
 use std::path::{Path, PathBuf};
 use std::str::FromStr;
+use thiserror::Error;
 
 use uv_state::{StateBucket, StateStore};
 
-// TODO(zanieb): Separate download and managed error types
+use crate::downloads::Error as DownloadError;
-pub use crate::downloads::Error;
+use crate::implementation::Error as ImplementationError;
+use crate::platform::Error as PlatformError;
 use crate::platform::{Arch, Libc, Os};
 use crate::python_version::PythonVersion;
+use uv_fs::Simplified;
 
+#[derive(Error, Debug)]
+pub enum Error {
+    #[error(transparent)]
+    IO(#[from] io::Error),
+    #[error(transparent)]
+    Download(#[from] DownloadError),
+    #[error(transparent)]
+    PlatformError(#[from] PlatformError),
+    #[error(transparent)]
+    ImplementationError(#[from] ImplementationError),
+    #[error("Invalid python version: {0}")]
+    InvalidPythonVersion(String),
+    #[error(transparent)]
+    ExtractError(#[from] uv_extract::Error),
+    #[error("Failed to copy to: {0}", to.user_display())]
+    CopyError {
+        to: PathBuf,
+        #[source]
+        err: io::Error,
+    },
+    #[error("Failed to read toolchain directory: {0}", dir.user_display())]
+    ReadError {
+        dir: PathBuf,
+        #[source]
+        err: io::Error,
+    },
+    #[error("Failed to parse toolchain directory name: {0}")]
+    NameError(String),
+}
 /// A collection of uv-managed Python toolchains installed on the current system.
 #[derive(Debug, Clone)]
 pub struct InstalledToolchains {
@@ -22,31 +54,35 @@ pub struct InstalledToolchains {
 
 impl InstalledToolchains {
     /// A directory for installed toolchains at `root`.
-    pub fn from_path(root: impl Into<PathBuf>) -> Result<Self, io::Error> {
+    fn from_path(root: impl Into<PathBuf>) -> Self {
-        Ok(Self { root: root.into() })
+        Self { root: root.into() }
     }
 
     /// Prefer, in order:
     /// 1. The specific toolchain directory specified by the user, i.e., `UV_TOOLCHAIN_DIR`
     /// 2. A directory in the system-appropriate user-level data directory, e.g., `~/.local/uv/toolchains`
     /// 3. A directory in the local data directory, e.g., `./.uv/toolchains`
-    pub fn from_settings() -> Result<Self, io::Error> {
+    pub fn from_settings() -> Result<Self, Error> {
         if let Some(toolchain_dir) = std::env::var_os("UV_TOOLCHAIN_DIR") {
-            Self::from_path(toolchain_dir)
+            Ok(Self::from_path(toolchain_dir))
         } else {
-            Self::from_path(StateStore::from_settings(None)?.bucket(StateBucket::Toolchains))
+            Ok(Self::from_path(
+                StateStore::from_settings(None)?.bucket(StateBucket::Toolchains),
+            ))
         }
     }
 
     /// Create a temporary installed toolchain directory.
-    pub fn temp() -> Result<Self, io::Error> {
+    pub fn temp() -> Result<Self, Error> {
-        Self::from_path(StateStore::temp()?.bucket(StateBucket::Toolchains))
+        Ok(Self::from_path(
+            StateStore::temp()?.bucket(StateBucket::Toolchains),
+        ))
     }
 
     /// Initialize the installed toolchain directory.
     ///
     /// Ensures the directory is created.
-    pub fn init(self) -> Result<Self, io::Error> {
+    pub fn init(self) -> Result<Self, Error> {
         let root = &self.root;
 
         // Create the cache directory, if it doesn't exist.
@@ -60,7 +96,7 @@ impl InstalledToolchains {
         {
             Ok(mut file) => file.write_all(b"*")?,
             Err(err) if err.kind() == io::ErrorKind::AlreadyExists => (),
-            Err(err) => return Err(err),
+            Err(err) => return Err(err.into()),
         }
 
         Ok(self)

crates/uv-toolchain/src/toolchain.rs

@@ -1,3 +1,5 @@
+use tracing::{debug, info};
+use uv_client::BaseClientBuilder;
 use uv_configuration::PreviewMode;
 
 use uv_cache::Cache;
@@ -6,6 +8,8 @@ use crate::discovery::{
     find_best_toolchain, find_default_toolchain, find_toolchain, SystemPython, ToolchainRequest,
     ToolchainSources,
 };
+use crate::downloads::{DownloadResult, PythonDownload, PythonDownloadRequest};
+use crate::managed::{InstalledToolchain, InstalledToolchains};
 use crate::{Error, Interpreter, ToolchainSource};
 
 /// A Python interpreter and accompanying tools.
@@ -114,6 +118,60 @@ impl Toolchain {
         Ok(toolchain)
     }
 
+    /// Find or fetch a [`Toolchain`].
+    ///
+    /// Unlike [`Toolchain::find`], if the toolchain is not installed it will be installed automatically.
+    pub async fn find_or_fetch<'a>(
+        python: Option<&str>,
+        system: SystemPython,
+        preview: PreviewMode,
+        client_builder: BaseClientBuilder<'a>,
+        cache: &Cache,
+    ) -> Result<Self, Error> {
+        // Perform a find first
+        match Self::find(python, system, preview, cache) {
+            Ok(venv) => Ok(venv),
+            Err(Error::NotFound(_)) if system.is_allowed() && preview.is_enabled() => {
+                debug!("Requested Python not found, checking for available download...");
+                let request = if let Some(request) = python {
+                    ToolchainRequest::parse(request)
+                } else {
+                    ToolchainRequest::default()
+                };
+                Self::fetch(request, client_builder, cache).await
+            }
+            Err(err) => Err(err),
+        }
+    }
+
+    pub async fn fetch<'a>(
+        request: ToolchainRequest,
+        client_builder: BaseClientBuilder<'a>,
+        cache: &Cache,
+    ) -> Result<Self, Error> {
+        let toolchains = InstalledToolchains::from_settings()?.init()?;
+        let toolchain_dir = toolchains.root();
+
+        let request = PythonDownloadRequest::from_request(request)?.fill()?;
+        let download = PythonDownload::from_request(&request)?;
+        let client = client_builder.build();
+
+        info!("Fetching requested toolchain...");
+        let result = download.fetch(&client, toolchain_dir).await?;
+
+        let path = match result {
+            DownloadResult::AlreadyAvailable(path) => path,
+            DownloadResult::Fetched(path) => path,
+        };
+
+        let installed = InstalledToolchain::new(path)?;
+
+        Ok(Self {
+            source: ToolchainSource::Managed,
+            interpreter: Interpreter::query(installed.executable(), cache)?,
+        })
+    }
+
     /// Create a [`Toolchain`] from an existing [`Interpreter`].
     pub fn from_interpreter(interpreter: Interpreter) -> Self {
         Self {

crates/uv/src/commands/venv.rs

@@ -14,7 +14,7 @@ use install_wheel_rs::linker::LinkMode;
 use pypi_types::Requirement;
 use uv_auth::store_credentials_from_url;
 use uv_cache::Cache;
-use uv_client::{Connectivity, FlatIndexClient, RegistryClientBuilder};
+use uv_client::{BaseClientBuilder, Connectivity, FlatIndexClient, RegistryClientBuilder};
 use uv_configuration::{Concurrency, KeyringProviderType, PreviewMode};
 use uv_configuration::{ConfigSettings, IndexStrategy, NoBinary, NoBuild, SetupPyStrategy};
 use uv_dispatch::BuildDispatch;
@@ -119,8 +119,19 @@ async fn venv_impl(
     cache: &Cache,
     printer: Printer,
 ) -> miette::Result<ExitStatus> {
+    let client_builder = BaseClientBuilder::default()
+        .connectivity(connectivity)
+        .native_tls(native_tls);
+
     // Locate the Python interpreter to use in the environment
-    let interpreter = Toolchain::find(python_request, SystemPython::Required, preview, cache)
+    let interpreter = Toolchain::find_or_fetch(
+        python_request,
+        SystemPython::Required,
+        preview,
+        client_builder,
+        cache,
+    )
+    .await
     .into_diagnostic()?
     .into_interpreter();