Support inline optional tables in uv add and uv remove (#6787)

## Summary Closes #6785.
astral-sh · Aug 29, 2024 · 933d4ef · 933d4ef
1 parent c166e65
commit 933d4ef
Show file tree

Hide file tree

Showing 2 changed files with 124 additions and 9 deletions.
diff --git a/crates/uv-workspace/src/pyproject_mut.rs b/crates/uv-workspace/src/pyproject_mut.rs
@@ -2,10 +2,11 @@ use std::path::Path;
 use std::str::FromStr;
 use std::{fmt, mem};
 
-use pep440_rs::{Version, VersionSpecifier, VersionSpecifiers};
-use pep508_rs::{ExtraName, MarkerTree, PackageName, Requirement, VersionOrUrl};
 use thiserror::Error;
 use toml_edit::{Array, DocumentMut, Item, RawString, Table, TomlError, Value};
+
+use pep440_rs::{Version, VersionSpecifier, VersionSpecifiers};
+use pep508_rs::{ExtraName, MarkerTree, PackageName, Requirement, VersionOrUrl};
 use uv_fs::PortablePath;
 
 use crate::pyproject::{DependencyType, Source};
@@ -196,7 +197,7 @@ impl PyProjectTomlMut {
             .doc()?
             .entry("optional-dependencies")
             .or_insert(Item::Table(Table::new()))
-            .as_table_mut()
+            .as_table_like_mut()
             .ok_or(Error::MalformedDependencies)?;
 
         let group = optional_dependencies
@@ -208,6 +209,8 @@ impl PyProjectTomlMut {
         let name = req.name.clone();
         let added = add_dependency(req, group, source.is_some())?;
 
+        optional_dependencies.fmt();
+
         if let Some(source) = source {
             self.add_source(&name, source)?;
         }
@@ -348,7 +351,11 @@ impl PyProjectTomlMut {
         let Some(dependencies) = self
             .doc_mut()?
             .and_then(|project| project.get_mut("dependencies"))
-            .map(|dependencies| dependencies.as_array_mut().ok_or(Error::MalformedSources))
+            .map(|dependencies| {
+                dependencies
+                    .as_array_mut()
+                    .ok_or(Error::MalformedDependencies)
+            })
             .transpose()?
         else {
             return Ok(Vec::new());
@@ -366,13 +373,17 @@ impl PyProjectTomlMut {
         let Some(dev_dependencies) = self
             .doc
             .get_mut("tool")
-            .map(|tool| tool.as_table_mut().ok_or(Error::MalformedSources))
+            .map(|tool| tool.as_table_mut().ok_or(Error::MalformedDependencies))
             .transpose()?
             .and_then(|tool| tool.get_mut("uv"))
-            .map(|tool_uv| tool_uv.as_table_mut().ok_or(Error::MalformedSources))
+            .map(|tool_uv| tool_uv.as_table_mut().ok_or(Error::MalformedDependencies))
             .transpose()?
             .and_then(|tool_uv| tool_uv.get_mut("dev-dependencies"))
-            .map(|dependencies| dependencies.as_array_mut().ok_or(Error::MalformedSources))
+            .map(|dependencies| {
+                dependencies
+                    .as_array_mut()
+                    .ok_or(Error::MalformedDependencies)
+            })
             .transpose()?
         else {
             return Ok(Vec::new());
@@ -394,10 +405,18 @@ impl PyProjectTomlMut {
         let Some(optional_dependencies) = self
             .doc_mut()?
             .and_then(|project| project.get_mut("optional-dependencies"))
-            .map(|extras| extras.as_table_mut().ok_or(Error::MalformedSources))
+            .map(|extras| {
+                extras
+                    .as_table_like_mut()
+                    .ok_or(Error::MalformedDependencies)
+            })
             .transpose()?
             .and_then(|extras| extras.get_mut(group.as_ref()))
-            .map(|dependencies| dependencies.as_array_mut().ok_or(Error::MalformedSources))
+            .map(|dependencies| {
+                dependencies
+                    .as_array_mut()
+                    .ok_or(Error::MalformedDependencies)
+            })
             .transpose()?
         else {
             return Ok(Vec::new());

diff --git a/crates/uv/tests/edit.rs b/crates/uv/tests/edit.rs
@@ -1373,6 +1373,102 @@ fn add_remove_optional() -> Result<()> {
     Ok(())
 }
 
+#[test]
+fn add_remove_inline_optional() -> Result<()> {
+    let context = TestContext::new("3.12");
+
+    let pyproject_toml = context.temp_dir.child("pyproject.toml");
+    pyproject_toml.write_str(indoc! {r#"
+        [project]
+        name = "project"
+        version = "0.1.0"
+        requires-python = ">=3.12"
+        dependencies = []
+        optional-dependencies = { io = [
+            "anyio==3.7.0",
+        ] }
+
+        [build-system]
+        requires = ["setuptools>=42", "wheel"]
+        build-backend = "setuptools.build_meta"
+    "#})?;
+
+    uv_snapshot!(context.filters(), context.add(&["typing-extensions"]).arg("--optional=types"), @r###"
+    success: false
+    exit_code: 2
+    ----- stdout -----
+
+    ----- stderr -----
+    Resolved 5 packages in [TIME]
+    error: Dependencies in `pyproject.toml` are malformed
+    "###);
+
+    let pyproject_toml = fs_err::read_to_string(context.temp_dir.join("pyproject.toml"))?;
+
+    insta::with_settings!({
+        filters => context.filters(),
+    }, {
+        assert_snapshot!(
+            pyproject_toml, @r###"
+        [project]
+        name = "project"
+        version = "0.1.0"
+        requires-python = ">=3.12"
+        dependencies = []
+        optional-dependencies = { io = [
+            "anyio==3.7.0",
+        ], types = [
+            "typing-extensions",
+        ] }
+
+        [build-system]
+        requires = ["setuptools>=42", "wheel"]
+        build-backend = "setuptools.build_meta"
+        "###
+        );
+    });
+
+    uv_snapshot!(context.filters(), context.remove(&["typing-extensions"]).arg("--optional=types"), @r###"
+    success: true
+    exit_code: 0
+    ----- stdout -----
+
+    ----- stderr -----
+    Resolved 4 packages in [TIME]
+    Prepared 4 packages in [TIME]
+    Installed 4 packages in [TIME]
+     + anyio==3.7.0
+     + idna==3.6
+     + project==0.1.0 (from file://[TEMP_DIR]/)
+     + sniffio==1.3.1
+    "###);
+
+    let pyproject_toml = fs_err::read_to_string(context.temp_dir.join("pyproject.toml"))?;
+
+    insta::with_settings!({
+        filters => context.filters(),
+    }, {
+        assert_snapshot!(
+            pyproject_toml, @r###"
+        [project]
+        name = "project"
+        version = "0.1.0"
+        requires-python = ">=3.12"
+        dependencies = []
+        optional-dependencies = { io = [
+            "anyio==3.7.0",
+        ], types = [] }
+
+        [build-system]
+        requires = ["setuptools>=42", "wheel"]
+        build-backend = "setuptools.build_meta"
+        "###
+        );
+    });
+
+    Ok(())
+}
+
 /// Add and remove a workspace dependency.
 #[test]
 fn add_remove_workspace() -> Result<()> {