Skip to content

Commit

Permalink
allow uv add to specify optional dependency groups
Browse files Browse the repository at this point in the history
  • Loading branch information
@ibraheemdev
ibraheemdev committed Jun 28, 2024
1 parent 9b38450 commit d122994
Show file tree
Hide file tree
Showing 8 changed files with 438 additions and 111 deletions.
12 changes: 10 additions & 2 deletions crates/uv-cli/src/lib.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1763,9 +1763,13 @@ pub struct AddArgs {
pub requirements: Vec<String>,

/// Add the requirements as development dependencies.
#[arg(long)]
#[arg(long, conflicts_with("optional"))]
pub dev: bool,

/// Add the requirements to the specified optional dependency group.
#[arg(long, conflicts_with("dev"))]
pub optional: Option<ExtraName>,

/// Add the requirements as editables.
#[arg(long, default_missing_value = "true", num_args(0..=1))]
pub editable: Option<bool>,
Expand Down Expand Up @@ -1829,9 +1833,13 @@ pub struct RemoveArgs {
pub requirements: Vec<PackageName>,

/// Remove the requirements from development dependencies.
#[arg(long)]
#[arg(long, conflicts_with("optional"))]
pub dev: bool,

/// Remove the requirements from the specified optional dependency group.
#[arg(long, conflicts_with("dev"))]
pub optional: Option<ExtraName>,

/// Remove the dependency from a specific package in the workspace.
#[arg(long, conflicts_with = "isolated")]
pub package: Option<PackageName>,
Expand Down
11 changes: 11 additions & 0 deletions crates/uv-distribution/src/pyproject.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down
244 changes: 179 additions & 65 deletions crates/uv-distribution/src/pyproject_mut.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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`.
///
Expand Down Expand Up @@ -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(())
Expand All @@ -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())
Expand All @@ -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()
Expand All @@ -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);
self.remove_source(req)?;

Ok(requirements)
}

// 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());
/// 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.
Expand Down
Loading

0 comments on commit d122994

Please sign in to comment.