Improve yaml handling of comments in sequences.

openrewrite · Oct 16, 2024 · 6602fe7 · 6602fe7
1 parent 9da3321
commit 6602fe7
Show file tree

Hide file tree

Showing 4 changed files with 54 additions and 3 deletions.
diff --git a/rewrite-yaml/src/main/java/org/openrewrite/yaml/MergeYaml.java b/rewrite-yaml/src/main/java/org/openrewrite/yaml/MergeYaml.java
@@ -99,6 +99,9 @@ public TreeVisitor<?, ExecutionContext> getVisitor() {
                         if(doc.getBlock() instanceof Yaml.Mapping) {
                             Yaml.Mapping m = (Yaml.Mapping) doc.getBlock();
                             return m.withEntries(ListUtils.mapFirst(m.getEntries(), entry -> entry.withPrefix(doc.getPrefix())));
+                        } else if (doc.getBlock() instanceof Yaml.Sequence) {
+                            Yaml.Sequence s = (Yaml.Sequence) doc.getBlock();
+                            return s.withEntries(ListUtils.mapFirst(s.getEntries(), entry -> entry.withPrefix(doc.getPrefix())));
                         }
                         return doc.getBlock().withPrefix(doc.getPrefix());
                     })

diff --git a/rewrite-yaml/src/main/java/org/openrewrite/yaml/MergeYamlVisitor.java b/rewrite-yaml/src/main/java/org/openrewrite/yaml/MergeYamlVisitor.java
@@ -51,6 +51,9 @@ public MergeYamlVisitor(Yaml scope, @Language("yml") String yamlString, boolean
                             if(doc.getBlock() instanceof Yaml.Mapping) {
                                 Yaml.Mapping m = (Yaml.Mapping) doc.getBlock();
                                 return m.withEntries(ListUtils.mapFirst(m.getEntries(), entry -> entry.withPrefix(doc.getPrefix())));
+                            } else if (doc.getBlock() instanceof Yaml.Sequence) {
+                                Yaml.Sequence s = (Yaml.Sequence) doc.getBlock();
+                                return s.withEntries(ListUtils.mapFirst(s.getEntries(), entry -> entry.withPrefix(doc.getPrefix())));
                             }
                             return doc.getBlock().withPrefix(doc.getPrefix());
                         })
@@ -172,11 +175,20 @@ private Yaml.Sequence mergeSequence(Yaml.Sequence s1, Yaml.Sequence s2, P p, Cur
                     for (Yaml.Sequence.Entry existingEntry : s1.getEntries()) {
                         Yaml.Mapping existingMapping = (Yaml.Mapping) existingEntry.getBlock();
                         if (keyMatches(existingMapping, incomingMapping)) {
-                            return existingEntry.withBlock(mergeMapping(existingMapping, incomingMapping, p, cursor));
+                            Yaml.Sequence.Entry e1 = existingEntry.withBlock(mergeMapping(existingMapping, incomingMapping, p, cursor));
+                            if(e1 == existingEntry) {
+                                // Made no change, no need to consider the entry "mutated"
+                                //noinspection DataFlowIssue
+                                return null;
+                            }
+                            return e1;
                         }
                     }
                     return entry;
                 });
+                if (mutatedEntries.isEmpty()) {
+                    return s1;
+                }
 
                 List<Yaml.Sequence.Entry> entries = ListUtils.concatAll(
                         s1.getEntries().stream().filter(entry -> !mutatedEntries.contains(entry))

diff --git a/rewrite-yaml/src/main/java/org/openrewrite/yaml/format/IndentsVisitor.java b/rewrite-yaml/src/main/java/org/openrewrite/yaml/format/IndentsVisitor.java
@@ -97,7 +97,8 @@ public IndentsVisitor(IndentsStyle style, @Nullable Tree stopAfter) {
     }
 
     private boolean isUnindentedTopLevel() {
-        return getCursor().getParentOrThrow().getValue() instanceof Yaml.Document ||
+        return getCursor().getValue() instanceof Yaml.Documents ||
+                getCursor().getParentOrThrow().getValue() instanceof Yaml.Document ||
                 getCursor().getParentOrThrow(2).getValue() instanceof Yaml.Document;
     }
 
@@ -123,10 +124,12 @@ private String indentTo(String prefix, int column) {
         }
 
         int indent = findIndent(prefix);
-        prefix = indentComments(prefix, indent);
         if (indent != column) {
             int shift = column - indent;
             prefix = indent(prefix, shift);
+            prefix = indentComments(prefix, shift);
+        } else {
+            prefix = indentComments(prefix, indent);
         }
 
         return prefix;

diff --git a/rewrite-yaml/src/test/java/org/openrewrite/yaml/MergeYamlTest.java b/rewrite-yaml/src/test/java/org/openrewrite/yaml/MergeYamlTest.java
@@ -1168,4 +1168,37 @@ void comment() {
           )
         );
     }
+
+    @Test
+    void commentInList() {
+        rewriteRun(
+          spec -> spec.recipe(
+            new MergeYaml(
+              "$.groups",
+              //language=yaml
+              """
+                
+                # comment
+                - id: 3
+                """,
+              false,
+              "id",
+              null
+            )),
+          yaml(
+            """
+              groups:
+                - id: 1
+                - id: 2
+              """,
+            """
+              groups:
+                - id: 1
+                - id: 2
+                # comment
+                - id: 3
+              """
+          )
+        );
+    }
 }