From 45df889fe44b656197154d3d94cffec776d0c8d4 Mon Sep 17 00:00:00 2001 From: Zanie Blue Date: Mon, 10 Jun 2024 10:10:45 -0400 Subject: [PATCH] Implement `Toolchain::find_or_fetch` and use in `uv venv --preview` (#4138) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 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. --- Cargo.lock | 1 - crates/uv-dev/Cargo.toml | 1 - crates/uv-dev/src/fetch_python.rs | 53 ++------------- crates/uv-toolchain/src/discovery.rs | 11 ++++ crates/uv-toolchain/src/downloads.rs | 97 +++++++++++++++++++++------- crates/uv-toolchain/src/lib.rs | 6 ++ crates/uv-toolchain/src/managed.rs | 60 +++++++++++++---- crates/uv-toolchain/src/toolchain.rs | 58 +++++++++++++++++ crates/uv/src/commands/venv.rs | 19 ++++-- 9 files changed, 219 insertions(+), 87 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 9d728d525776..4f146c5ad535 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4593,7 +4593,6 @@ dependencies = [ "fs-err", "futures", "install-wheel-rs", - "itertools 0.13.0", "mimalloc", "owo-colors", "pep508_rs", diff --git a/crates/uv-dev/Cargo.toml b/crates/uv-dev/Cargo.toml index dc4ab0099f99..564a095df8aa 100644 --- a/crates/uv-dev/Cargo.toml +++ b/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" } diff --git a/crates/uv-dev/src/fetch_python.rs b/crates/uv-dev/src/fetch_python.rs index 5abd52ab88dc..3303a323b6b8 100644 --- a/crates/uv-dev/src/fetch_python.rs +++ b/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::, Error>>()?; let downloads = requests .iter() - .map(|request| match PythonDownload::from_request(request) { - Some(download) => download, - None => panic!("No download found for request {request:?}"), - }) - .collect::>(); + .map(PythonDownload::from_request) + .collect::, 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 = 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(()) diff --git a/crates/uv-toolchain/src/discovery.rs b/crates/uv-toolchain/src/discovery.rs index 24fa90a56bc1..fbf68894787a 100644 --- a/crates/uv-toolchain/src/discovery.rs +++ b/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 { diff --git a/crates/uv-toolchain/src/downloads.rs b/crates/uv-toolchain/src/downloads.rs index 57fd45d41416..d10d07872b96 100644 --- a/crates/uv-toolchain/src/downloads.rs +++ b/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, Default)] pub struct PythonDownloadRequest { - version: Option, + version: Option, implementation: Option, arch: Option, os: Option, @@ -77,7 +83,7 @@ pub struct PythonDownloadRequest { impl PythonDownloadRequest { pub fn new( - version: Option, + version: Option, implementation: Option, arch: Option, os: Option, @@ -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 { + 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 { 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 { // 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 diff --git a/crates/uv-toolchain/src/lib.rs b/crates/uv-toolchain/src/lib.rs index 6fcbed5cb9ad..7927a165fce1 100644 --- a/crates/uv-toolchain/src/lib.rs +++ b/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), } diff --git a/crates/uv-toolchain/src/managed.rs b/crates/uv-toolchain/src/managed.rs index 3fe11efbf7d7..db8d6e9feeb3 100644 --- a/crates/uv-toolchain/src/managed.rs +++ b/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 -pub use crate::downloads::Error; +use crate::downloads::Error as DownloadError; +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) -> Result { - Ok(Self { root: root.into() }) + fn from_path(root: impl Into) -> Self { + 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 { + pub fn from_settings() -> Result { 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::from_path(StateStore::temp()?.bucket(StateBucket::Toolchains)) + pub fn temp() -> Result { + Ok(Self::from_path( + StateStore::temp()?.bucket(StateBucket::Toolchains), + )) } /// Initialize the installed toolchain directory. /// /// Ensures the directory is created. - pub fn init(self) -> Result { + pub fn init(self) -> Result { 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) diff --git a/crates/uv-toolchain/src/toolchain.rs b/crates/uv-toolchain/src/toolchain.rs index 65809c0c3f43..e3b5d300b444 100644 --- a/crates/uv-toolchain/src/toolchain.rs +++ b/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 { + // 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 { + 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 { diff --git a/crates/uv/src/commands/venv.rs b/crates/uv/src/commands/venv.rs index e6ea44d939b7..95b9ce871d3d 100644 --- a/crates/uv/src/commands/venv.rs +++ b/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,10 +119,21 @@ async fn venv_impl( cache: &Cache, printer: Printer, ) -> miette::Result { + 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) - .into_diagnostic()? - .into_interpreter(); + let interpreter = Toolchain::find_or_fetch( + python_request, + SystemPython::Required, + preview, + client_builder, + cache, + ) + .await + .into_diagnostic()? + .into_interpreter(); // Add all authenticated sources to the cache. for url in index_locations.urls() {