From bdef77c3fe0f386b6d66a237f498e3d219c4e0c8 Mon Sep 17 00:00:00 2001
From: Christopher Tee <christopherfytee@gmail.com>
Date: Fri, 21 Mar 2025 22:16:15 +0800
Subject: [PATCH] Retain end-of-line comment position when adding dependency
 (#12360)

<!--
Thank you for contributing to uv! To help us out with reviewing, please
consider the following:

- Does this pull request include a summary of the change? (See below.)
- Does this pull request include a descriptive title?
- Does this pull request include references to any relevant issues?
-->

## Summary

<!-- What's the purpose of the change? What does it do, and why? -->

This fixes a case described in #12333, where trailing comments in
dependencies can be unexpectedly shifted when a new dependency is added.

Fixes #12333.

## Test Plan

<!-- How was it tested? -->

`cargo test` (Added a snapshot test)
---
 crates/uv-workspace/src/pyproject_mut.rs | 71 +++++++++++++++++-------
 crates/uv/tests/it/edit.rs               | 56 +++++++++++++++++++
 2 files changed, 106 insertions(+), 21 deletions(-)

diff --git a/crates/uv-workspace/src/pyproject_mut.rs b/crates/uv-workspace/src/pyproject_mut.rs
index 022c54d23..97b826ab9 100644
--- a/crates/uv-workspace/src/pyproject_mut.rs
+++ b/crates/uv-workspace/src/pyproject_mut.rs
@@ -1004,27 +1004,56 @@ pub fn add_dependency(
 
             let decor = value.decor_mut();
 
-            // If we're adding to the end of the list, treat trailing comments as leading comments
-            // on the added dependency.
-            //
-            // For example, given:
-            // ```toml
-            // dependencies = [
-            //     "anyio", # trailing comment
-            // ]
-            // ```
-            //
-            // If we add `flask` to the end, we want to retain the comment on `anyio`:
-            // ```toml
-            // dependencies = [
-            //     "anyio", # trailing comment
-            //     "flask",
-            // ]
-            // ```
-            if index == deps.len() {
-                decor.set_prefix(deps.trailing().clone());
-                deps.set_trailing("");
-            }
+            // Ensure comments remain on the correct line, post-insertion
+            match index {
+                val if val == deps.len() => {
+                    // If we're adding to the end of the list, treat trailing comments as leading comments
+                    // on the added dependency.
+                    //
+                    // For example, given:
+                    // ```toml
+                    // dependencies = [
+                    //     "anyio", # trailing comment
+                    // ]
+                    // ```
+                    //
+                    // If we add `flask` to the end, we want to retain the comment on `anyio`:
+                    // ```toml
+                    // dependencies = [
+                    //     "anyio", # trailing comment
+                    //     "flask",
+                    // ]
+                    // ```
+                    decor.set_prefix(deps.trailing().clone());
+                    deps.set_trailing("");
+                }
+                0 => {
+                    // If the dependency is prepended to a non-empty list, do nothing
+                }
+                val => {
+                    // Retain position of end-of-line comments when a dependency is inserted right below it.
+                    //
+                    // For example, given:
+                    // ```toml
+                    // dependencies = [
+                    //     "anyio", # end-of-line comment
+                    //     "flask",
+                    // ]
+                    // ```
+                    //
+                    // If we add `pydantic` (between `anyio` and `flask`), we want to retain the comment on `anyio`:
+                    // ```toml
+                    // dependencies = [
+                    //     "anyio", # end-of-line comment
+                    //     "pydantic",
+                    //     "flask",
+                    // ]
+                    // ```
+                    let targeted_decor = deps.get_mut(val).unwrap().decor_mut();
+                    decor.set_prefix(targeted_decor.prefix().unwrap().clone());
+                    targeted_decor.set_prefix(""); // Re-formatted later by `reformat_array_multiline`
+                }
+            };
 
             deps.insert_formatted(index, value);
 
diff --git a/crates/uv/tests/it/edit.rs b/crates/uv/tests/it/edit.rs
index fba4bc355..dbcf6c07d 100644
--- a/crates/uv/tests/it/edit.rs
+++ b/crates/uv/tests/it/edit.rs
@@ -8829,6 +8829,62 @@ fn add_preserves_end_of_line_comments() -> Result<()> {
     Ok(())
 }
 
+#[test]
+fn add_preserves_end_of_line_comment_on_non_last_deps() -> 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 = [ # comment
+            "anyio==3.7.0", # comment 1
+            "sniffio==1.3.1",
+        ]
+    "#})?;
+
+    uv_snapshot!(context.filters(), context.add().arg("requests==2.31.0"), @r"
+    success: true
+    exit_code: 0
+    ----- stdout -----
+
+    ----- stderr -----
+    Resolved 8 packages in [TIME]
+    Prepared 7 packages in [TIME]
+    Installed 7 packages in [TIME]
+     + anyio==3.7.0
+     + certifi==2024.2.2
+     + charset-normalizer==3.3.2
+     + idna==3.6
+     + requests==2.31.0
+     + sniffio==1.3.1
+     + urllib3==2.2.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 = [ # comment
+            "anyio==3.7.0", # comment 1
+            "requests==2.31.0",
+            "sniffio==1.3.1",
+        ]
+        "###
+        );
+    });
+    Ok(())
+}
+
 #[test]
 fn add_direct_url_subdirectory() -> Result<()> {
     let context = TestContext::new("3.12");