If the snippet begins with a comment, preserve it in the output. This…

… time with no changes to Yaml LST
Gepardec · Aug 28, 2024 · e4fb82b · e4fb82b
1 parent f73054b
commit e4fb82b
Show file tree

Hide file tree

Showing 3 changed files with 69 additions and 17 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
@@ -87,20 +87,28 @@ public String getDescription() {
 
     @Override
     public TreeVisitor<?, ExecutionContext> getVisitor() {
-        JsonPathMatcher matcher = new JsonPathMatcher(key);
-        Yaml incoming = new YamlParser().parse(yaml)
-                .findFirst()
-                .map(Yaml.Documents.class::cast)
-                .orElseThrow(() -> new IllegalArgumentException("Could not parse as YAML"))
-                .getDocuments().get(0).getBlock();
-
-
         return Preconditions.check(new FindSourceFiles(filePattern), new YamlIsoVisitor<ExecutionContext>() {
+            final JsonPathMatcher matcher = new JsonPathMatcher(key);
+
+            final Yaml incoming = new YamlParser().parse(yaml)
+                    .findFirst()
+                    .map(Yaml.Documents.class::cast)
+                    .map(docs -> {
+                        // Any comments will have been put on the parent Document node, preserve by copying to the mapping
+                        Yaml.Document doc = docs.getDocuments().get(0);
+                        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())));
+                        }
+                        return doc.getBlock().withPrefix(doc.getPrefix());
+                    })
+                    .orElseThrow(() -> new IllegalArgumentException("Could not parse as YAML"));
+
             @Override
             public Yaml.Document visitDocument(Yaml.Document document, ExecutionContext ctx) {
                 if ("$".equals(key)) {
                     return document.withBlock((Yaml.Block) new MergeYamlVisitor<>(document.getBlock(), yaml,
-                            Boolean.TRUE.equals(acceptTheirs), objectIdentifyingProperty).visit(document.getBlock(),
+                            Boolean.TRUE.equals(acceptTheirs), objectIdentifyingProperty).visitNonNull(document.getBlock(),
                             ctx, getCursor()));
                 }
                 Yaml.Document d = super.visitDocument(document, ctx);
@@ -120,7 +128,7 @@ public Yaml.Document visitDocument(Yaml.Document document, ExecutionContext ctx)
                     // No matching element already exists, so it must be constructed
                     //noinspection LanguageMismatch
                     return d.withBlock((Yaml.Block) new MergeYamlVisitor<>(d.getBlock(), snippet,
-                            Boolean.TRUE.equals(acceptTheirs), objectIdentifyingProperty).visit(d.getBlock(),
+                            Boolean.TRUE.equals(acceptTheirs), objectIdentifyingProperty).visitNonNull(d.getBlock(),
                             ctx, getCursor()));
                 }
                 return d;
@@ -192,7 +200,7 @@ public Yaml.Sequence visitSequence(Yaml.Sequence sequence, ExecutionContext ctx)
                     return sequence.withEntries(ListUtils.map(sequence.getEntries(),
                             entry -> entry.withBlock((Yaml.Block) new MergeYamlVisitor<>(entry.getBlock(), incoming,
                                     Boolean.TRUE.equals(acceptTheirs), objectIdentifyingProperty)
-                                    .visit(entry.getBlock(), ctx, new Cursor(getCursor(), entry)))));
+                                    .visitNonNull(entry.getBlock(), ctx, new Cursor(getCursor(), entry)))));
                 }
                 return super.visitSequence(sequence, ctx);
             }

diff --git a/rewrite-yaml/src/main/java/org/openrewrite/yaml/MergeYamlVisitor.java b/rewrite-yaml/src/main/java/org/openrewrite/yaml/MergeYamlVisitor.java
@@ -41,11 +41,22 @@ public class MergeYamlVisitor<P> extends YamlVisitor<P> {
     private boolean shouldAutoFormat = true;
 
     public MergeYamlVisitor(Yaml scope, @Language("yml") String yamlString, boolean acceptTheirs, @Nullable String objectIdentifyingProperty) {
-        this(scope, new YamlParser().parse(yamlString)
-                .findFirst()
-                .map(Yaml.Documents.class::cast)
-                .orElseThrow(() -> new IllegalArgumentException("Could not parse as YAML"))
-                .getDocuments().get(0).getBlock(), acceptTheirs, objectIdentifyingProperty);
+        this(scope,
+                new YamlParser().parse(yamlString)
+                        .findFirst()
+                        .map(Yaml.Documents.class::cast)
+                        .map(docs -> {
+                            // Any comments will have been put on the parent Document node, preserve by copying to the mapping
+                            Yaml.Document doc = docs.getDocuments().get(0);
+                            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())));
+                            }
+                            return doc.getBlock().withPrefix(doc.getPrefix());
+                        })
+                        .orElseThrow(() -> new IllegalArgumentException("Could not parse as YAML")),
+                acceptTheirs,
+                objectIdentifyingProperty);
     }
 
     @Override
@@ -105,7 +116,7 @@ private Yaml.Mapping mergeMapping(Yaml.Mapping m1, Yaml.Mapping m2, P p, Cursor
                 if (keyMatches(existingEntry, incomingEntry)) {
                     return existingEntry.withValue((Yaml.Block) new MergeYamlVisitor<>(existingEntry.getValue(),
                             incomingEntry.getValue(), acceptTheirs, objectIdentifyingProperty, shouldAutoFormat)
-                            .visit(existingEntry.getValue(), p, new Cursor(cursor, existingEntry)));
+                            .visitNonNull(existingEntry.getValue(), p, new Cursor(cursor, existingEntry)));
                 }
             }
             return existingEntry;

diff --git a/rewrite-yaml/src/test/java/org/openrewrite/yaml/MergeYamlTest.java b/rewrite-yaml/src/test/java/org/openrewrite/yaml/MergeYamlTest.java
@@ -1135,4 +1135,37 @@ void mergeEmptyStructureFollowedByCopyValue() {
           )
         );
     }
+
+    @Test
+    void comment() {
+        rewriteRun(
+          spec -> spec.recipe(
+              new MergeYaml(
+                "$",
+                //language=yaml
+                """
+                  
+                  # new stuff
+                  new-property: value
+                  """,
+                false,
+                null,
+                null
+              )),
+          yaml(
+            """
+              # config
+              activate-auto: true
+              activate-mep: true
+              """,
+            """
+              # config
+              activate-auto: true
+              activate-mep: true
+              # new stuff
+              new-property: value
+              """
+          )
+        );
+    }
 }