Fix indentation formatting for variable initializers

openrewrite · Oct 2, 2023 · 4352c42 · 4352c42
1 parent 2c1aad1
commit 4352c42
Show file tree

Hide file tree

Showing 2 changed files with 50 additions and 1 deletion.
diff --git a/src/main/java/org/openrewrite/kotlin/format/TabsAndIndentsVisitor.java b/src/main/java/org/openrewrite/kotlin/format/TabsAndIndentsVisitor.java
@@ -24,6 +24,7 @@
 import org.openrewrite.java.tree.*;
 import org.openrewrite.kotlin.KotlinIsoVisitor;
 import org.openrewrite.kotlin.marker.Implicit;
+import org.openrewrite.kotlin.marker.SingleExpressionBlock;
 import org.openrewrite.kotlin.style.TabsAndIndentsStyle;
 import org.openrewrite.kotlin.style.WrappingAndBracesStyle;
 import org.openrewrite.kotlin.tree.K;
@@ -94,7 +95,9 @@ public J preVisit(@Nullable J tree, P p) {
                 tree instanceof J.ArrayDimension ||
                 tree instanceof J.ClassDeclaration) {
             getCursor().putMessage("indentType", IndentType.ALIGN);
-        } else if (tree instanceof J.Block ||
+        } else if (tree instanceof  J.Block && tree.getMarkers().findFirst(SingleExpressionBlock.class).isPresent()) {
+            getCursor().putMessage("indentType", wrappingStyle.getExpressionBodyFunctions().getUseContinuationIndent() ? IndentType.CONTINUATION_INDENT : IndentType.INDENT);
+        } else if (tree instanceof  J.Block ||
                 tree instanceof J.If ||
                 tree instanceof J.If.Else ||
                 tree instanceof J.ForLoop ||
@@ -213,6 +216,17 @@ && getCursor().getValue() instanceof J.ClassDeclaration) {
         return s;
     }
 
+    @Override
+    public <T> JLeftPadded<T> visitLeftPadded(@Nullable JLeftPadded<T> left, JLeftPadded.Location loc, P p) {
+        if (loc == JLeftPadded.Location.VARIABLE_INITIALIZER) {
+            // this formatting option also applies to variable declarations
+            getCursor().putMessage("indentType",
+                    wrappingStyle.getExpressionBodyFunctions().getUseContinuationIndent() ? IndentType.CONTINUATION_INDENT : IndentType.INDENT);
+        }
+
+        return super.visitLeftPadded(left, loc, p);
+    }
+
     @Override
     public <T> JRightPadded<T> visitRightPadded(@Nullable JRightPadded<T> right, JRightPadded.Location loc, P p) {
         if (right == null) {

diff --git a/src/test/java/org/openrewrite/kotlin/format/TabsAndIndentsTest.java b/src/test/java/org/openrewrite/kotlin/format/TabsAndIndentsTest.java
@@ -56,6 +56,16 @@ private static Consumer<RecipeSpec> tabsAndIndents(UnaryOperator<TabsAndIndentsS
           )));
     }
 
+    private static Consumer<RecipeSpec> wrappingAndBraces(UnaryOperator<WrappingAndBracesStyle> with) {
+        return spec -> spec.recipe(toRecipe(() -> new TabsAndIndentsVisitor<>(IntelliJ.tabsAndIndents(), with.apply(IntelliJ.wrappingAndBraces()))))
+          .parser(KotlinParser.builder().styles(singletonList(
+            new NamedStyles(
+              Tree.randomId(), "test", "test", "test", emptySet(),
+              singletonList(with.apply(IntelliJ.wrappingAndBraces()))
+            )
+          )));
+    }
+
     private static Consumer<RecipeSpec> style(UnaryOperator<TabsAndIndentsStyle> tabsAndIndents, UnaryOperator<WrappingAndBracesStyle> wrappingAndBraces) {
         return spec -> spec.recipe(toRecipe(() -> new TabsAndIndentsVisitor<>(
             tabsAndIndents.apply(IntelliJ.tabsAndIndents()), wrappingAndBraces.apply(IntelliJ.wrappingAndBraces()))))
@@ -693,6 +703,31 @@ fun method(s: String) {
         );
     }
 
+    @Test
+    void expressionBodyFunctions() {
+        rewriteRun(
+          kotlin(
+            """
+              val i =
+                  1
+              fun f() =
+                  2
+              """
+          )
+        );
+        rewriteRun(
+          wrappingAndBraces(style -> style.withExpressionBodyFunctions(style.getExpressionBodyFunctions().withUseContinuationIndent(true))),
+          kotlin(
+            """
+              val i =
+                      1
+              fun f() =
+                      2
+              """
+          )
+        );
+    }
+
     /**
      * Slight renaming but structurally the same as IntelliJ's code style view.
      */