Preserve end-of-line comments on previous entries when removing dependencies

crates/uv-workspace/src/pyproject_mut.rs

@@ -1545,11 +1545,51 @@ fn update_requirement(old: &mut Requirement, new: &Requirement, has_source: bool
 
 /// Removes all occurrences of dependencies with the given name from the given `deps` array.
 fn remove_dependency(name: &PackageName, deps: &mut Array) -> Vec<Requirement> {
-    // Remove matching dependencies.
+    // Remove in reverse to preserve indices. Before each removal, transfer the item's
+    // prefix (which may contain end-of-line comments belonging to the previous line) to
+    // the next item or array trailing so comments are not lost.
+    //
+    // For example, in:
+    // ```toml
+    // dependencies = [
+    //     "numpy>=2.4.3", # essential comment
+    //     "requests>=2.32.5",
+    // ]
+    // ```
+    //
+    // The comment `# essential comment` is stored by `toml_edit` in the prefix of
+    // `requests`. When `requests` is removed, we transfer it so it remains on the
+    // `numpy` line.
     let removed = find_dependencies(name, None, deps)
         .into_iter()
-        .rev() // Reverse to preserve indices as we remove them.
+        .rev()
         .filter_map(|(i, _)| {
+            if let Some(prefix) = deps
+                .get(i)
+                .and_then(|item| item.decor().prefix().and_then(|s| s.as_str()))
+                .filter(|s| !s.is_empty())
+            {
+                let prefix = prefix.to_string();
+                if let Some(next) = deps.get(i + 1)
+                    && let Some(existing) = next.decor().prefix().and_then(|s| s.as_str())
+                {
+                    // Transfer removed item's prefix to the next item's prefix.
+                    let existing = existing.to_string();
+                    deps.get_mut(i + 1)
+                        .unwrap()
+                        .decor_mut()
+                        .set_prefix(format!("{prefix}{existing}"));
+                } else if let Some(next) = deps.get_mut(i + 1) {
+                    // Next item exists but has no prefix; use ours directly.
+                    next.decor_mut().set_prefix(&prefix);
+                } else if let Some(existing) = deps.trailing().as_str() {
+                    // No next item; move comments to the array trailing.
+                    deps.set_trailing(format!("{prefix}{existing}"));
+                } else {
+                    deps.set_trailing(&prefix);
+                }
+            }
+
             deps.remove(i)
                 .as_str()
                 .and_then(|req| Requirement::from_str(req).ok())
@@ -1750,9 +1790,11 @@ fn split_specifiers(req: &str) -> (&str, &str) {
 
 #[cfg(test)]
 mod test {
-    use super::{AddBoundsKind, reformat_array_multiline, split_specifiers};
+    use super::{AddBoundsKind, reformat_array_multiline, remove_dependency, split_specifiers};
+    use insta::assert_snapshot;
     use std::str::FromStr;
     use toml_edit::DocumentMut;
+    use uv_normalize::PackageName;
     use uv_pep440::Version;
 
     #[test]
@@ -1927,4 +1969,191 @@ dependencies = [
             assert_eq!(actual, expected, "{version}");
         }
     }
+
+    #[test]
+    fn remove_preserves_end_of_line_comment_on_previous_item() {
+        let toml = r#"
+[project]
+dependencies = [
+    "numpy>=2.4.3", # this comment is clearly essential
+    "requests>=2.32.5",
+]
+"#;
+        let mut doc: DocumentMut = toml.parse().unwrap();
+        let deps = doc["project"]["dependencies"]
+            .as_array_mut()
+            .expect("dependencies array");
+
+        let name = PackageName::from_str("requests").unwrap();
+        remove_dependency(&name, deps);
+
+        assert_snapshot!(
+            doc.to_string(),
+            @r#"
+[project]
+dependencies = [
+    "numpy>=2.4.3", # this comment is clearly essential
+]
+"#
+        );
+    }
+
+    #[test]
+    fn remove_preserves_end_of_line_comment_on_previous_item_middle() {
+        let toml = r#"
+[project]
+dependencies = [
+    "numpy>=2.4.3", # numpy comment
+    "requests>=2.32.5",
+    "flask>=3.0.0",
+]
+"#;
+        let mut doc: DocumentMut = toml.parse().unwrap();
+        let deps = doc["project"]["dependencies"]
+            .as_array_mut()
+            .expect("dependencies array");
+
+        let name = PackageName::from_str("requests").unwrap();
+        remove_dependency(&name, deps);
+
+        assert_snapshot!(
+            doc.to_string(),
+            @r#"
+[project]
+dependencies = [
+    "numpy>=2.4.3", # numpy comment
+    "flask>=3.0.0",
+]
+"#
+        );
+    }
+
+    #[test]
+    fn remove_preserves_own_line_comment_above_removed_item() {
+        let toml = r#"
+[project]
+dependencies = [
+    "numpy>=2.4.3",
+    # This is a comment about requests
+    "requests>=2.32.5",
+]
+"#;
+        let mut doc: DocumentMut = toml.parse().unwrap();
+        let deps = doc["project"]["dependencies"]
+            .as_array_mut()
+            .expect("dependencies array");
+
+        let name = PackageName::from_str("requests").unwrap();
+        remove_dependency(&name, deps);
+
+        assert_snapshot!(
+            doc.to_string(),
+            @r#"
+[project]
+dependencies = [
+    "numpy>=2.4.3",
+    # This is a comment about requests
+]
+"#
+        );
+    }
+
+    #[test]
+    fn remove_item_with_trailing_comment_last() {
+        // When the removed item itself has an end-of-line comment and is the last item,
+        // toml_edit stores the comment in the array trailing. The comment is preserved
+        // (as an own-line comment in the trailing section) but moves position since it
+        // can no longer be on the removed item's line.
+        let toml = r#"
+[project]
+dependencies = [
+    "requests>=2.32.5",
+    "numpy>=2.4.3", # comment on numpy
+]
+"#;
+        let mut doc: DocumentMut = toml.parse().unwrap();
+        let deps = doc["project"]["dependencies"]
+            .as_array_mut()
+            .expect("dependencies array");
+
+        let name = PackageName::from_str("numpy").unwrap();
+        remove_dependency(&name, deps);
+
+        assert_snapshot!(
+            doc.to_string(),
+            @r#"
+[project]
+dependencies = [
+    "requests>=2.32.5",
+    # comment on numpy
+]
+"#
+        );
+    }
+
+    #[test]
+    fn remove_item_with_trailing_comment_middle() {
+        // When the removed item has an end-of-line comment and is in the middle,
+        // toml_edit stores the comment in the next item's prefix. After removal,
+        // reformat_array_multiline repositions it as an own-line comment.
+        let toml = r#"
+[project]
+dependencies = [
+    "requests>=2.32.5",
+    "numpy>=2.4.3", # comment on numpy
+    "flask>=3.0.0",
+]
+"#;
+        let mut doc: DocumentMut = toml.parse().unwrap();
+        let deps = doc["project"]["dependencies"]
+            .as_array_mut()
+            .expect("dependencies array");
+
+        let name = PackageName::from_str("numpy").unwrap();
+        remove_dependency(&name, deps);
+
+        assert_snapshot!(
+            doc.to_string(),
+            @r#"
+[project]
+dependencies = [
+    "requests>=2.32.5",
+    # comment on numpy
+    "flask>=3.0.0",
+]
+"#
+        );
+    }
+
+    #[test]
+    fn remove_multiple_adjacent_matches_preserves_comment_order() {
+        let toml = r#"
+[project]
+dependencies = [
+    "iniconfig>=2.0.0", # comment on iniconfig
+    "typing-extensions>=4.0.0 ; python_version < '3.11'", # comment on first typing-extensions
+    "typing-extensions>=4.0.0 ; python_version >= '3.11'",
+    "sniffio>=1.3.0",
+]
+"#;
+        let mut doc: DocumentMut = toml.parse().unwrap();
+        let deps = doc["project"]["dependencies"]
+            .as_array_mut()
+            .expect("dependencies array");
+
+        let name = PackageName::from_str("typing-extensions").unwrap();
+        remove_dependency(&name, deps);
+
+        assert_snapshot!(
+            doc.to_string(),
+            @r#"
+[project]
+dependencies = [
+    "iniconfig>=2.0.0", # comment on iniconfig
+    # comment on first typing-extensions
+    "sniffio>=1.3.0",
+]
+"#
+        );
+    }
 }

crates/uv/tests/it/edit.rs

@@ -12587,6 +12587,115 @@ fn remove_all_with_comments() -> Result<()> {
     Ok(())
 }
 
+/// Removing a dependency should preserve end-of-line comments on nearby lines.
+///
+/// See: <https://github.com/astral-sh/uv/issues/18555>
+#[test]
+fn remove_preserves_nearby_end_of_line_comments() -> Result<()> {
+    let context = uv_test::test_context!("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 = [
+            "iniconfig>=2.0.0", # this comment is clearly essential
+            "typing-extensions>=4.0.0",
+        ]
+    "#})?;
+
+    uv_snapshot!(context.filters(), context.remove().arg("typing-extensions"), @"
+    success: true
+    exit_code: 0
+    ----- stdout -----
+
+    ----- stderr -----
+    Resolved 2 packages in [TIME]
+    Prepared 1 package in [TIME]
+    Installed 1 package in [TIME]
+     + iniconfig==2.0.0
+    ");
+
+    let pyproject_toml = context.read("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 = [
+            "iniconfig>=2.0.0", # this comment is clearly essential
+        ]
+        "#
+        );
+    });
+
+    Ok(())
+}
+
+/// Removing multiple adjacent matching dependencies should preserve comment order.
+///
+/// See: <https://github.com/astral-sh/uv/issues/18555>
+#[test]
+fn remove_preserves_comment_order_for_multiple_adjacent_matches() -> Result<()> {
+    let context = uv_test::test_context!("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 = [
+            "iniconfig>=2.0.0", # comment on iniconfig
+            "typing-extensions>=4.0.0 ; python_version < '3.11'", # comment on first typing-extensions
+            "typing-extensions>=4.0.0 ; python_version >= '3.11'",
+            "sniffio>=1.3.0",
+        ]
+    "#})?;
+
+    uv_snapshot!(context.filters(), context.remove().arg("typing-extensions"), @"
+    success: true
+    exit_code: 0
+    ----- stdout -----
+
+    ----- stderr -----
+    Resolved 3 packages in [TIME]
+    Prepared 2 packages in [TIME]
+    Installed 2 packages in [TIME]
+     + iniconfig==2.0.0
+     + sniffio==1.3.1
+    ");
+
+    let pyproject_toml = context.read("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 = [
+            "iniconfig>=2.0.0", # comment on iniconfig
+            # comment on first typing-extensions
+            "sniffio>=1.3.0",
+        ]
+        "#
+        );
+    });
+
+    Ok(())
+}
+
 /// If multiple indexes are provided on the CLI, the first-provided index should take precedence
 /// during resolution, and should appear first in the `pyproject.toml` file.
 ///