Update ChangePropertyKey to auto-format modified mappings. (issue #1873)

openrewrite · Jun 21, 2022 · 88fd0df · 88fd0df
1 parent b0ab650
commit 88fd0df
Show file tree

Hide file tree

Showing 3 changed files with 65 additions and 38 deletions.
diff --git a/rewrite-yaml/src/main/java/org/openrewrite/yaml/ChangePropertyKey.java b/rewrite-yaml/src/main/java/org/openrewrite/yaml/ChangePropertyKey.java
@@ -121,7 +121,6 @@ public Yaml.Mapping.Entry visitMappingEntry(Yaml.Mapping.Entry entry, P p) {
                                 propertyToTest,
                                 entry
                         ));
-                        doAfterVisit(new DeleteProperty.DeletePropertyVisitor<>(entry));
                         break;
                     }
                     propertyToTest = propertyToTest.substring(value.length() + 1);
@@ -148,9 +147,6 @@ public Yaml.Mapping visitMapping(Yaml.Mapping mapping, P p) {
             Yaml.Mapping m = super.visitMapping(mapping, p);
             if (m.getEntries().contains(scope)) {
                 String newEntryPrefix = scope.getPrefix();
-                if (newEntryPrefix.isEmpty()) {
-                    newEntryPrefix = "\n";
-                }
                 Yaml.Mapping.Entry newEntry = new Yaml.Mapping.Entry(randomId(),
                         newEntryPrefix,
                         Markers.EMPTY,
@@ -167,12 +163,52 @@ public Yaml.Mapping visitMapping(Yaml.Mapping mapping, P p) {
                         return e;
                     }));
                 } else {
-                    m = m.withEntries(ListUtils.concat(m.getEntries(), newEntry));
-                    doAfterVisit(new DeleteProperty.DeletePropertyVisitor<>(entryToReplace));
+                    m = (Yaml.Mapping) new DeletePropertyVisitor<>(entryToReplace).visitNonNull(m, p);
+                    m = maybeAutoFormat(m, m.withEntries(ListUtils.concat(m.getEntries(), newEntry)), p, getCursor().getParentOrThrow());
                 }
             }
 
             return m;
         }
     }
+    private static class DeletePropertyVisitor<P> extends YamlIsoVisitor<P> {
+        private final Yaml.Mapping.Entry scope;
+
+        private DeletePropertyVisitor(Yaml.Mapping.Entry scope) {
+            this.scope = scope;
+        }
+
+        @Override
+        public Yaml.Mapping visitMapping(Yaml.Mapping mapping, P p) {
+            Yaml.Mapping m = super.visitMapping(mapping, p);
+
+            boolean changed = false;
+            List<Yaml.Mapping.Entry> entries = new ArrayList<>();
+            for (Yaml.Mapping.Entry entry : m.getEntries()) {
+                if (entry == scope || (entry.getValue() instanceof Yaml.Mapping && ((Yaml.Mapping) entry.getValue()).getEntries().isEmpty())) {
+                    changed = true;
+                } else {
+                    entries.add(entry);
+                }
+            }
+
+            if (entries.size() == 1) {
+                entries = ListUtils.map(entries, e -> e.withPrefix(""));
+            }
+
+            if (changed) {
+                m = m.withEntries(entries);
+
+                if (getCursor().getParentOrThrow().getValue() instanceof Yaml.Document) {
+                    Yaml.Document document = getCursor().getParentOrThrow().getValue();
+                    if (!document.isExplicit()) {
+                        m = m.withEntries(m.getEntries());
+                    }
+                }
+            }
+            return m;
+        }
+    }
+
+
 }
diff --git a/rewrite-yaml/src/main/java/org/openrewrite/yaml/DeleteProperty.java b/rewrite-yaml/src/main/java/org/openrewrite/yaml/DeleteProperty.java
@@ -111,10 +111,10 @@ public Yaml.Mapping.Entry visitMappingEntry(Yaml.Mapping.Entry entry, ExecutionC
         };
     }
 
-    public static class DeletePropertyVisitor<P> extends YamlVisitor<P> {
+    private static class DeletePropertyVisitor<P> extends YamlVisitor<P> {
         private final Yaml.Mapping.Entry scope;
 
-        public DeletePropertyVisitor(Yaml.Mapping.Entry scope) {
+        private DeletePropertyVisitor(Yaml.Mapping.Entry scope) {
             this.scope = scope;
         }
 

diff --git a/rewrite-yaml/src/test/kotlin/org/openrewrite/yaml/ChangePropertyKeyTest.kt b/rewrite-yaml/src/test/kotlin/org/openrewrite/yaml/ChangePropertyKeyTest.kt
@@ -16,7 +16,6 @@
 package org.openrewrite.yaml
 
 import org.assertj.core.api.Assertions.assertThat
-import org.junit.jupiter.api.Disabled
 import org.junit.jupiter.api.Test
 import org.junit.jupiter.api.io.TempDir
 import org.junit.jupiter.params.ParameterizedTest
@@ -34,47 +33,39 @@ class ChangePropertyKeyTest : YamlRecipeTest {
             null
         )
 
-    @Disabled
     @Issue("https://github.com/openrewrite/rewrite/issues/1873")
     @Test
     fun `shorter new key with indented config`() = assertChanged(
         recipe = ChangePropertyKey("a.b.c.d.e", "x.y", null, null),
-        before = """
-    a:
-      b:
-        c:
-          d:
-            e:
-              child: true
-    """,
-        after = """
-    x.y:
-      child: true
-    """
-        /* actual:
-        x.y:
+        before =
+        """
+        a:
+          b:
+            c:
+              d:
+                e:
                   child: true
-        */
+        """,
+            after = """
+        x.y:
+          child: true
+        """
     )
 
-    @Disabled
     @Issue("https://github.com/openrewrite/rewrite/issues/1873")
     @Test
     fun `longer new key with indented config`() = assertChanged(
         recipe = ChangePropertyKey("x.y", "a.b.c.d.e",  null, null),
-        before = """
-    x:
-      y:
-        child: true
-    """,
-        after = """
-    a.b.c.d.e:
-      child: true
-    """
-        /* actual:
-        a.b.c.d.e:
+        before =
+        """
+        x:
+          y:
             child: true
-        */
+        """,
+            after = """
+        a.b.c.d.e:
+          child: true
+        """
     )
 
     @Test