Normalize extra and group names in uv add and uv remove (#12586)

## Summary

Closes https://github.com/astral-sh/uv/issues/12585.

crates/uv-workspace/src/pyproject_mut.rs

@@ -406,11 +406,24 @@ impl PyProjectTomlMut {
             .as_table_like_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)?;
+        // Try to find the existing group.
+        let existing_group = optional_dependencies.iter_mut().find_map(|(key, value)| {
+            if ExtraName::from_str(key.get()).is_ok_and(|g| g == *group) {
+                Some(value)
+            } else {
+                None
+            }
+        });
+
+        // If the group doesn't exist, create it.
+        let group = match existing_group {
+            Some(value) => value,
+            None => optional_dependencies
+                .entry(group.as_ref())
+                .or_insert(Item::Value(Value::Array(Array::new()))),
+        }
+        .as_array_mut()
+        .ok_or(Error::MalformedDependencies)?;
 
         let added = add_dependency(req, group, source.is_some())?;
 
@@ -457,11 +470,24 @@ impl PyProjectTomlMut {
             .map(|k| k.get())
             .is_sorted();
 
-        let group = dependency_groups
-            .entry(group.as_ref())
-            .or_insert(Item::Value(Value::Array(Array::new())))
-            .as_array_mut()
-            .ok_or(Error::MalformedDependencies)?;
+        // Try to find the existing group.
+        let existing_group = dependency_groups.iter_mut().find_map(|(key, value)| {
+            if GroupName::from_str(key.get()).is_ok_and(|g| g == *group) {
+                Some(value)
+            } else {
+                None
+            }
+        });
+
+        // If the group doesn't exist, create it.
+        let group = match existing_group {
+            Some(value) => value,
+            None => dependency_groups
+                .entry(group.as_ref())
+                .or_insert(Item::Value(Value::Array(Array::new()))),
+        }
+        .as_array_mut()
+        .ok_or(Error::MalformedDependencies)?;
 
         let added = add_dependency(req, group, source.is_some())?;
 
@@ -573,11 +599,24 @@ impl PyProjectTomlMut {
             .as_table_like_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)?;
+        // Try to find the existing group.
+        let existing_group = optional_dependencies.iter_mut().find_map(|(key, value)| {
+            if ExtraName::from_str(key.get()).is_ok_and(|g| g == *group) {
+                Some(value)
+            } else {
+                None
+            }
+        });
+
+        // If the group doesn't exist, create it.
+        let group = match existing_group {
+            Some(value) => value,
+            None => optional_dependencies
+                .entry(group.as_ref())
+                .or_insert(Item::Value(Value::Array(Array::new()))),
+        }
+        .as_array_mut()
+        .ok_or(Error::MalformedDependencies)?;
 
         let Some(req) = group.get(index) else {
             return Err(Error::MissingDependency(index));
@@ -610,11 +649,24 @@ impl PyProjectTomlMut {
             .as_table_like_mut()
             .ok_or(Error::MalformedDependencies)?;
 
-        let group = dependency_groups
-            .entry(group.as_ref())
-            .or_insert(Item::Value(Value::Array(Array::new())))
-            .as_array_mut()
-            .ok_or(Error::MalformedDependencies)?;
+        // Try to find the existing group.
+        let existing_group = dependency_groups.iter_mut().find_map(|(key, value)| {
+            if GroupName::from_str(key.get()).is_ok_and(|g| g == *group) {
+                Some(value)
+            } else {
+                None
+            }
+        });
+
+        // If the group doesn't exist, create it.
+        let group = match existing_group {
+            Some(value) => value,
+            None => dependency_groups
+                .entry(group.as_ref())
+                .or_insert(Item::Value(Value::Array(Array::new()))),
+        }
+        .as_array_mut()
+        .ok_or(Error::MalformedDependencies)?;
 
         let Some(req) = group.get(index) else {
             return Err(Error::MissingDependency(index));
@@ -724,7 +776,15 @@ impl PyProjectTomlMut {
                     .ok_or(Error::MalformedDependencies)
             })
             .transpose()?
-            .and_then(|extras| extras.get_mut(group.as_ref()))
+            .and_then(|extras| {
+                extras.iter_mut().find_map(|(key, value)| {
+                    if ExtraName::from_str(key.get()).is_ok_and(|g| g == *group) {
+                        Some(value)
+                    } else {
+                        None
+                    }
+                })
+            })
             .map(|dependencies| {
                 dependencies
                     .as_array_mut()
@@ -757,7 +817,15 @@ impl PyProjectTomlMut {
                     .ok_or(Error::MalformedDependencies)
             })
             .transpose()?
-            .and_then(|groups| groups.get_mut(group.as_ref()))
+            .and_then(|groups| {
+                groups.iter_mut().find_map(|(key, value)| {
+                    if GroupName::from_str(key.get()).is_ok_and(|g| g == *group) {
+                        Some(value)
+                    } else {
+                        None
+                    }
+                })
+            })
             .map(|dependencies| {
                 dependencies
                     .as_array_mut()

crates/uv/tests/it/edit.rs

@@ -4821,6 +4821,152 @@ fn add_group() -> Result<()> {
     Ok(())
 }
 
+/// Normalize group names when adding or removing.
+#[test]
+fn add_group_normalize() -> 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 = []
+
+        [dependency-groups]
+        cloud_export_to_parquet = [
+            "anyio==3.7.0",
+        ]
+    "#})?;
+
+    // Add with a non-normalized group name.
+    uv_snapshot!(context.filters(), context.add().arg("iniconfig").arg("--group").arg("cloud_export_to_parquet"), @r"
+    success: true
+    exit_code: 0
+    ----- stdout -----
+
+    ----- stderr -----
+    Resolved 5 packages in [TIME]
+    Prepared 4 packages in [TIME]
+    Installed 4 packages in [TIME]
+     + anyio==3.7.0
+     + idna==3.6
+     + iniconfig==2.0.0
+     + sniffio==1.3.1
+    ");
+
+    let pyproject_toml = context.read("pyproject.toml");
+
+    assert_snapshot!(pyproject_toml, @r#"
+    [project]
+    name = "project"
+    version = "0.1.0"
+    requires-python = ">=3.12"
+    dependencies = []
+
+    [dependency-groups]
+    cloud_export_to_parquet = [
+        "anyio==3.7.0",
+        "iniconfig>=2.0.0",
+    ]
+    "#
+    );
+
+    // Add with a normalized group name (which doesn't match the `pyproject.toml`).
+    uv_snapshot!(context.filters(), context.add().arg("typing-extensions").arg("--group").arg("cloud-export-to-parquet"), @r"
+    success: true
+    exit_code: 0
+    ----- stdout -----
+
+    ----- stderr -----
+    Resolved 6 packages in [TIME]
+    Prepared 1 package in [TIME]
+    Installed 1 package in [TIME]
+     + typing-extensions==4.10.0
+    ");
+
+    let pyproject_toml = context.read("pyproject.toml");
+
+    assert_snapshot!(pyproject_toml, @r#"
+    [project]
+    name = "project"
+    version = "0.1.0"
+    requires-python = ">=3.12"
+    dependencies = []
+
+    [dependency-groups]
+    cloud_export_to_parquet = [
+        "anyio==3.7.0",
+        "iniconfig>=2.0.0",
+        "typing-extensions>=4.10.0",
+    ]
+    "#
+    );
+
+    // Remove with a non-normalized group name.
+    uv_snapshot!(context.filters(), context.remove().arg("iniconfig").arg("--group").arg("cloud_export_to_parquet"), @r"
+    success: true
+    exit_code: 0
+    ----- stdout -----
+
+    ----- stderr -----
+    Resolved 5 packages in [TIME]
+    Uninstalled 5 packages in [TIME]
+     - anyio==3.7.0
+     - idna==3.6
+     - iniconfig==2.0.0
+     - sniffio==1.3.1
+     - typing-extensions==4.10.0
+    ");
+
+    let pyproject_toml = context.read("pyproject.toml");
+
+    assert_snapshot!(pyproject_toml, @r#"
+    [project]
+    name = "project"
+    version = "0.1.0"
+    requires-python = ">=3.12"
+    dependencies = []
+
+    [dependency-groups]
+    cloud_export_to_parquet = [
+        "anyio==3.7.0",
+        "typing-extensions>=4.10.0",
+    ]
+    "#
+    );
+
+    // Remove with a normalized group name (which doesn't match the `pyproject.toml`).
+    uv_snapshot!(context.filters(), context.remove().arg("typing-extensions").arg("--group").arg("cloud-export-to-parquet"), @r"
+    success: true
+    exit_code: 0
+    ----- stdout -----
+
+    ----- stderr -----
+    Resolved 4 packages in [TIME]
+    Audited in [TIME]
+    ");
+
+    let pyproject_toml = context.read("pyproject.toml");
+
+    assert_snapshot!(pyproject_toml, @r#"
+    [project]
+    name = "project"
+    version = "0.1.0"
+    requires-python = ">=3.12"
+    dependencies = []
+
+    [dependency-groups]
+    cloud_export_to_parquet = [
+        "anyio==3.7.0",
+    ]
+    "#
+    );
+
+    Ok(())
+}
+
 /// Add a requirement to a dependency group (sorted before the other groups).
 #[test]
 fn add_group_before_commented_groups() -> Result<()> {
@@ -10347,3 +10493,146 @@ fn add_ambiguous() -> Result<()> {
 
     Ok(())
 }
+
+/// Normalize extra names when adding or removing.
+#[test]
+fn add_optional_normalize() -> 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 = []
+
+        [project.optional-dependencies]
+        cloud_export_to_parquet = [
+            "anyio==3.7.0",
+        ]
+    "#})?;
+
+    // Add with a non-normalized group name.
+    uv_snapshot!(context.filters(), context.add().arg("iniconfig").arg("--optional").arg("cloud_export_to_parquet"), @r"
+    success: true
+    exit_code: 0
+    ----- stdout -----
+
+    ----- stderr -----
+    Resolved 5 packages in [TIME]
+    Prepared 4 packages in [TIME]
+    Installed 4 packages in [TIME]
+     + anyio==3.7.0
+     + idna==3.6
+     + iniconfig==2.0.0
+     + sniffio==1.3.1
+    ");
+
+    let pyproject_toml = context.read("pyproject.toml");
+
+    assert_snapshot!(pyproject_toml, @r#"
+    [project]
+    name = "project"
+    version = "0.1.0"
+    requires-python = ">=3.12"
+    dependencies = []
+
+    [project.optional-dependencies]
+    cloud_export_to_parquet = [
+        "anyio==3.7.0",
+        "iniconfig>=2.0.0",
+    ]
+    "#
+    );
+
+    // Add with a normalized group name (which doesn't match the `pyproject.toml`).
+    uv_snapshot!(context.filters(), context.add().arg("typing-extensions").arg("--optional").arg("cloud-export-to-parquet"), @r"
+    success: true
+    exit_code: 0
+    ----- stdout -----
+
+    ----- stderr -----
+    Resolved 6 packages in [TIME]
+    Prepared 1 package in [TIME]
+    Installed 1 package in [TIME]
+     + typing-extensions==4.10.0
+    ");
+
+    let pyproject_toml = context.read("pyproject.toml");
+
+    assert_snapshot!(pyproject_toml, @r#"
+    [project]
+    name = "project"
+    version = "0.1.0"
+    requires-python = ">=3.12"
+    dependencies = []
+
+    [project.optional-dependencies]
+    cloud_export_to_parquet = [
+        "anyio==3.7.0",
+        "iniconfig>=2.0.0",
+        "typing-extensions>=4.10.0",
+    ]
+    "#
+    );
+
+    // Remove with a non-normalized group name.
+    uv_snapshot!(context.filters(), context.remove().arg("iniconfig").arg("--optional").arg("cloud_export_to_parquet"), @r"
+    success: true
+    exit_code: 0
+    ----- stdout -----
+
+    ----- stderr -----
+    Resolved 5 packages in [TIME]
+    Uninstalled 1 package in [TIME]
+     - iniconfig==2.0.0
+    ");
+
+    let pyproject_toml = context.read("pyproject.toml");
+
+    assert_snapshot!(pyproject_toml, @r#"
+    [project]
+    name = "project"
+    version = "0.1.0"
+    requires-python = ">=3.12"
+    dependencies = []
+
+    [project.optional-dependencies]
+    cloud_export_to_parquet = [
+        "anyio==3.7.0",
+        "typing-extensions>=4.10.0",
+    ]
+    "#
+    );
+
+    // Remove with a normalized group name (which doesn't match the `pyproject.toml`).
+    uv_snapshot!(context.filters(), context.remove().arg("typing-extensions").arg("--optional").arg("cloud-export-to-parquet"), @r"
+    success: true
+    exit_code: 0
+    ----- stdout -----
+
+    ----- stderr -----
+    Resolved 4 packages in [TIME]
+    Uninstalled 1 package in [TIME]
+     - typing-extensions==4.10.0
+    ");
+
+    let pyproject_toml = context.read("pyproject.toml");
+
+    assert_snapshot!(pyproject_toml, @r#"
+    [project]
+    name = "project"
+    version = "0.1.0"
+    requires-python = ">=3.12"
+    dependencies = []
+
+    [project.optional-dependencies]
+    cloud_export_to_parquet = [
+        "anyio==3.7.0",
+    ]
+    "#
+    );
+
+    Ok(())
+}