From 7ef3dfeadb4085aed717d70619fc3d0e2bdb70b8 Mon Sep 17 00:00:00 2001 From: Charlie Marsh Date: Sat, 10 Aug 2024 21:07:11 -0400 Subject: [PATCH] Refactor --- crates/uv-scripts/src/lib.rs | 12 +-- crates/uv/src/commands/project/add.rs | 122 +++++++++++++---------- crates/uv/src/commands/project/remove.rs | 35 ++++--- 3 files changed, 92 insertions(+), 77 deletions(-) diff --git a/crates/uv-scripts/src/lib.rs b/crates/uv-scripts/src/lib.rs index 9cdd12c84f59..80c9e8adf764 100644 --- a/crates/uv-scripts/src/lib.rs +++ b/crates/uv-scripts/src/lib.rs @@ -211,13 +211,10 @@ impl ScriptTag { return Ok(None); } + // Extract the preceding content. + let prelude = std::str::from_utf8(&contents[..index])?; + // Decode as UTF-8. - let prelude = if index != 0 { - std::str::from_utf8(&contents[..index])? - } else { - "" - } - .to_string(); let contents = &contents[index..]; let contents = std::str::from_utf8(contents)?; @@ -235,6 +232,7 @@ impl ScriptTag { // > consists of only a single #). let mut toml = vec![]; + // Extract the content that follows the metadata block. let mut python_script = vec![]; while let Some(line) = lines.next() { @@ -260,6 +258,7 @@ impl ScriptTag { toml.push(line); } + // Find the closing `# ///`. The precedence is such that we need to identify the _last_ such // line. // @@ -293,6 +292,7 @@ impl ScriptTag { toml.truncate(index - 1); // Join the lines into a single string. + let prelude = prelude.to_string(); let metadata = toml.join("\n") + "\n"; let script = python_script.join("\n") + "\n"; diff --git a/crates/uv/src/commands/project/add.rs b/crates/uv/src/commands/project/add.rs index d81246818d39..10e02dfe09d9 100644 --- a/crates/uv/src/commands/project/add.rs +++ b/crates/uv/src/commands/project/add.rs @@ -3,9 +3,10 @@ use std::path::PathBuf; use anyhow::{Context, Result}; use owo_colors::OwoColorize; -use pep508_rs::{ExtraName, Requirement, VersionOrUrl}; use rustc_hash::{FxBuildHasher, FxHashMap}; use tracing::debug; + +use pep508_rs::{ExtraName, Requirement, VersionOrUrl}; use uv_auth::store_credentials_from_url; use uv_cache::Cache; use uv_client::{BaseClientBuilder, Connectivity, FlatIndexClient, RegistryClientBuilder}; @@ -17,8 +18,8 @@ use uv_distribution::DistributionDatabase; use uv_fs::CWD; use uv_normalize::PackageName; use uv_python::{ - request_from_version_file, EnvironmentPreference, PythonDownloads, PythonInstallation, - PythonPreference, PythonRequest, VersionRequest, + request_from_version_file, EnvironmentPreference, Interpreter, PythonDownloads, + PythonEnvironment, PythonInstallation, PythonPreference, PythonRequest, VersionRequest, }; use uv_requirements::{NamedRequirementsResolver, RequirementsSource, RequirementsSpecification}; use uv_resolver::{FlatIndex, RequiresPython}; @@ -38,12 +39,6 @@ use crate::commands::{pip, project, ExitStatus, SharedState}; use crate::printer::Printer; use crate::settings::ResolverInstallerSettings; -/// Represents the destination where dependencies are added, either to a project or a script. -enum DependencyDestination { - Project(VirtualProject), - Script(Pep723Script), -} - /// Add one or more packages to the project requirements. #[allow(clippy::fn_params_excessive_bools)] pub(crate) async fn add( @@ -75,8 +70,9 @@ pub(crate) async fn add( warn_user_once!("`uv add` is experimental and may change without warning"); } - let download_reporter = PythonDownloadReporter::single(printer); - let (dependency_destination, venv) = if let Some(script) = script { + let reporter = PythonDownloadReporter::single(printer); + + let target = if let Some(script) = script { // If we found a PEP 723 script and the user provided a project-only setting, warn. if !extras.is_empty() { warn_user_once!("Extras are not supported for Python scripts with inline metadata"); @@ -101,25 +97,27 @@ pub(crate) async fn add( "`--no_sync` is a no-op for Python scripts with inline metadata, which always run in isolation" ); } + let client_builder = BaseClientBuilder::new() .connectivity(connectivity) .native_tls(native_tls); + // If we found a script, add to the existing metadata. Otherwise, create a new inline + // metadata tag. let script = if let Some(script) = Pep723Script::read(&script).await? { script } else { - // As no metadata was found in the script, we will create a default metadata table for the script. - - // (1) Explicit request from user let python_request = if let Some(request) = python.as_deref() { + // (1) Explicit request from user PythonRequest::parse(request) - // (2) Request from `.python-version` } else if let Some(request) = request_from_version_file(&CWD).await? { + // (2) Request from `.python-version` request } else { + // (3) Assume any Python version PythonRequest::Any }; - let reporter = PythonDownloadReporter::single(printer); + let interpreter = PythonInstallation::find_or_download( Some(python_request), EnvironmentPreference::Any, @@ -131,18 +129,20 @@ pub(crate) async fn add( ) .await? .into_interpreter(); + let requires_python = RequiresPython::greater_than_equal_version(&interpreter.python_minor_version()); Pep723Script::create(&script, requires_python.specifiers()).await? }; - // (1) Explicit request from user + let python_request = if let Some(request) = python.as_deref() { + // (1) Explicit request from user Some(PythonRequest::parse(request)) - // (2) Request from `.python-version` } else if let Some(request) = request_from_version_file(&CWD).await? { + // (2) Request from `.python-version` Some(request) - // (3) `Requires-Python` in `pyproject.toml` } else { + // (3) `Requires-Python` in `pyproject.toml` script .metadata .requires_python @@ -159,23 +159,12 @@ pub(crate) async fn add( python_downloads, &client_builder, cache, - Some(&download_reporter), + Some(&reporter), ) .await? .into_interpreter(); - // Create a virtual environment. - let temp_dir = cache.environment()?; - let venv = uv_virtualenv::create_venv( - temp_dir.path(), - interpreter, - uv_virtualenv::Prompt::None, - false, - false, - false, - )?; - - (DependencyDestination::Script(script), venv) + Target::Script(script, Box::new(interpreter)) } else { // Find the project in the workspace. let project = if let Some(package) = package { @@ -214,7 +203,8 @@ pub(crate) async fn add( printer, ) .await?; - (DependencyDestination::Project(project), venv) + + Target::Project(project, venv) }; let client_builder = BaseClientBuilder::new() @@ -236,7 +226,7 @@ pub(crate) async fn add( // Determine the environment for the resolution. let (tags, markers) = - resolution_environment(python_version, python_platform, venv.interpreter())?; + resolution_environment(python_version, python_platform, target.interpreter())?; // Add all authenticated sources to the cache. for url in settings.index_locations.urls() { @@ -248,7 +238,7 @@ pub(crate) async fn add( .index_urls(settings.index_locations.index_urls()) .index_strategy(settings.index_strategy) .markers(&markers) - .platform(venv.interpreter().platform()) + .platform(target.interpreter().platform()) .build(); // Initialize any shared state. @@ -269,7 +259,7 @@ pub(crate) async fn add( &client, cache, &build_constraints, - venv.interpreter(), + target.interpreter(), &settings.index_locations, &flat_index, &state.index, @@ -298,15 +288,15 @@ pub(crate) async fn add( .resolve() .await?; - // Add the requirements to the `pyproject.toml`. - let mut toml = match &dependency_destination { - DependencyDestination::Script(script) => { + // Add the requirements to the `pyproject.toml` or script. + let mut toml = match &target { + Target::Script(script, _) => { PyProjectTomlMut::from_toml(&script.metadata.raw, DependencyTarget::Script) } - DependencyDestination::Project(project) => { - let raw = project.pyproject_toml().raw.clone(); - PyProjectTomlMut::from_toml(&raw, DependencyTarget::PyProjectToml) - } + Target::Project(project, _) => PyProjectTomlMut::from_toml( + &project.pyproject_toml().raw, + DependencyTarget::PyProjectToml, + ), }?; let mut edits = Vec::with_capacity(requirements.len()); for mut requirement in requirements { @@ -315,11 +305,11 @@ pub(crate) async fn add( requirement.extras.sort_unstable(); requirement.extras.dedup(); - let (requirement, source) = match dependency_destination { - DependencyDestination::Script(_) | DependencyDestination::Project(_) if raw_sources => { + let (requirement, source) = match target { + Target::Script(_, _) | Target::Project(_, _) if raw_sources => { (pep508_rs::Requirement::from(requirement), None) } - DependencyDestination::Script(_) => resolve_and_process_requirement( + Target::Script(_, _) => resolve_requirement( requirement, false, editable, @@ -327,12 +317,12 @@ pub(crate) async fn add( tag.clone(), branch.clone(), )?, - DependencyDestination::Project(ref project) => { + Target::Project(ref project, _) => { let workspace = project .workspace() .packages() .contains_key(&requirement.name); - resolve_and_process_requirement( + resolve_requirement( requirement, workspace, editable, @@ -342,6 +332,7 @@ pub(crate) async fn add( )? } }; + // Update the `pyproject.toml`. let edit = match dependency_type { DependencyType::Production => toml.add_dependency(&requirement, source.as_ref())?, @@ -360,10 +351,11 @@ pub(crate) async fn add( }); } - // Save the modified `pyproject.toml`. let content = toml.to_string(); - let modified = match &dependency_destination { - DependencyDestination::Script(script) => { + + // Save the modified `pyproject.toml` or script. + let modified = match &target { + Target::Script(script, _) => { if content == script.metadata.raw { debug!("No changes to dependencies; skipping update"); false @@ -372,7 +364,7 @@ pub(crate) async fn add( true } } - DependencyDestination::Project(project) => { + Target::Project(project, _) => { if content == *project.pyproject_toml().raw { debug!("No changes to dependencies; skipping update"); false @@ -390,11 +382,12 @@ pub(crate) async fn add( } // If `--script`, exit early. There's no reason to lock and sync. - let DependencyDestination::Project(project) = dependency_destination else { + let Target::Project(project, venv) = target else { return Ok(ExitStatus::Success); }; let existing = project.pyproject_toml(); + // Update the `pypackage.toml` in-memory. let project = project .clone() @@ -560,14 +553,14 @@ pub(crate) async fn add( } /// Resolves the source for a requirement and processes it into a PEP 508 compliant format. -fn resolve_and_process_requirement( +fn resolve_requirement( requirement: pypi_types::Requirement, workspace: bool, editable: Option, rev: Option, tag: Option, branch: Option, -) -> Result<(pep508_rs::Requirement, Option), anyhow::Error> { +) -> Result<(Requirement, Option), anyhow::Error> { let result = Source::from_requirement( &requirement.name, requirement.source.clone(), @@ -596,6 +589,25 @@ fn resolve_and_process_requirement( Ok((processed_requirement, source)) } +/// Represents the destination where dependencies are added, either to a project or a script. +#[derive(Debug)] +enum Target { + /// A PEP 723 script, with inline metadata. + Script(Pep723Script, Box), + /// A project with a `pyproject.toml`. + Project(VirtualProject, PythonEnvironment), +} + +impl Target { + /// Returns the [`Interpreter`] for the target. + fn interpreter(&self) -> &Interpreter { + match self { + Self::Script(_, interpreter) => interpreter, + Self::Project(_, venv) => venv.interpreter(), + } + } +} + #[derive(Debug, Clone)] struct DependencyEdit<'a> { dependency_type: &'a DependencyType, diff --git a/crates/uv/src/commands/project/remove.rs b/crates/uv/src/commands/project/remove.rs index 9ed08c09615c..5342fc2b0504 100644 --- a/crates/uv/src/commands/project/remove.rs +++ b/crates/uv/src/commands/project/remove.rs @@ -18,12 +18,6 @@ use crate::commands::{project, ExitStatus, SharedState}; use crate::printer::Printer; use crate::settings::ResolverInstallerSettings; -/// Represents the destination where dependencies are added, either to a project or a script. -enum DependencyDestination { - Project(VirtualProject), - Script(Pep723Script), -} - /// Remove one or more packages from the project requirements. #[allow(clippy::fn_params_excessive_bools)] pub(crate) async fn remove( @@ -49,7 +43,7 @@ pub(crate) async fn remove( warn_user_once!("`uv remove` is experimental and may change without warning"); } - let dependency_destination = if let Some(script) = script { + let target = if let Some(script) = script { // If we found a PEP 723 script and the user provided a project-only setting, warn. if package.is_some() { warn_user_once!( @@ -71,7 +65,7 @@ pub(crate) async fn remove( "`--no_sync` is a no-op for Python scripts with inline metadata, which always run in isolation" ); } - DependencyDestination::Script(script) + Target::Script(script) } else { // Find the project in the workspace. let project = if let Some(package) = package { @@ -85,14 +79,14 @@ pub(crate) async fn remove( VirtualProject::discover(&CWD, &DiscoveryOptions::default()).await? }; - DependencyDestination::Project(project) + Target::Project(project) }; - let mut toml = match &dependency_destination { - DependencyDestination::Script(script) => { + let mut toml = match &target { + Target::Script(script) => { PyProjectTomlMut::from_toml(&script.metadata.raw, DependencyTarget::Script) } - DependencyDestination::Project(project) => PyProjectTomlMut::from_toml( + Target::Project(project) => PyProjectTomlMut::from_toml( project.pyproject_toml().raw.as_ref(), DependencyTarget::PyProjectToml, ), @@ -131,11 +125,11 @@ pub(crate) async fn remove( } // Save the modified dependencies. - match &dependency_destination { - DependencyDestination::Script(script) => { + match &target { + Target::Script(script) => { script.replace_metadata(&toml.to_string()).await?; } - DependencyDestination::Project(project) => { + Target::Project(project) => { let pyproject_path = project.root().join("pyproject.toml"); fs_err::write(pyproject_path, toml.to_string())?; } @@ -148,7 +142,7 @@ pub(crate) async fn remove( } // If `--script`, exit early. There's no reason to lock and sync. - let DependencyDestination::Project(project) = dependency_destination else { + let Target::Project(project) = target else { return Ok(ExitStatus::Success); }; @@ -216,6 +210,15 @@ pub(crate) async fn remove( Ok(ExitStatus::Success) } +/// Represents the destination where dependencies are added, either to a project or a script. +#[derive(Debug)] +enum Target { + /// A PEP 723 script, with inline metadata. + Project(VirtualProject), + /// A project with a `pyproject.toml`. + Script(Pep723Script), +} + /// Emit a warning if a dependency with the given name is present as any dependency type. /// /// This is useful when a dependency of the user-specified type was not found, but it may be present