From fb3f365d10f06d96f0945c854796acfa324fe55a Mon Sep 17 00:00:00 2001 From: Charlie Marsh Date: Fri, 15 Nov 2024 23:03:39 -0500 Subject: [PATCH] Support overrides and constraints in PEP 723 scripts (#9162) ## Summary Closes https://github.com/astral-sh/uv/issues/9141. --- crates/uv-requirements/src/specification.rs | 24 +++++++ crates/uv-scripts/src/lib.rs | 10 ++- crates/uv/src/commands/project/run.rs | 44 +++++++++++- crates/uv/src/lib.rs | 4 +- crates/uv/tests/it/run.rs | 78 +++++++++++++++++++++ 5 files changed, 154 insertions(+), 6 deletions(-) diff --git a/crates/uv-requirements/src/specification.rs b/crates/uv-requirements/src/specification.rs index bb491cb19..8b5f0a3af 100644 --- a/crates/uv-requirements/src/specification.rs +++ b/crates/uv-requirements/src/specification.rs @@ -338,6 +338,30 @@ impl RequirementsSpecification { } } + /// Initialize a [`RequirementsSpecification`] from a list of [`Requirement`], including + /// constraints and overrides. + pub fn from_constraints( + requirements: Vec, + constraints: Vec, + overrides: Vec, + ) -> Self { + Self { + requirements: requirements + .into_iter() + .map(UnresolvedRequirementSpecification::from) + .collect(), + constraints: constraints + .into_iter() + .map(NameRequirementSpecification::from) + .collect(), + overrides: overrides + .into_iter() + .map(UnresolvedRequirementSpecification::from) + .collect(), + ..Self::default() + } + } + /// Return true if the specification does not include any requirements to install. pub fn is_empty(&self) -> bool { self.requirements.is_empty() && self.source_trees.is_empty() && self.overrides.is_empty() diff --git a/crates/uv-scripts/src/lib.rs b/crates/uv-scripts/src/lib.rs index 488439206..fb9118a2a 100644 --- a/crates/uv-scripts/src/lib.rs +++ b/crates/uv-scripts/src/lib.rs @@ -1,11 +1,13 @@ -use memchr::memmem::Finder; -use serde::Deserialize; use std::collections::BTreeMap; use std::io; use std::path::{Path, PathBuf}; use std::str::FromStr; use std::sync::LazyLock; + +use memchr::memmem::Finder; +use serde::Deserialize; use thiserror::Error; + use uv_distribution_types::Index; use uv_pep440::VersionSpecifiers; use uv_pep508::PackageName; @@ -264,12 +266,14 @@ pub struct Tool { } #[derive(Debug, Deserialize, Clone)] -#[serde(deny_unknown_fields)] +#[serde(deny_unknown_fields, rename_all = "kebab-case")] pub struct ToolUv { #[serde(flatten)] pub globals: GlobalOptions, #[serde(flatten)] pub top_level: ResolverInstallerOptions, + pub override_dependencies: Option>>, + pub constraint_dependencies: Option>>, pub sources: Option>, pub indexes: Option>, } diff --git a/crates/uv/src/commands/project/run.rs b/crates/uv/src/commands/project/run.rs index 7002c59d7..f5bb2b627 100644 --- a/crates/uv/src/commands/project/run.rs +++ b/crates/uv/src/commands/project/run.rs @@ -274,7 +274,49 @@ pub(crate) async fn run( .map_ok(LoweredRequirement::into_inner) }) .collect::>()?; - let spec = RequirementsSpecification::from_requirements(requirements); + let constraints = script + .tool + .as_ref() + .and_then(|tool| tool.uv.as_ref()) + .and_then(|uv| uv.constraint_dependencies.as_ref()) + .into_iter() + .flatten() + .cloned() + .flat_map(|requirement| { + LoweredRequirement::from_non_workspace_requirement( + requirement, + script_dir.as_ref(), + script_sources, + script_indexes, + &settings.index_locations, + LowerBound::Allow, + ) + .map_ok(LoweredRequirement::into_inner) + }) + .collect::, _>>()?; + let overrides = script + .tool + .as_ref() + .and_then(|tool| tool.uv.as_ref()) + .and_then(|uv| uv.override_dependencies.as_ref()) + .into_iter() + .flatten() + .cloned() + .flat_map(|requirement| { + LoweredRequirement::from_non_workspace_requirement( + requirement, + script_dir.as_ref(), + script_sources, + script_indexes, + &settings.index_locations, + LowerBound::Allow, + ) + .map_ok(LoweredRequirement::into_inner) + }) + .collect::, _>>()?; + + let spec = + RequirementsSpecification::from_constraints(requirements, constraints, overrides); let result = CachedEnvironment::get_or_create( EnvironmentSpecification::from(spec), interpreter, diff --git a/crates/uv/src/lib.rs b/crates/uv/src/lib.rs index d59ac90c5..9da6b293f 100644 --- a/crates/uv/src/lib.rs +++ b/crates/uv/src/lib.rs @@ -1497,7 +1497,7 @@ async fn run_project( Pep723Item::Remote(_) => unreachable!("`uv remove` does not support remote files"), }); - commands::remove( + Box::pin(commands::remove( project_dir, args.locked, args.frozen, @@ -1518,7 +1518,7 @@ async fn run_project( no_config, &cache, printer, - ) + )) .await } ProjectCommand::Tree(args) => { diff --git a/crates/uv/tests/it/run.rs b/crates/uv/tests/it/run.rs index 82502c437..2a12d9846 100644 --- a/crates/uv/tests/it/run.rs +++ b/crates/uv/tests/it/run.rs @@ -654,6 +654,84 @@ fn run_pep723_script_metadata() -> Result<()> { Ok(()) } +/// Run a PEP 723-compatible script with `tool.uv` constraints. +#[test] +fn run_pep723_script_constraints() -> Result<()> { + let context = TestContext::new("3.12"); + + let test_script = context.temp_dir.child("main.py"); + test_script.write_str(indoc! { r#" + # /// script + # requires-python = ">=3.11" + # dependencies = [ + # "anyio>=3", + # ] + # + # [tool.uv] + # constraint-dependencies = ["idna<=3"] + # /// + + import anyio + "# + })?; + + uv_snapshot!(context.filters(), context.run().arg("main.py"), @r###" + success: true + exit_code: 0 + ----- stdout ----- + + ----- stderr ----- + Reading inline script metadata from `main.py` + Resolved 3 packages in [TIME] + Prepared 3 packages in [TIME] + Installed 3 packages in [TIME] + + anyio==4.3.0 + + idna==3.0 + + sniffio==1.3.1 + "###); + + Ok(()) +} + +/// Run a PEP 723-compatible script with `tool.uv` overrides. +#[test] +fn run_pep723_script_overrides() -> Result<()> { + let context = TestContext::new("3.12"); + + let test_script = context.temp_dir.child("main.py"); + test_script.write_str(indoc! { r#" + # /// script + # requires-python = ">=3.11" + # dependencies = [ + # "anyio>=3", + # ] + # + # [tool.uv] + # override-dependencies = ["idna<=2"] + # /// + + import anyio + "# + })?; + + uv_snapshot!(context.filters(), context.run().arg("main.py"), @r###" + success: true + exit_code: 0 + ----- stdout ----- + + ----- stderr ----- + Reading inline script metadata from `main.py` + Resolved 3 packages in [TIME] + Prepared 3 packages in [TIME] + Installed 3 packages in [TIME] + + anyio==4.3.0 + + idna==2.0 + + sniffio==1.3.1 + "###); + + Ok(()) +} + /// With `managed = false`, we should avoid installing the project itself. #[test] fn run_managed_false() -> Result<()> {