Support PEP 723 scripts in uv add and uv remove (#5995)

## Summary

Resolves https://github.com/astral-sh/uv/issues/4667

## Test Plan

`cargo test`
This commit is contained in:
Ahmed Ilyas 2024-08-11 03:40:59 +02:00 committed by GitHub
parent 9b8c07bf18
commit 2d53e35e39
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
12 changed files with 1215 additions and 289 deletions

View file

@ -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,
target: DependencyTarget,
}
#[derive(Error, Debug)]
@ -47,11 +48,21 @@ 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 {
/// A PEP 723 script, with inline metadata.
Script,
/// A project with a `pyproject.toml`.
PyProjectToml,
}
impl PyProjectTomlMut {
/// Initialize a [`PyProjectTomlMut`] from a [`PyProjectToml`].
pub fn from_toml(pyproject: &PyProjectToml) -> Result<Self, Error> {
/// Initialize a [`PyProjectTomlMut`] from a [`str`].
pub fn from_toml(raw: &str, target: DependencyTarget) -> Result<Self, Error> {
Ok(Self {
doc: pyproject.raw.parse().map_err(Box::new)?,
doc: raw.parse().map_err(Box::new)?,
target,
})
}
@ -83,6 +94,34 @@ 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 Table, Error> {
let doc = match self.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<Option<&mut Table>, Error> {
let doc = match self.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 +132,7 @@ impl PyProjectTomlMut {
) -> Result<ArrayEdit, 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()
@ -158,11 +193,7 @@ impl PyProjectTomlMut {
) -> Result<ArrayEdit, 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()
@ -192,11 +223,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 +292,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 +346,7 @@ impl PyProjectTomlMut {
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")
.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 +392,7 @@ impl PyProjectTomlMut {
) -> 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()?
.doc_mut()?
.and_then(|project| project.get_mut("optional-dependencies"))
.map(|extras| extras.as_table_mut().ok_or(Error::MalformedSources))
.transpose()?

View file

@ -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.