From 5b428eaabac88c9fe21fd38c7c1811d8dfe51d96 Mon Sep 17 00:00:00 2001 From: Ahmed Ilyas Date: Tue, 2 Jul 2024 16:20:23 +0200 Subject: [PATCH] Support PEP 723 scripts in `uv add` and `uv remove` --- crates/uv-cli/src/lib.rs | 7 + crates/uv-scripts/src/lib.rs | 177 +++++++++++----- crates/uv-workspace/src/pyproject_mut.rs | 77 ++++--- crates/uv-workspace/src/workspace.rs | 2 +- crates/uv/src/commands/project/add.rs | 257 +++++++++++++++-------- crates/uv/src/commands/project/init.rs | 7 +- crates/uv/src/commands/project/remove.rs | 79 +++++-- crates/uv/src/lib.rs | 15 ++ crates/uv/src/settings.rs | 6 + crates/uv/tests/edit.rs | 124 +++++++++++ docs/reference/cli.md | 4 + 11 files changed, 560 insertions(+), 195 deletions(-) diff --git a/crates/uv-cli/src/lib.rs b/crates/uv-cli/src/lib.rs index 7e5a7b5c3462b..77360190fd2da 100644 --- a/crates/uv-cli/src/lib.rs +++ b/crates/uv-cli/src/lib.rs @@ -2462,6 +2462,10 @@ pub struct AddArgs { help_heading = "Python options" )] pub python: Option, + + /// Specifies the Python script where the dependency will be added. + #[arg(long)] + pub script: Option, } #[derive(Args)] @@ -2509,6 +2513,9 @@ pub struct RemoveArgs { #[arg(long, conflicts_with = "isolated")] pub package: Option, + /// Specifies the Python script where the dependency will be removed. + #[arg(long)] + pub script: Option, /// The Python interpreter to use for resolving and syncing. /// /// See `uv help python` for details on Python discovery and supported diff --git a/crates/uv-scripts/src/lib.rs b/crates/uv-scripts/src/lib.rs index c2f66ac26ea42..2359829fa877c 100644 --- a/crates/uv-scripts/src/lib.rs +++ b/crates/uv-scripts/src/lib.rs @@ -19,6 +19,7 @@ static FINDER: LazyLock = LazyLock::new(|| Finder::new(b"# /// script")) pub struct Pep723Script { pub path: PathBuf, pub metadata: Pep723Metadata, + pub data: String, } impl Pep723Script { @@ -26,12 +27,35 @@ impl Pep723Script { /// /// See: pub async fn read(file: impl AsRef) -> Result, Pep723Error> { - let metadata = Pep723Metadata::read(&file).await?; - Ok(metadata.map(|metadata| Self { + let contents = match fs_err::tokio::read(&file).await { + Ok(contents) => contents, + Err(err) if err.kind() == io::ErrorKind::NotFound => return Ok(None), + Err(err) => return Err(err.into()), + }; + + // Extract the `script` tag. + let Some((metadata, data)) = extract_script_tag(&contents)? else { + return Ok(None); + }; + + // Parse the metadata. + let metadata = Pep723Metadata::from_string(metadata)?; + + Ok(Some(Self { path: file.as_ref().to_path_buf(), metadata, + data, })) } + + /// Replace the existing metadata in the file with new metadata and write the updated content. + pub async fn replace_metadata(&self, new_metadata: &str) -> Result<(), Pep723Error> { + let new_content = format!("{}{}", serialize_metadata(new_metadata), self.data); + + fs_err::tokio::write(&self.path, new_content) + .await + .map_err(std::convert::Into::into) + } } /// PEP 723 metadata as parsed from a `script` comment block. @@ -43,28 +67,16 @@ pub struct Pep723Metadata { pub dependencies: Option>>, pub requires_python: Option, pub tool: Option, + /// The raw unserialized document. + #[serde(skip)] + pub raw: String, } impl Pep723Metadata { - /// Read the PEP 723 `script` metadata from a Python file, if it exists. - /// - /// See: - pub async fn read(file: impl AsRef) -> Result, Pep723Error> { - let contents = match fs_err::tokio::read(file).await { - Ok(contents) => contents, - Err(err) if err.kind() == io::ErrorKind::NotFound => return Ok(None), - Err(err) => return Err(err.into()), - }; - - // Extract the `script` tag. - let Some(contents) = extract_script_tag(&contents)? else { - return Ok(None); - }; - - // Parse the metadata. - let metadata = toml::from_str(&contents)?; - - Ok(Some(metadata)) + /// Parse `Pep723Metadata` from a raw TOML string. + pub fn from_string(raw: String) -> Result { + let metadata = toml::from_str(&raw)?; + Ok(Pep723Metadata { raw, ..metadata }) } } @@ -94,34 +106,11 @@ pub enum Pep723Error { Toml(#[from] toml::de::Error), } -/// Read the PEP 723 `script` metadata from a Python file, if it exists. -/// -/// See: -pub async fn read_pep723_metadata( - file: impl AsRef, -) -> Result, Pep723Error> { - let contents = match fs_err::tokio::read(file).await { - Ok(contents) => contents, - Err(err) if err.kind() == io::ErrorKind::NotFound => return Ok(None), - Err(err) => return Err(err.into()), - }; - - // Extract the `script` tag. - let Some(contents) = extract_script_tag(&contents)? else { - return Ok(None); - }; - - // Parse the metadata. - let metadata = toml::from_str(&contents)?; - - Ok(Some(metadata)) -} - /// Given the contents of a Python file, extract the `script` metadata block, with leading comment -/// hashes removed. +/// hashes removed and the python script. /// /// See: -fn extract_script_tag(contents: &[u8]) -> Result, Pep723Error> { +fn extract_script_tag(contents: &[u8]) -> Result, Pep723Error> { // Identify the opening pragma. let Some(index) = FINDER.find(contents) else { return Ok(None); @@ -149,9 +138,14 @@ fn extract_script_tag(contents: &[u8]) -> Result, Pep723Error> { // > second character is a space, otherwise just the first character (which means the line // > consists of only a single #). let mut toml = vec![]; - for line in lines { + + let mut python_script = vec![]; + + while let Some(line) = lines.next() { // Remove the leading `#`. let Some(line) = line.strip_prefix('#') else { + python_script.push(line); + python_script.extend(lines); break; }; @@ -163,11 +157,13 @@ fn extract_script_tag(contents: &[u8]) -> Result, Pep723Error> { // Otherwise, the line _must_ start with ` `. let Some(line) = line.strip_prefix(' ') else { + python_script.push(line); + python_script.extend(lines); break; }; + toml.push(line); } - // Find the closing `# ///`. The precedence is such that we need to identify the _last_ such // line. // @@ -202,12 +198,36 @@ fn extract_script_tag(contents: &[u8]) -> Result, Pep723Error> { // Join the lines into a single string. let toml = toml.join("\n") + "\n"; + let python_script = python_script.join("\n") + "\n"; + + Ok(Some((toml, python_script))) +} + +/// Formats the provided metadata by prefixing each line with `#` and wrapping it with script markers. +fn serialize_metadata(metadata: &str) -> String { + let mut output = String::with_capacity(metadata.len() + 2); + + output.push_str("# /// script\n"); - Ok(Some(toml)) + for line in metadata.lines() { + if line.is_empty() { + output.push('\n'); + } else { + output.push_str("# "); + output.push_str(line); + output.push('\n'); + } + } + + output.push_str("# ///\n"); + + output } #[cfg(test)] mod tests { + use crate::serialize_metadata; + #[test] fn missing_space() { let contents = indoc::indoc! {r" @@ -269,9 +289,15 @@ mod tests { # 'rich', # ] # /// + + import requests + from rich.pretty import pprint + + resp = requests.get('https://peps.python.org/api/peps.json') + data = resp.json() "}; - let expected = indoc::indoc! {r" + let expected_metadata = indoc::indoc! {r" requires-python = '>=3.11' dependencies = [ 'requests<3', @@ -279,11 +305,21 @@ mod tests { ] "}; + let expected_data = indoc::indoc! {r" + + import requests + from rich.pretty import pprint + + resp = requests.get('https://peps.python.org/api/peps.json') + data = resp.json() + "}; + let actual = super::extract_script_tag(contents.as_bytes()) .unwrap() .unwrap(); - assert_eq!(actual, expected); + assert_eq!(actual.0, expected_metadata); + assert_eq!(actual.1, expected_data); } #[test] @@ -312,7 +348,8 @@ mod tests { let actual = super::extract_script_tag(contents.as_bytes()) .unwrap() - .unwrap(); + .unwrap() + .0; assert_eq!(actual, expected); } @@ -341,8 +378,42 @@ mod tests { let actual = super::extract_script_tag(contents.as_bytes()) .unwrap() - .unwrap(); + .unwrap() + .0; assert_eq!(actual, expected); } + + #[test] + fn test_serialize_metadata_formatting() { + let metadata = indoc::indoc! {r" + requires-python = '>=3.11' + dependencies = [ + 'requests<3', + 'rich', + ] + "}; + + let expected_output = indoc::indoc! {r" + # /// script + # requires-python = '>=3.11' + # dependencies = [ + # 'requests<3', + # 'rich', + # ] + # /// + "}; + + let result = serialize_metadata(metadata); + assert_eq!(result, expected_output); + } + + #[test] + fn test_serialize_metadata_empty() { + let metadata = ""; + let expected_output = "# /// script\n# ///\n"; + + let result = serialize_metadata(metadata); + assert_eq!(result, expected_output); + } } diff --git a/crates/uv-workspace/src/pyproject_mut.rs b/crates/uv-workspace/src/pyproject_mut.rs index b105c98c1adf3..cb55e0619ec98 100644 --- a/crates/uv-workspace/src/pyproject_mut.rs +++ b/crates/uv-workspace/src/pyproject_mut.rs @@ -8,7 +8,7 @@ use thiserror::Error; use toml_edit::{Array, DocumentMut, Item, RawString, Table, TomlError, Value}; use uv_fs::PortablePath; -use crate::pyproject::{DependencyType, PyProjectToml, Source}; +use crate::pyproject::{DependencyType, Source}; /// Raw and mutable representation of a `pyproject.toml`. /// @@ -16,6 +16,7 @@ use crate::pyproject::{DependencyType, PyProjectToml, Source}; /// preserving comments and other structure, such as `uv add` and `uv remove`. pub struct PyProjectTomlMut { doc: DocumentMut, + dependency_target: DependencyTarget, } #[derive(Error, Debug)] @@ -47,11 +48,19 @@ pub enum ArrayEdit { Add(usize), } +/// Specifies whether dependencies are added to a script file or a `pyproject.toml` file. +#[derive(Debug, Copy, Clone, PartialEq, Eq)] +pub enum DependencyTarget { + Script, + PyProjectToml, +} + impl PyProjectTomlMut { - /// Initialize a [`PyProjectTomlMut`] from a [`PyProjectToml`]. - pub fn from_toml(pyproject: &PyProjectToml) -> Result { + /// Initialize a [`PyProjectTomlMut`] from a [`str`]. + pub fn from_toml(raw: &str, dependency_target: DependencyTarget) -> Result { Ok(Self { - doc: pyproject.raw.parse().map_err(Box::new)?, + doc: raw.parse().map_err(Box::new)?, + dependency_target, }) } @@ -83,6 +92,32 @@ impl PyProjectTomlMut { Ok(()) } + /// Retrieves a mutable reference to the root `Table` of the TOML document, creating the `project` table if necessary. + fn doc(&mut self) -> Result<&mut toml_edit::Table, Error> { + let doc = match self.dependency_target { + DependencyTarget::Script => self.doc.as_table_mut(), + DependencyTarget::PyProjectToml => self + .doc + .entry("project") + .or_insert(Item::Table(Table::new())) + .as_table_mut() + .ok_or(Error::MalformedDependencies)?, + }; + Ok(doc) + } + + /// Retrieves an optional mutable reference to the `project` `Table`, returning `None` if it doesn't exist. + fn doc_mut(&mut self) -> Result, Error> { + let doc = match self.dependency_target { + DependencyTarget::Script => Some(self.doc.as_table_mut()), + DependencyTarget::PyProjectToml => self + .doc + .get_mut("project") + .map(|project| project.as_table_mut().ok_or(Error::MalformedSources)) + .transpose()?, + }; + Ok(doc) + } /// Adds a dependency to `project.dependencies`. /// /// Returns `true` if the dependency was added, `false` if it was updated. @@ -93,11 +128,7 @@ impl PyProjectTomlMut { ) -> Result { // Get or create `project.dependencies`. let dependencies = self - .doc - .entry("project") - .or_insert(Item::Table(Table::new())) - .as_table_mut() - .ok_or(Error::MalformedDependencies)? + .doc()? .entry("dependencies") .or_insert(Item::Value(Value::Array(Array::new()))) .as_array_mut() @@ -158,11 +189,7 @@ impl PyProjectTomlMut { ) -> Result { // Get or create `project.optional-dependencies`. let optional_dependencies = self - .doc - .entry("project") - .or_insert(Item::Table(Table::new())) - .as_table_mut() - .ok_or(Error::MalformedDependencies)? + .doc()? .entry("optional-dependencies") .or_insert(Item::Table(Table::new())) .as_table_mut() @@ -192,11 +219,7 @@ impl PyProjectTomlMut { ) -> Result<(), Error> { // Get or create `project.dependencies`. let dependencies = self - .doc - .entry("project") - .or_insert(Item::Table(Table::new())) - .as_table_mut() - .ok_or(Error::MalformedDependencies)? + .doc()? .entry("dependencies") .or_insert(Item::Value(Value::Array(Array::new()))) .as_array_mut() @@ -265,11 +288,7 @@ impl PyProjectTomlMut { ) -> Result<(), Error> { // Get or create `project.optional-dependencies`. let optional_dependencies = self - .doc - .entry("project") - .or_insert(Item::Table(Table::new())) - .as_table_mut() - .ok_or(Error::MalformedDependencies)? + .doc()? .entry("optional-dependencies") .or_insert(Item::Table(Table::new())) .as_table_mut() @@ -323,10 +342,7 @@ impl PyProjectTomlMut { pub fn remove_dependency(&mut self, req: &PackageName) -> Result, Error> { // Try to get `project.dependencies`. let Some(dependencies) = self - .doc - .get_mut("project") - .map(|project| project.as_table_mut().ok_or(Error::MalformedSources)) - .transpose()? + .doc_mut()? .and_then(|project| project.get_mut("dependencies")) .map(|dependencies| dependencies.as_array_mut().ok_or(Error::MalformedSources)) .transpose()? @@ -372,10 +388,7 @@ impl PyProjectTomlMut { ) -> Result, Error> { // Try to get `project.optional-dependencies.`. let Some(optional_dependencies) = self - .doc - .get_mut("project") - .map(|project| project.as_table_mut().ok_or(Error::MalformedSources)) - .transpose()? + .doc_mut()? .and_then(|project| project.get_mut("optional-dependencies")) .map(|extras| extras.as_table_mut().ok_or(Error::MalformedSources)) .transpose()? diff --git a/crates/uv-workspace/src/workspace.rs b/crates/uv-workspace/src/workspace.rs index 6dad72e6e8c72..8414a446a3a46 100644 --- a/crates/uv-workspace/src/workspace.rs +++ b/crates/uv-workspace/src/workspace.rs @@ -104,7 +104,7 @@ impl Workspace { let pyproject_path = project_path.join("pyproject.toml"); let contents = fs_err::tokio::read_to_string(&pyproject_path).await?; - let pyproject_toml = PyProjectToml::from_string(contents) + let pyproject_toml = PyProjectToml::from_string(contents.clone()) .map_err(|err| WorkspaceError::Toml(pyproject_path.clone(), Box::new(err)))?; // Check if the project is explicitly marked as unmanaged. diff --git a/crates/uv/src/commands/project/add.rs b/crates/uv/src/commands/project/add.rs index 30d9a3eaf5c60..e689d6c1c40fe 100644 --- a/crates/uv/src/commands/project/add.rs +++ b/crates/uv/src/commands/project/add.rs @@ -15,24 +15,34 @@ use uv_dispatch::BuildDispatch; use uv_distribution::DistributionDatabase; use uv_fs::CWD; use uv_normalize::PackageName; -use uv_python::{PythonDownloads, PythonPreference, PythonRequest}; +use uv_python::{ + request_from_version_file, EnvironmentPreference, PythonDownloads, PythonInstallation, + PythonPreference, PythonRequest, VersionRequest, +}; use uv_requirements::{NamedRequirementsResolver, RequirementsSource, RequirementsSpecification}; use uv_resolver::FlatIndex; +use uv_scripts::Pep723Script; use uv_types::{BuildIsolation, HashStrategy}; use uv_warnings::warn_user_once; use uv_workspace::pyproject::{DependencyType, Source, SourceError}; -use uv_workspace::pyproject_mut::{ArrayEdit, PyProjectTomlMut}; +use uv_workspace::pyproject_mut::{ArrayEdit, DependencyTarget, PyProjectTomlMut}; use uv_workspace::{DiscoveryOptions, VirtualProject, Workspace}; use crate::commands::pip::loggers::{DefaultInstallLogger, DefaultResolveLogger}; use crate::commands::pip::operations::Modifications; use crate::commands::pip::resolution_environment; use crate::commands::project::ProjectError; -use crate::commands::reporters::ResolverReporter; +use crate::commands::reporters::{PythonDownloadReporter, ResolverReporter}; 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( @@ -50,6 +60,7 @@ pub(crate) async fn add( package: Option, python: Option, settings: ResolverInstallerSettings, + script: Option, python_preference: PythonPreference, python_downloads: PythonDownloads, preview: PreviewMode, @@ -63,43 +74,93 @@ pub(crate) async fn add( warn_user_once!("`uv add` is experimental and may change without warning"); } - // Find the project in the workspace. - let project = if let Some(package) = package { - VirtualProject::Project( - Workspace::discover(&CWD, &DiscoveryOptions::default()) - .await? - .with_current_project(package.clone()) - .with_context(|| format!("Package `{package}` not found in workspace"))?, + let download_reporter = PythonDownloadReporter::single(printer); + let (dependency_destination, venv) = if let Some(script) = script { + // (1) Explicit request from user + let python_request = if let Some(request) = python.as_deref() { + Some(PythonRequest::parse(request)) + // (2) Request from `.python-version` + } else if let Some(request) = request_from_version_file(&CWD).await? { + Some(request) + // (3) `Requires-Python` in `pyproject.toml` + } else { + script + .metadata + .requires_python + .clone() + .map(|requires_python| { + PythonRequest::Version(VersionRequest::Range(requires_python)) + }) + }; + + let client_builder = BaseClientBuilder::new() + .connectivity(connectivity) + .native_tls(native_tls); + + let interpreter = PythonInstallation::find_or_download( + python_request, + EnvironmentPreference::Any, + python_preference, + python_downloads, + &client_builder, + cache, + Some(&download_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) } else { - VirtualProject::discover(&CWD, &DiscoveryOptions::default()).await? - }; + // Find the project in the workspace. + let project = if let Some(package) = package { + VirtualProject::Project( + Workspace::discover(&CWD, &DiscoveryOptions::default()) + .await? + .with_current_project(package.clone()) + .with_context(|| format!("Package `{package}` not found in workspace"))?, + ) + } else { + VirtualProject::discover(&CWD, &DiscoveryOptions::default()).await? + }; - // For virtual projects, allow dev dependencies, but nothing else. - if project.is_virtual() { - match dependency_type { - DependencyType::Production => { - anyhow::bail!("Found a virtual workspace root, but virtual projects do not support production dependencies (instead, use: `{}`)", "uv add --dev".green()) - } - DependencyType::Optional(_) => { - anyhow::bail!("Found a virtual workspace root, but virtual projects do not support optional dependencies (instead, use: `{}`)", "uv add --dev".green()) + // For virtual projects, allow dev dependencies, but nothing else. + if project.is_virtual() { + match dependency_type { + DependencyType::Production => { + anyhow::bail!("Found a virtual workspace root, but virtual projects do not support production dependencies (instead, use: `{}`)", "uv add --dev".green()) + } + DependencyType::Optional(_) => { + anyhow::bail!("Found a virtual workspace root, but virtual projects do not support optional dependencies (instead, use: `{}`)", "uv add --dev".green()) + } + DependencyType::Dev => (), } - DependencyType::Dev => (), } - } - // Discover or create the virtual environment. - let venv = project::get_or_init_environment( - project.workspace(), - python.as_deref().map(PythonRequest::parse), - python_preference, - python_downloads, - connectivity, - native_tls, - cache, - printer, - ) - .await?; + // Discover or create the virtual environment. + let venv = project::get_or_init_environment( + project.workspace(), + python.as_deref().map(PythonRequest::parse), + python_preference, + python_downloads, + connectivity, + native_tls, + cache, + printer, + ) + .await?; + (DependencyDestination::Project(project), venv) + }; let client_builder = BaseClientBuilder::new() .connectivity(connectivity) @@ -183,8 +244,15 @@ pub(crate) async fn add( .await?; // Add the requirements to the `pyproject.toml`. - let existing = project.pyproject_toml(); - let mut pyproject = PyProjectTomlMut::from_toml(existing)?; + let mut toml = match &dependency_destination { + DependencyDestination::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) + } + }?; let mut edits = Vec::with_capacity(requirements.len()); for mut requirement in requirements { // Add the specified extras. @@ -192,48 +260,48 @@ pub(crate) async fn add( requirement.extras.sort_unstable(); requirement.extras.dedup(); - let (requirement, source) = if raw_sources { - // Use the PEP 508 requirement directly. - (pep508_rs::Requirement::from(requirement), None) - } else { - // Otherwise, try to construct the source. - let workspace = project - .workspace() - .packages() - .contains_key(&requirement.name); - let result = Source::from_requirement( - &requirement.name, - requirement.source.clone(), - workspace, - editable, - rev.clone(), - tag.clone(), - branch.clone(), - ); - - let source = match result { - Ok(source) => source, - Err(SourceError::UnresolvedReference(rev)) => { - anyhow::bail!("Cannot resolve Git reference `{rev}` for requirement `{name}`. Specify the reference with one of `--tag`, `--branch`, or `--rev`, or use the `--raw-sources` flag.", name = requirement.name) - } - Err(err) => return Err(err.into()), - }; + let (requirement, source) = match dependency_destination { + DependencyDestination::Script(_) => (pep508_rs::Requirement::from(requirement), None), + DependencyDestination::Project(_) if raw_sources => { + (pep508_rs::Requirement::from(requirement), None) + } + DependencyDestination::Project(ref project) => { + // Otherwise, try to construct the source. + let workspace = project + .workspace() + .packages() + .contains_key(&requirement.name); + let result = Source::from_requirement( + &requirement.name, + requirement.source.clone(), + workspace, + editable, + rev.clone(), + tag.clone(), + branch.clone(), + ); + + let source = match result { + Ok(source) => source, + Err(SourceError::UnresolvedReference(rev)) => { + anyhow::bail!("Cannot resolve Git reference `{rev}` for requirement `{name}`. Specify the reference with one of `--tag`, `--branch`, or `--rev`, or use the `--raw-sources` flag.", name = requirement.name) + } + Err(err) => return Err(err.into()), + }; - // Ignore the PEP 508 source. - let mut requirement = pep508_rs::Requirement::from(requirement); - requirement.clear_url(); + // Ignore the PEP 508 source. + let mut requirement = pep508_rs::Requirement::from(requirement); + requirement.clear_url(); - (requirement, source) + (requirement, source) + } }; - // Update the `pyproject.toml`. let edit = match dependency_type { - DependencyType::Production => { - pyproject.add_dependency(&requirement, source.as_ref())? - } - DependencyType::Dev => pyproject.add_dev_dependency(&requirement, source.as_ref())?, + DependencyType::Production => toml.add_dependency(&requirement, source.as_ref())?, + DependencyType::Dev => toml.add_dev_dependency(&requirement, source.as_ref())?, DependencyType::Optional(ref group) => { - pyproject.add_optional_dependency(group, &requirement, source.as_ref())? + toml.add_optional_dependency(group, &requirement, source.as_ref())? } }; @@ -247,21 +315,40 @@ pub(crate) async fn add( } // Save the modified `pyproject.toml`. - let mut modified = false; - let content = pyproject.to_string(); - if content == existing.raw { - debug!("No changes to `pyproject.toml`; skipping update"); - } else { - fs_err::write(project.root().join("pyproject.toml"), &content)?; - modified = true; - } - + let content = toml.to_string(); + let modified = match &dependency_destination { + DependencyDestination::Script(script) => { + if content == script.metadata.raw { + debug!("No changes to dependencies; skipping update"); + false + } else { + script.replace_metadata(&content).await?; + true + } + } + DependencyDestination::Project(project) => { + if content == *project.pyproject_toml().raw { + debug!("No changes to dependencies; skipping update"); + false + } else { + let pyproject_path = project.root().join("pyproject.toml"); + fs_err::write(pyproject_path, &content)?; + true + } + } + }; // If `--frozen`, exit early. There's no reason to lock and sync, and we don't need a `uv.lock` // to exist at all. if frozen { return Ok(ExitStatus::Success); } + // If `--script`, exit early. There's no reason to lock and sync. + let DependencyDestination::Project(project) = dependency_destination else { + return Ok(ExitStatus::Success); + }; + + let existing = project.pyproject_toml(); // Update the `pypackage.toml` in-memory. let project = project .clone() @@ -357,13 +444,13 @@ pub(crate) async fn add( match edit.dependency_type { DependencyType::Production => { - pyproject.set_dependency_minimum_version(*index, minimum)?; + toml.set_dependency_minimum_version(*index, minimum)?; } DependencyType::Dev => { - pyproject.set_dev_dependency_minimum_version(*index, minimum)?; + toml.set_dev_dependency_minimum_version(*index, minimum)?; } DependencyType::Optional(ref group) => { - pyproject.set_optional_dependency_minimum_version(group, *index, minimum)?; + toml.set_optional_dependency_minimum_version(group, *index, minimum)?; } } @@ -374,7 +461,7 @@ pub(crate) async fn add( // string content, since the above loop _must_ change an empty specifier to a non-empty // specifier. if modified { - fs_err::write(project.root().join("pyproject.toml"), pyproject.to_string())?; + fs_err::write(project.root().join("pyproject.toml"), toml.to_string())?; } } diff --git a/crates/uv/src/commands/project/init.rs b/crates/uv/src/commands/project/init.rs index abe75f6865b05..c469346fd7810 100644 --- a/crates/uv/src/commands/project/init.rs +++ b/crates/uv/src/commands/project/init.rs @@ -16,7 +16,7 @@ use uv_python::{ }; use uv_resolver::RequiresPython; use uv_warnings::warn_user_once; -use uv_workspace::pyproject_mut::PyProjectTomlMut; +use uv_workspace::pyproject_mut::{DependencyTarget, PyProjectTomlMut}; use uv_workspace::{check_nested_workspaces, DiscoveryOptions, Workspace, WorkspaceError}; use crate::commands::project::find_requires_python; @@ -315,7 +315,10 @@ async fn init_project( )?; } else { // Add the package to the workspace. - let mut pyproject = PyProjectTomlMut::from_toml(workspace.pyproject_toml())?; + let mut pyproject = PyProjectTomlMut::from_toml( + &workspace.pyproject_toml().raw, + DependencyTarget::PyProjectToml, + )?; pyproject.add_workspace(path.strip_prefix(workspace.install_path())?)?; // Save the modified `pyproject.toml`. diff --git a/crates/uv/src/commands/project/remove.rs b/crates/uv/src/commands/project/remove.rs index c0129ec6ec768..970eeb31f1e2d 100644 --- a/crates/uv/src/commands/project/remove.rs +++ b/crates/uv/src/commands/project/remove.rs @@ -6,10 +6,11 @@ use uv_client::Connectivity; use uv_configuration::{Concurrency, ExtrasSpecification, PreviewMode}; use uv_fs::CWD; use uv_python::{PythonDownloads, PythonPreference, PythonRequest}; +use uv_scripts::Pep723Script; use uv_warnings::{warn_user, warn_user_once}; use uv_workspace::pyproject::DependencyType; -use uv_workspace::pyproject_mut::PyProjectTomlMut; -use uv_workspace::{DiscoveryOptions, ProjectWorkspace, VirtualProject, Workspace}; +use uv_workspace::pyproject_mut::{DependencyTarget, PyProjectTomlMut}; +use uv_workspace::{DiscoveryOptions, VirtualProject, Workspace}; use crate::commands::pip::loggers::{DefaultInstallLogger, DefaultResolveLogger}; use crate::commands::pip::operations::Modifications; @@ -17,6 +18,12 @@ 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( @@ -28,6 +35,7 @@ pub(crate) async fn remove( package: Option, python: Option, settings: ResolverInstallerSettings, + script: Option, python_preference: PythonPreference, python_downloads: PythonDownloads, preview: PreviewMode, @@ -41,41 +49,58 @@ pub(crate) async fn remove( warn_user_once!("`uv remove` is experimental and may change without warning"); } - // Find the project in the workspace. - let project = if let Some(package) = package { - Workspace::discover(&CWD, &DiscoveryOptions::default()) - .await? - .with_current_project(package.clone()) - .with_context(|| format!("Package `{package}` not found in workspace"))? + let dependency_destination = if let Some(script) = script { + DependencyDestination::Script(script) } else { - ProjectWorkspace::discover(&CWD, &DiscoveryOptions::default()).await? + // Find the project in the workspace. + let project = if let Some(package) = package { + VirtualProject::Project( + Workspace::discover(&CWD, &DiscoveryOptions::default()) + .await? + .with_current_project(package.clone()) + .with_context(|| format!("Package `{package}` not found in workspace"))?, + ) + } else { + VirtualProject::discover(&CWD, &DiscoveryOptions::default()).await? + }; + + DependencyDestination::Project(project) }; - let mut pyproject = PyProjectTomlMut::from_toml(project.current_project().pyproject_toml())?; + let mut toml = match &dependency_destination { + DependencyDestination::Script(script) => { + PyProjectTomlMut::from_toml(&script.metadata.raw, DependencyTarget::Script) + } + DependencyDestination::Project(project) => PyProjectTomlMut::from_toml( + project.pyproject_toml().raw.as_ref(), + DependencyTarget::PyProjectToml, + ), + }?; + for package in packages { match dependency_type { DependencyType::Production => { - let deps = pyproject.remove_dependency(&package)?; + let deps = toml.remove_dependency(&package)?; if deps.is_empty() { - warn_if_present(&package, &pyproject); + warn_if_present(&package, &toml); anyhow::bail!( "The dependency `{package}` could not be found in `dependencies`" ); } } DependencyType::Dev => { - let deps = pyproject.remove_dev_dependency(&package)?; + let deps = toml.remove_dev_dependency(&package)?; if deps.is_empty() { - warn_if_present(&package, &pyproject); + warn_if_present(&package, &toml); anyhow::bail!( "The dependency `{package}` could not be found in `dev-dependencies`" ); } } DependencyType::Optional(ref group) => { - let deps = pyproject.remove_optional_dependency(&package, group)?; + let deps = toml.remove_optional_dependency(&package, group)?; if deps.is_empty() { - warn_if_present(&package, &pyproject); + warn_if_present(&package, &toml); anyhow::bail!( "The dependency `{package}` could not be found in `optional-dependencies`" ); @@ -84,11 +109,16 @@ pub(crate) async fn remove( } } - // Save the modified `pyproject.toml`. - fs_err::write( - project.current_project().root().join("pyproject.toml"), - pyproject.to_string(), - )?; + // Save the modified dependencies. + match &dependency_destination { + DependencyDestination::Script(script) => { + script.replace_metadata(&toml.to_string()).await?; + } + DependencyDestination::Project(project) => { + let pyproject_path = project.root().join("pyproject.toml"); + fs_err::write(pyproject_path, toml.to_string())?; + } + }; // If `--frozen`, exit early. There's no reason to lock and sync, and we don't need a `uv.lock` // to exist at all. @@ -96,6 +126,11 @@ pub(crate) async fn remove( return Ok(ExitStatus::Success); } + // If `--script`, exit early. There's no reason to lock and sync. + let DependencyDestination::Project(project) = dependency_destination else { + return Ok(ExitStatus::Success); + }; + // Discover or create the virtual environment. let venv = project::get_or_init_environment( project.workspace(), @@ -139,7 +174,7 @@ pub(crate) async fn remove( let state = SharedState::default(); project::sync::do_sync( - &VirtualProject::Project(project), + &project, &venv, &lock.lock, &extras, diff --git a/crates/uv/src/lib.rs b/crates/uv/src/lib.rs index c5f093cf0744e..c364bbb47f55d 100644 --- a/crates/uv/src/lib.rs +++ b/crates/uv/src/lib.rs @@ -1141,6 +1141,12 @@ async fn run_project( .combine(Refresh::from(args.settings.reinstall.clone())) .combine(Refresh::from(args.settings.upgrade.clone())), ); + // If the target is a PEP 723 script, parse it. + let script = if let Some(script) = args.script { + Pep723Script::read(&script).await? + } else { + None + }; commands::add( args.locked, @@ -1157,6 +1163,7 @@ async fn run_project( args.package, args.python, args.settings, + script, globals.python_preference, globals.python_downloads, globals.preview, @@ -1179,6 +1186,12 @@ async fn run_project( .combine(Refresh::from(args.settings.reinstall.clone())) .combine(Refresh::from(args.settings.upgrade.clone())), ); + // If the target is a PEP 723 script, parse it. + let script = if let Some(script) = args.script { + Pep723Script::read(&script).await? + } else { + None + }; commands::remove( args.locked, @@ -1189,6 +1202,7 @@ async fn run_project( args.package, args.python, args.settings, + script, globals.python_preference, globals.python_downloads, globals.preview, @@ -1362,3 +1376,4 @@ where } } } + diff --git a/crates/uv/src/settings.rs b/crates/uv/src/settings.rs index be54b1de8edda..8e2d3856cc8ea 100644 --- a/crates/uv/src/settings.rs +++ b/crates/uv/src/settings.rs @@ -706,6 +706,7 @@ pub(crate) struct AddSettings { pub(crate) python: Option, pub(crate) refresh: Refresh, pub(crate) settings: ResolverInstallerSettings, + pub(crate) script: Option, } impl AddSettings { @@ -731,6 +732,7 @@ impl AddSettings { refresh, package, python, + script, } = args; let requirements = requirements @@ -765,6 +767,7 @@ impl AddSettings { resolver_installer_options(installer, build), filesystem, ), + script, } } } @@ -782,6 +785,7 @@ pub(crate) struct RemoveSettings { pub(crate) python: Option, pub(crate) refresh: Refresh, pub(crate) settings: ResolverInstallerSettings, + pub(crate) script: Option, } impl RemoveSettings { @@ -800,6 +804,7 @@ impl RemoveSettings { refresh, package, python, + script, } = args; let dependency_type = if let Some(group) = optional { @@ -823,6 +828,7 @@ impl RemoveSettings { resolver_installer_options(installer, build), filesystem, ), + script, } } } diff --git a/crates/uv/tests/edit.rs b/crates/uv/tests/edit.rs index a11ef0d9e88e3..22a9da3dbec07 100644 --- a/crates/uv/tests/edit.rs +++ b/crates/uv/tests/edit.rs @@ -2912,3 +2912,127 @@ fn add_repeat() -> Result<()> { Ok(()) } + +/// Add to a PEP732 script +#[test] +fn add_script() -> Result<()> { + let context = TestContext::new("3.12"); + + let script = context.temp_dir.child("script.py"); + script.write_str(indoc! {r#" + # /// script + # requires-python = ">=3.11" + # dependencies = [ + # "requests<3", + # "rich", + # ] + # /// + + + import requests + from rich.pretty import pprint + + resp = requests.get("https://peps.python.org/api/peps.json") + data = resp.json() + pprint([(k, v["title"]) for k, v in data.items()][:10]) + "#})?; + + uv_snapshot!(context.filters(), context.add(&["anyio"]).arg("--script").arg(script.path()), @r###" + success: true + exit_code: 0 + ----- stdout ----- + + ----- stderr ----- + warning: `uv add` is experimental and may change without warning + "###); + + let script_content = fs_err::read_to_string(context.temp_dir.join("script.py"))?; + + insta::with_settings!({ + filters => context.filters(), + }, { + assert_snapshot!( + script_content, @r###" + # /// script + # requires-python = ">=3.11" + # dependencies = [ + # "requests<3", + # "rich", + # "anyio", + # ] + # /// + + + import requests + from rich.pretty import pprint + + resp = requests.get("https://peps.python.org/api/peps.json") + data = resp.json() + pprint([(k, v["title"]) for k, v in data.items()][:10]) + "### + ); + }); + Ok(()) +} + +/// Remove from a PEP732 script +#[test] +fn remove_script() -> Result<()> { + let context = TestContext::new("3.12"); + + let script = context.temp_dir.child("script.py"); + script.write_str(indoc! {r#" + # /// script + # requires-python = ">=3.11" + # dependencies = [ + # "requests<3", + # "rich", + # "anyio", + # ] + # /// + + + import requests + from rich.pretty import pprint + + resp = requests.get("https://peps.python.org/api/peps.json") + data = resp.json() + pprint([(k, v["title"]) for k, v in data.items()][:10]) + "#})?; + + uv_snapshot!(context.filters(), context.remove(&["anyio"]).arg("--script").arg(script.path()), @r###" + success: true + exit_code: 0 + ----- stdout ----- + + ----- stderr ----- + warning: `uv remove` is experimental and may change without warning + "###); + + let script_content = fs_err::read_to_string(context.temp_dir.join("script.py"))?; + + insta::with_settings!({ + filters => context.filters(), + }, { + assert_snapshot!( + script_content, @r###" + # /// script + # requires-python = ">=3.11" + # dependencies = [ + # "requests<3", + # "rich", + # ] + # /// + + + import requests + from rich.pretty import pprint + + resp = requests.get("https://peps.python.org/api/peps.json") + data = resp.json() + pprint([(k, v["title"]) for k, v in data.items()][:10]) + "### + ); + }); + Ok(()) +} diff --git a/docs/reference/cli.md b/docs/reference/cli.md index 284031ba05cdc..550631310d5fd 100644 --- a/docs/reference/cli.md +++ b/docs/reference/cli.md @@ -711,6 +711,8 @@ uv add [OPTIONS] ...
--rev rev

Commit to use when adding a dependency from Git

+
--script script

Specifies the Python script where the dependency will be added

+
--tag tag

Tag to use when adding a dependency from Git

--upgrade, -U

Allow package upgrades, ignoring pinned versions in any existing output file. Implies --refresh

@@ -967,6 +969,8 @@ uv remove [OPTIONS] ...
  • lowest-direct: Resolve the lowest compatible version of any direct dependencies, and the highest compatible version of any transitive dependencies
    • +
    --script script

    Specifies the Python script where the dependency will be removed

    +
    --upgrade, -U

    Allow package upgrades, ignoring pinned versions in any existing output file. Implies --refresh

    --upgrade-package, -P upgrade-package

    Allow upgrades for a specific package, ignoring pinned versions in any existing output file. Implies --refresh-package