From 14a1ea460d297008762be2ce45e3bb1ee041fee1 Mon Sep 17 00:00:00 2001
From: Alexander Gherm <aleksandr_germ@epam.com>
Date: Tue, 16 Jul 2024 01:13:22 +0200
Subject: [PATCH] Rework reformatting in PyProjectTomlMut to respect original
 indentation (#5075)

## Summary

So this PR introduces change to how `Array` of dependencies
representation is reformatted while `PyProjectTomlMut` is manipulated.
These changes are here for it to respect the original indentation.
Closes https://github.com/astral-sh/uv/issues/5009

## Test Plan
Using `pyproject.toml` like
```
[project]
name = "project"
version = "0.1.0"
requires-python = ">=3.12"
dependencies = [
  "requests"
]
```
Executed
```
$ uv add httpx
```
And expected in `pyproject.toml`
```
[project]
name = "project"
version = "0.1.0"
requires-python = ">=3.12"
dependencies = [
  "requests",
  "httpx",
]
```
Preserving original indentation
---
 crates/uv-distribution/src/pyproject_mut.rs |  26 ++++-
 crates/uv/tests/edit.rs                     | 114 ++++++++++++++++++++
 2 files changed, 138 insertions(+), 2 deletions(-)

diff --git a/crates/uv-distribution/src/pyproject_mut.rs b/crates/uv-distribution/src/pyproject_mut.rs
index 24ac450a4..39b7d5ca5 100644
--- a/crates/uv-distribution/src/pyproject_mut.rs
+++ b/crates/uv-distribution/src/pyproject_mut.rs
@@ -424,14 +424,36 @@ fn reformat_array_multiline(deps: &mut Array) {
             })
     }
 
+    let mut indentation_prefix = None;
+
     for item in deps.iter_mut() {
         let decor = item.decor_mut();
         let mut prefix = String::new();
+        // calculating the indentation prefix as the indentation of the first dependency entry
+        if indentation_prefix.is_none() {
+            let decor_prefix = decor
+                .prefix()
+                .and_then(|s| s.as_str())
+                .map(|s| s.split('#').next().unwrap_or("").to_string())
+                .unwrap_or(String::new())
+                .trim_start_matches('\n')
+                .to_string();
+
+            // if there is no indentation then apply a default one
+            indentation_prefix = Some(if decor_prefix.is_empty() {
+                "    ".to_string()
+            } else {
+                decor_prefix
+            });
+        }
+
+        let indentation_prefix_str = format!("\n{}", indentation_prefix.as_ref().unwrap());
+
         for comment in find_comments(decor.prefix()).chain(find_comments(decor.suffix())) {
-            prefix.push_str("\n    ");
+            prefix.push_str(&indentation_prefix_str);
             prefix.push_str(comment);
         }
-        prefix.push_str("\n    ");
+        prefix.push_str(&indentation_prefix_str);
         decor.set_prefix(prefix);
         decor.set_suffix("");
     }
diff --git a/crates/uv/tests/edit.rs b/crates/uv/tests/edit.rs
index 660643dc1..26c2679be 100644
--- a/crates/uv/tests/edit.rs
+++ b/crates/uv/tests/edit.rs
@@ -1710,3 +1710,117 @@ fn remove_registry() -> Result<()> {
 
     Ok(())
 }
+
+#[test]
+fn add_preserves_indentation_in_pyproject_toml() -> 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 = [
+          "anyio==3.7.0"
+        ]
+    "#})?;
+
+    uv_snapshot!(context.filters(), context.add(&["requests==2.31.0"]), @r###"
+    success: true
+    exit_code: 0
+    ----- stdout -----
+
+    ----- stderr -----
+    warning: `uv add` is experimental and may change without warning.
+    Resolved 8 packages in [TIME]
+    Prepared 8 packages in [TIME]
+    Installed 8 packages in [TIME]
+     + anyio==3.7.0
+     + certifi==2024.2.2
+     + charset-normalizer==3.3.2
+     + idna==3.6
+     + project==0.1.0 (from file://[TEMP_DIR]/)
+     + requests==2.31.0
+     + sniffio==1.3.1
+     + urllib3==2.2.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 = [
+          "anyio==3.7.0",
+          "requests==2.31.0",
+        ]
+        "###
+        );
+    });
+    Ok(())
+}
+
+#[test]
+fn add_puts_default_indentation_in_pyproject_toml_if_not_observed() -> 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 = ["anyio==3.7.0"]
+    "#})?;
+
+    uv_snapshot!(context.filters(), context.add(&["requests==2.31.0"]), @r###"
+    success: true
+    exit_code: 0
+    ----- stdout -----
+
+    ----- stderr -----
+    warning: `uv add` is experimental and may change without warning.
+    Resolved 8 packages in [TIME]
+    Prepared 8 packages in [TIME]
+    Installed 8 packages in [TIME]
+     + anyio==3.7.0
+     + certifi==2024.2.2
+     + charset-normalizer==3.3.2
+     + idna==3.6
+     + project==0.1.0 (from file://[TEMP_DIR]/)
+     + requests==2.31.0
+     + sniffio==1.3.1
+     + urllib3==2.2.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 = [
+            "anyio==3.7.0",
+            "requests==2.31.0",
+        ]
+        "###
+        );
+    });
+    Ok(())
+}