Fix trailing newline bug (#1323)

Co-authored-by: William Woodruff <william@yossarian.net>

crates/yamlpatch/src/lib.rs

@@ -209,6 +209,35 @@ pub fn apply_yaml_patches(
     Ok(next_document)
 }
 
+/// Compute the replacement range for a byte span, including any trailing newline
+/// that exists immediately after the span in the document.
+fn compute_replacement_range_for_span(
+    byte_span: (usize, usize),
+    document: &yamlpath::Document,
+) -> std::ops::Range<usize> {
+    let start = byte_span.0;
+    let mut end = byte_span.1;
+    let content = document.source();
+
+    // Only include trailing newline if it's the last character in the document.
+    // This preserves document-level trailing newlines without interfering with
+    // newlines that are part of the document structure.
+    if end + 1 == content.len() && content.as_bytes()[end] == b'\n' {
+        end += 1;
+    }
+
+    start..end
+}
+
+/// Compute the replacement range for a feature, including any trailing newline
+/// that exists immediately after the feature in the document.
+fn compute_replacement_range(
+    feature: &yamlpath::Feature,
+    document: &yamlpath::Document,
+) -> std::ops::Range<usize> {
+    compute_replacement_range_for_span(feature.location.byte_span, document)
+}
+
 /// Apply a single YAML patch operation
 fn apply_single_patch(
     document: &yamlpath::Document,
@@ -244,10 +273,8 @@ fn apply_single_patch(
 
             // Finally, put our patch back into the overall content.
             let mut patched_content = content.to_string();
-            patched_content.replace_range(
-                feature.location.byte_span.0..feature.location.byte_span.1,
-                &patched_feature,
-            );
+            let replacement_range = compute_replacement_range(&feature, document);
+            patched_content.replace_range(replacement_range, &patched_feature);
 
             yamlpath::Document::new(patched_content).map_err(Error::from)
         }
@@ -272,8 +299,9 @@ fn apply_single_patch(
             };
 
             let mut result = content.to_string();
-            let span = comment_feature.location.byte_span;
-            result.replace_range(span.0..span.1, new);
+            let replacement_range =
+                compute_replacement_range_for_span(comment_feature.location.byte_span, document);
+            result.replace_range(replacement_range, new);
 
             yamlpath::Document::new(result).map_err(Error::from)
         }
@@ -300,7 +328,10 @@ fn apply_single_patch(
 
             // Replace the content
             let mut result = content.to_string();
-            result.replace_range(start_span..end_span, &replacement);
+            // Extend end_span to include trailing newline if present
+            let replacement_range =
+                compute_replacement_range_for_span((start_span, end_span), document);
+            result.replace_range(replacement_range, &replacement);
 
             yamlpath::Document::new(result).map_err(Error::from)
         }
@@ -350,10 +381,9 @@ fn apply_single_patch(
 
             // Replace the content in the document
             let mut result = content.to_string();
-            result.replace_range(
-                feature.location.byte_span.0..feature.location.byte_span.1,
-                &updated_feature,
-            );
+            let replacement_range = compute_replacement_range(&feature, document);
+            result.replace_range(replacement_range, &updated_feature);
+
             yamlpath::Document::new(result).map_err(Error::from)
         }
         Op::MergeInto { key, updates } => {

crates/yamlpatch/tests/unit_tests.rs

@@ -1,3 +1,4 @@
+use std::borrow::Cow;
 use yamlpatch::*;
 use yamlpath::route;
 
@@ -2308,3 +2309,241 @@ jobs:
                   persist-credentials: false
         "#);
 }
+
+#[test]
+fn test_preserve_trailing_newline_when_adding_at_end() {
+    // Test case for dependabot cooldown: adding cooldown config at the end of
+    // an ecosystem stanza should preserve the trailing newline
+    let original = r#"version: 2
+
+updates:
+  - package-ecosystem: pip
+    directory: /
+    schedule:
+      interval: daily
+    labels:
+      - A-deps
+"#;
+
+    let operations = vec![Patch {
+        route: route!("updates", 0),
+        operation: Op::Add {
+            key: "cooldown".to_string(),
+            value: serde_yaml::Value::Mapping({
+                let mut map = serde_yaml::Mapping::new();
+                map.insert(
+                    serde_yaml::Value::String("default-days".to_string()),
+                    serde_yaml::Value::Number(7.into()),
+                );
+                map
+            }),
+        },
+    }];
+
+    let result =
+        apply_yaml_patches(&yamlpath::Document::new(original).unwrap(), &operations).unwrap();
+
+    insta::assert_snapshot!(result.source(), @r#"
+        version: 2
+
+        updates:
+          - package-ecosystem: pip
+            directory: /
+            schedule:
+              interval: daily
+            labels:
+              - A-deps
+            cooldown:
+              default-days: 7
+        "#);
+}
+
+#[test]
+fn test_preserve_trailing_newline_replace_at_end() {
+    // Test Replace operation preserves trailing newline when replacing value at end of document
+    let original = r#"name: Test
+version: 1.0
+"#;
+
+    let operations = vec![Patch {
+        route: route!("version"),
+        operation: Op::Replace(serde_yaml::Value::String("2.0".to_string())),
+    }];
+
+    let result =
+        apply_yaml_patches(&yamlpath::Document::new(original).unwrap(), &operations).unwrap();
+
+    insta::assert_snapshot!(result.source(), @r#"
+        name: Test
+        version: '2.0'
+        "#);
+}
+
+#[test]
+fn test_preserve_trailing_newline_replace_comment_at_end() {
+    // Test ReplaceComment operation preserves trailing newline when replacing comment at end
+    let original = r#"name: Test
+version: 1.0  # old version
+"#;
+
+    let operations = vec![Patch {
+        route: route!("version"),
+        operation: Op::ReplaceComment {
+            new: Cow::Owned("# updated version".to_string()),
+        },
+    }];
+
+    let result =
+        apply_yaml_patches(&yamlpath::Document::new(original).unwrap(), &operations).unwrap();
+
+    insta::assert_snapshot!(result.source(), @r#"
+        name: Test
+        version: 1.0  # updated version
+        "#);
+}
+
+#[test]
+fn test_preserve_trailing_newline_rewrite_fragment_at_end() {
+    // Test RewriteFragment operation preserves trailing newline when rewriting at end of document
+    let original = r#"run: |
+  echo "Hello ${{ env.NAME }}"
+"#;
+
+    let operations = vec![Patch {
+        route: route!("run"),
+        operation: Op::RewriteFragment {
+            from: subfeature::Subfeature::new(0, "${{ env.NAME }}"),
+            to: Cow::Borrowed("${NAME}"),
+        },
+    }];
+
+    let result =
+        apply_yaml_patches(&yamlpath::Document::new(original).unwrap(), &operations).unwrap();
+
+    insta::assert_snapshot!(result.source(), @r#"
+        run: |
+          echo "Hello ${NAME}"
+        "#);
+}
+
+#[test]
+fn test_preserve_trailing_newline_add_simple_at_end() {
+    // Test Add operation preserves trailing newline when adding to mapping at end of document
+    let original = r#"name: Test
+key: value
+"#;
+
+    let operations = vec![Patch {
+        route: route!(),
+        operation: Op::Add {
+            key: "newkey".to_string(),
+            value: serde_yaml::Value::String("newvalue".to_string()),
+        },
+    }];
+
+    let result =
+        apply_yaml_patches(&yamlpath::Document::new(original).unwrap(), &operations).unwrap();
+
+    insta::assert_snapshot!(result.source(), @r#"
+        name: Test
+        key: value
+        newkey: newvalue
+        "#);
+}
+
+#[test]
+fn test_preserve_trailing_newline_replace_nested_at_end() {
+    // Test Replace operation preserves trailing newline when replacing nested value at end
+    let original = r#"jobs:
+  test:
+    runs-on: ubuntu-latest
+    env:
+      VAR: old
+"#;
+
+    let operations = vec![Patch {
+        route: route!("jobs", "test", "env", "VAR"),
+        operation: Op::Replace(serde_yaml::Value::String("new".to_string())),
+    }];
+
+    let result =
+        apply_yaml_patches(&yamlpath::Document::new(original).unwrap(), &operations).unwrap();
+
+    insta::assert_snapshot!(result.source(), @r#"
+        jobs:
+          test:
+            runs-on: ubuntu-latest
+            env:
+              VAR: new
+        "#);
+}
+
+#[test]
+fn test_preserve_trailing_newline_add_to_nested_mapping_at_end() {
+    // Test Add operation preserves trailing newline when adding to nested mapping at end
+    let original = r#"jobs:
+  test:
+    runs-on: ubuntu-latest
+    steps:
+      - run: echo "test"
+"#;
+
+    let operations = vec![Patch {
+        route: route!("jobs", "test", "steps", 0),
+        operation: Op::Add {
+            key: "name".to_string(),
+            value: serde_yaml::Value::String("Test step".to_string()),
+        },
+    }];
+
+    let result =
+        apply_yaml_patches(&yamlpath::Document::new(original).unwrap(), &operations).unwrap();
+
+    insta::assert_snapshot!(result.source(), @r#"
+        jobs:
+          test:
+            runs-on: ubuntu-latest
+            steps:
+              - run: echo "test"
+                name: Test step
+        "#);
+}
+
+#[test]
+fn test_preserve_trailing_newline_replace_multiline_at_end() {
+    // Test Replace operation preserves trailing newline when replacing multiline value at end
+    let original = r#"description: |
+  Line 1
+  Line 2
+"#;
+
+    let operations = vec![Patch {
+        route: route!("description"),
+        operation: Op::Replace(serde_yaml::Value::String("New description".to_string())),
+    }];
+
+    let result =
+        apply_yaml_patches(&yamlpath::Document::new(original).unwrap(), &operations).unwrap();
+
+    insta::assert_snapshot!(result.source(), @r#"
+        description: New description
+        "#);
+}
+
+#[test]
+fn test_preserve_trailing_newline_no_newline_original() {
+    // Test that operations don't add trailing newline if original document doesn't have one
+    let original = r#"name: Test
+key: value"#;
+
+    let operations = vec![Patch {
+        route: route!("key"),
+        operation: Op::Replace(serde_yaml::Value::String("newvalue".to_string())),
+    }];
+
+    let result =
+        apply_yaml_patches(&yamlpath::Document::new(original).unwrap(), &operations).unwrap();
+
+    insta::assert_snapshot!(result.source(), @"name: Test
+key: newvalue");
+}