From 0c22252934b93552852cff6fdc00c92982663aaf Mon Sep 17 00:00:00 2001 From: Zanie Blue Date: Fri, 7 Jun 2024 11:30:39 -0500 Subject: [PATCH] Implement `Toolchain::find_or_fetch` and use in `uv venv --preview` --- 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 | 86 +++++++++++++++++++++------- crates/uv-toolchain/src/lib.rs | 6 ++ crates/uv-toolchain/src/managed.rs | 48 +++++++++++++--- crates/uv-toolchain/src/toolchain.rs | 58 +++++++++++++++++++ crates/uv/src/commands/venv.rs | 19 ++++-- 9 files changed, 200 insertions(+), 83 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 7a77a40db66bb..e191984d5c5b3 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4592,7 +4592,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 dc4ab0099f99d..564a095df8aa2 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 5abd52ab88dc4..3303a323b6b86 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 b8d59c2f8cad3..0b411953ed3bc 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 57fd45d41416c..09efb46611d88 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)] 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,32 @@ impl PythonDownloadRequest { self } + pub fn from_request(request: ToolchainRequest) -> Result { + 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 { 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 { // 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 diff --git a/crates/uv-toolchain/src/lib.rs b/crates/uv-toolchain/src/lib.rs index 63d4e0d41ad35..f718797b4c3ef 100644 --- a/crates/uv-toolchain/src/lib.rs +++ b/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), } diff --git a/crates/uv-toolchain/src/managed.rs b/crates/uv-toolchain/src/managed.rs index 3fe11efbf7d77..a57eafc50c98d 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,7 +54,7 @@ pub struct InstalledToolchains { impl InstalledToolchains { /// A directory for installed toolchains at `root`. - pub fn from_path(root: impl Into) -> Result { + pub fn from_path(root: impl Into) -> Result { Ok(Self { root: root.into() }) } @@ -30,7 +62,7 @@ impl InstalledToolchains { /// 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) } else { @@ -39,14 +71,14 @@ impl InstalledToolchains { } /// Create a temporary installed toolchain directory. - pub fn temp() -> Result { + pub fn temp() -> Result { 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 +92,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 a277150eed362..9fca8fb9cd8f2 100644 --- a/crates/uv-toolchain/src/toolchain.rs +++ b/crates/uv-toolchain/src/toolchain.rs @@ -1,8 +1,12 @@ +use tracing::{debug, info}; +use uv_client::BaseClientBuilder; use uv_configuration::PreviewMode; use uv_cache::Cache; use crate::discovery::{SystemPython, ToolchainRequest, ToolchainSources}; +use crate::downloads::{DownloadResult, PythonDownload, PythonDownloadRequest}; +use crate::managed::{InstalledToolchain, InstalledToolchains}; use crate::{ find_best_toolchain, find_default_toolchain, find_toolchain, Error, Interpreter, ToolchainSource, @@ -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 33e256fb45617..38c0f4a54dff9 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() {