Improvement

openrewrite · Nov 18, 2024 · 4bc691b · 4bc691b
1 parent 2c3f52f
commit 4bc691b
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 12 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
@@ -84,6 +84,7 @@ public String getDescription() {
     }
 
     final static String FOUND_MATCHING_ELEMENT = "FOUND_MATCHING_ELEMENT";
+    final static String REMOVE_PREFIX = "REMOVE_PREFIX";
 
     @Override
     public TreeVisitor<?, ExecutionContext> getVisitor() {
@@ -114,12 +115,7 @@ public Yaml.Document visitDocument(Yaml.Document document, ExecutionContext ctx)
                             new MergeYamlVisitor<>(document.getBlock(), yaml, Boolean.TRUE.equals(acceptTheirs), objectIdentifyingProperty)
                                     .visitNonNull(document.getBlock(), ctx, getCursor())
                     );
-
-                    if (getCursor().getMessage("RemovePrefix", false)) {
-                        return d.withEnd(d.getEnd().withPrefix(""));
-                    }
-
-                    return d;
+                    return getCursor().getMessage(REMOVE_PREFIX, false) ? d.withEnd(d.getEnd().withPrefix("")) : d;
                 }
                 Yaml.Document d = super.visitDocument(document, ctx);
                 if (d == document && !getCursor().getMessage(FOUND_MATCHING_ELEMENT, false)) {

diff --git a/rewrite-yaml/src/main/java/org/openrewrite/yaml/MergeYamlVisitor.java b/rewrite-yaml/src/main/java/org/openrewrite/yaml/MergeYamlVisitor.java
@@ -34,6 +34,7 @@
 import java.util.stream.Collectors;
 
 import static org.openrewrite.internal.ListUtils.*;
+import static org.openrewrite.yaml.MergeYaml.REMOVE_PREFIX;
 
 @AllArgsConstructor
 @RequiredArgsConstructor
@@ -103,7 +104,7 @@ public Yaml visitMapping(Mapping existingMapping, P p) {
         if (existing.isScope(existingMapping) && incoming instanceof Mapping) {
             Mapping mapping = mergeMapping(existingMapping, (Mapping) incoming, p, getCursor());
 
-            if (getCursor().getMessage("RemovePrefix", false)) {
+            if (getCursor().getMessage(REMOVE_PREFIX, false)) {
                 List<Entry> entries = ((Mapping) getCursor().getValue()).getEntries();
                 return mapping.withEntries(mapLast(mapping.getEntries(), it -> it.withPrefix("\n" + grabPartLineBreak(entries.get(entries.size() - 1), 1))));
             }
@@ -149,7 +150,7 @@ private Mapping mergeMapping(Mapping m1, Mapping m2, P p, Cursor cursor) {
                 }
             }
             // workaround: autoFormat put sometimes extra spaces before elements
-            if (!mergedEntries.isEmpty() && it.getValue() instanceof Scalar && hasLineBreakPieces(mergedEntries.get(0), 2)) {
+            if (!mergedEntries.isEmpty() && it.getValue() instanceof Scalar && hasLineBreak(mergedEntries.get(0), 2)) {
                 return it.withPrefix("\n" + grabPartLineBreak(mergedEntries.get(0), 1));
             }
             return shouldAutoFormat ? autoFormat(it, p, cursor) : it;
@@ -187,14 +188,14 @@ private Mapping mergeMapping(Mapping m1, Mapping m2, P p, Cursor cursor) {
                             break;
                         }
                     }
-                    if (comment.isEmpty() && hasLineBreakPieces(entries.get(entries.size() - 1), 1)) {
+                    if (comment.isEmpty() && hasLineBreak(entries.get(entries.size() - 1), 1)) {
                         comment = grabPartLineBreak(entries.get(entries.size() - 1), 0);
                     }
                 }
 
                 Entry last = mutatedEntries.get(mutatedEntries.size() - 1);
                 mutatedEntries.set(mutatedEntries.size() - 1, last.withPrefix(comment + last.getPrefix()));
-                c.putMessage("RemovePrefix", true);
+                c.putMessage(REMOVE_PREFIX, true);
             }
         }
 
@@ -277,9 +278,9 @@ private Yaml.Sequence mergeSequence(Yaml.Sequence s1, Yaml.Sequence s2, P p, Cur
         }
     }
 
-    private boolean hasLineBreakPieces(Entry entry, int atLeast) {
+    private boolean hasLineBreak(Entry entry, int atLeastParts) {
         boolean a = LINE_BREAK.matcher(entry.getPrefix()).find();
-        return a && !grabPartLineBreak(entry, atLeast - 1).isEmpty();
+        return a && !grabPartLineBreak(entry, atLeastParts - 1).isEmpty();
     }
 
     private String grabPartLineBreak(Entry entry, int index) {