Allow uv add to specify optional dependency groups (#4607)

## Summary

Implements `uv add --optional <group>`, which adds a dependency to
`project.optional-dependency.<group>`.

Resolves https://github.com/astral-sh/uv/issues/4585.
This commit is contained in:
Ibraheem Ahmed 2024-06-27 21:24:21 -04:00 committed by GitHub
parent 9b38450998
commit bbd59ff455
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
8 changed files with 443 additions and 113 deletions

View file

@ -289,6 +289,17 @@ impl Source {
}
}
/// The type of a dependency in a `pyproject.toml`.
#[derive(Debug, Clone)]
pub enum DependencyType {
/// A dependency in `project.dependencies`.
Production,
/// A dependency in `tool.uv.dev-dependencies`.
Dev,
/// A dependency in `project.optional-dependencies.{0}`.
Optional(ExtraName),
}
/// <https://github.com/serde-rs/serde/issues/1316#issue-332908452>
mod serde_from_and_to_string {
use std::fmt::Display;

View file

@ -4,9 +4,9 @@ use std::{fmt, mem};
use thiserror::Error;
use toml_edit::{Array, DocumentMut, Item, RawString, Table, TomlError, Value};
use pep508_rs::{PackageName, Requirement, VersionOrUrl};
use pep508_rs::{ExtraName, PackageName, Requirement, VersionOrUrl};
use crate::pyproject::{PyProjectToml, Source};
use crate::pyproject::{DependencyType, PyProjectToml, Source};
/// Raw and mutable representation of a `pyproject.toml`.
///
@ -60,23 +60,7 @@ impl PyProjectTomlMut {
add_dependency(req, dependencies, source.is_some())?;
if let Some(source) = source {
// Get or create `tool.uv.sources`.
let sources = self
.doc
.entry("tool")
.or_insert(implicit())
.as_table_mut()
.ok_or(Error::MalformedSources)?
.entry("uv")
.or_insert(implicit())
.as_table_mut()
.ok_or(Error::MalformedSources)?
.entry("sources")
.or_insert(Item::Table(Table::new()))
.as_table_mut()
.ok_or(Error::MalformedSources)?;
add_source(&name, &source, sources)?;
self.add_source(&name, &source)?;
}
Ok(())
@ -88,8 +72,8 @@ impl PyProjectTomlMut {
req: Requirement,
source: Option<Source>,
) -> Result<(), Error> {
// Get or create `tool.uv`.
let tool_uv = self
// Get or create `tool.uv.dev-dependencies`.
let dev_dependencies = self
.doc
.entry("tool")
.or_insert(implicit())
@ -98,10 +82,7 @@ impl PyProjectTomlMut {
.entry("uv")
.or_insert(Item::Table(Table::new()))
.as_table_mut()
.ok_or(Error::MalformedSources)?;
// Get or create the `tool.uv.dev-dependencies` array.
let dev_dependencies = tool_uv
.ok_or(Error::MalformedSources)?
.entry("dev-dependencies")
.or_insert(Item::Value(Value::Array(Array::new())))
.as_array_mut()
@ -111,82 +92,215 @@ impl PyProjectTomlMut {
add_dependency(req, dev_dependencies, source.is_some())?;
if let Some(source) = source {
// Get or create `tool.uv.sources`.
let sources = tool_uv
.entry("sources")
.or_insert(Item::Table(Table::new()))
.as_table_mut()
.ok_or(Error::MalformedSources)?;
add_source(&name, &source, sources)?;
self.add_source(&name, &source)?;
}
Ok(())
}
/// Adds a dependency to `project.optional-dependencies`.
pub fn add_optional_dependency(
&mut self,
req: Requirement,
group: &ExtraName,
source: Option<Source>,
) -> 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)?
.entry("optional-dependencies")
.or_insert(Item::Table(Table::new()))
.as_table_mut()
.ok_or(Error::MalformedDependencies)?;
let group = optional_dependencies
.entry(group.as_ref())
.or_insert(Item::Value(Value::Array(Array::new())))
.as_array_mut()
.ok_or(Error::MalformedDependencies)?;
let name = req.name.clone();
add_dependency(req, group, source.is_some())?;
if let Some(source) = source {
self.add_source(&name, &source)?;
}
Ok(())
}
/// Adds a source to `tool.uv.sources`.
fn add_source(&mut self, name: &PackageName, source: &Source) -> Result<(), Error> {
// Get or create `tool.uv.sources`.
let sources = self
.doc
.entry("tool")
.or_insert(implicit())
.as_table_mut()
.ok_or(Error::MalformedSources)?
.entry("uv")
.or_insert(implicit())
.as_table_mut()
.ok_or(Error::MalformedSources)?
.entry("sources")
.or_insert(Item::Table(Table::new()))
.as_table_mut()
.ok_or(Error::MalformedSources)?;
add_source(name, source, sources)?;
Ok(())
}
/// Removes all occurrences of dependencies with the given name.
pub fn remove_dependency(&mut self, req: &PackageName) -> Result<Vec<Requirement>, Error> {
// Try to get `project.dependencies`.
let Some(dependencies) = self
.doc
.get_mut("project")
.and_then(Item::as_table_mut)
.map(|project| project.as_table_mut().ok_or(Error::MalformedSources))
.transpose()?
.and_then(|project| project.get_mut("dependencies"))
.map(|dependencies| dependencies.as_array_mut().ok_or(Error::MalformedSources))
.transpose()?
else {
return Ok(Vec::new());
};
let dependencies = dependencies
.as_array_mut()
.ok_or(Error::MalformedDependencies)?;
let requirements = remove_dependency(req, dependencies);
// Remove a matching source from `tool.uv.sources`, if it exists.
if let Some(sources) = self
.doc
.get_mut("tool")
.and_then(Item::as_table_mut)
.and_then(|tool| tool.get_mut("uv"))
.and_then(Item::as_table_mut)
.and_then(|tool_uv| tool_uv.get_mut("sources"))
{
let sources = sources.as_table_mut().ok_or(Error::MalformedSources)?;
sources.remove(req.as_ref());
}
self.remove_source(req)?;
Ok(requirements)
}
/// Removes all occurrences of development dependencies with the given name.
pub fn remove_dev_dependency(&mut self, req: &PackageName) -> Result<Vec<Requirement>, Error> {
let Some(tool_uv) = self
// Try to get `tool.uv.dev-dependencies`.
let Some(dev_dependencies) = self
.doc
.get_mut("tool")
.and_then(Item::as_table_mut)
.map(|tool| tool.as_table_mut().ok_or(Error::MalformedSources))
.transpose()?
.and_then(|tool| tool.get_mut("uv"))
.and_then(Item::as_table_mut)
.map(|tool_uv| tool_uv.as_table_mut().ok_or(Error::MalformedSources))
.transpose()?
.and_then(|tool_uv| tool_uv.get_mut("dev-dependencies"))
.map(|dependencies| dependencies.as_array_mut().ok_or(Error::MalformedSources))
.transpose()?
else {
return Ok(Vec::new());
};
// Try to get `tool.uv.dev-dependencies`.
let Some(dev_dependencies) = tool_uv.get_mut("dev-dependencies") else {
return Ok(Vec::new());
};
let dev_dependencies = dev_dependencies
.as_array_mut()
.ok_or(Error::MalformedDependencies)?;
let requirements = remove_dependency(req, dev_dependencies);
// Remove a matching source from `tool.uv.sources`, if it exists.
if let Some(sources) = tool_uv.get_mut("sources") {
let sources = sources.as_table_mut().ok_or(Error::MalformedSources)?;
sources.remove(req.as_ref());
};
self.remove_source(req)?;
Ok(requirements)
}
/// Removes all occurrences of optional dependencies in the group with the given name.
pub fn remove_optional_dependency(
&mut self,
req: &PackageName,
group: &ExtraName,
) -> Result<Vec<Requirement>, Error> {
// Try to get `project.optional-dependencies.<group>`.
let Some(optional_dependencies) = self
.doc
.get_mut("project")
.map(|project| project.as_table_mut().ok_or(Error::MalformedSources))
.transpose()?
.and_then(|project| project.get_mut("optional-dependencies"))
.map(|extras| extras.as_table_mut().ok_or(Error::MalformedSources))
.transpose()?
.and_then(|extras| extras.get_mut(group.as_ref()))
.map(|dependencies| dependencies.as_array_mut().ok_or(Error::MalformedSources))
.transpose()?
else {
return Ok(Vec::new());
};
let requirements = remove_dependency(req, optional_dependencies);
self.remove_source(req)?;
Ok(requirements)
}
// Remove a matching source from `tool.uv.sources`, if it exists.
fn remove_source(&mut self, name: &PackageName) -> Result<(), Error> {
if let Some(sources) = self
.doc
.get_mut("tool")
.map(|tool| tool.as_table_mut().ok_or(Error::MalformedSources))
.transpose()?
.and_then(|tool| tool.get_mut("uv"))
.map(|tool_uv| tool_uv.as_table_mut().ok_or(Error::MalformedSources))
.transpose()?
.and_then(|tool_uv| tool_uv.get_mut("sources"))
.map(|sources| sources.as_table_mut().ok_or(Error::MalformedSources))
.transpose()?
{
sources.remove(name.as_ref());
}
Ok(())
}
/// Returns all the places in this `pyproject.toml` that contain a dependency with the given
/// name.
///
/// This method searches `project.dependencies`, `tool.uv.dev-dependencies`, and
/// `tool.uv.optional-dependencies`.
pub fn find_dependency(&self, name: &PackageName) -> Vec<DependencyType> {
let mut types = Vec::new();
if let Some(project) = self.doc.get("project").and_then(Item::as_table) {
// Check `project.dependencies`.
if let Some(dependencies) = project.get("dependencies").and_then(Item::as_array) {
if !find_dependencies(name, dependencies).is_empty() {
types.push(DependencyType::Production);
}
}
// Check `project.optional-dependencies`.
if let Some(extras) = project
.get("optional-dependencies")
.and_then(Item::as_table)
{
for (extra, dependencies) in extras {
let Some(dependencies) = dependencies.as_array() else {
continue;
};
let Ok(extra) = ExtraName::new(extra.to_string()) else {
continue;
};
if !find_dependencies(name, dependencies).is_empty() {
types.push(DependencyType::Optional(extra));
}
}
}
}
// Check `tool.uv.dev-dependencies`.
if let Some(dev_dependencies) = self
.doc
.get("tool")
.and_then(Item::as_table)
.and_then(|tool| tool.get("uv"))
.and_then(Item::as_table)
.and_then(|tool| tool.get("dev-dependencies"))
.and_then(Item::as_array)
{
if !find_dependencies(name, dev_dependencies).is_empty() {
types.push(DependencyType::Dev);
}
}
types
}
}
/// Returns an implicit table.