From 29b6c94afe1de0da16eaf1842892ebc682aa3ecb Mon Sep 17 00:00:00 2001
From: nickreid <nickreid@google.com>
Date: Tue, 2 Apr 2024 10:18:44 -0700
Subject: [PATCH] Automatically manage trailing commas in enums (#449)

Summary: Pull Request resolved: https://github.com/facebook/ktfmt/pull/449

Reviewed By: cortinico

Differential Revision: D55495213

Pulled By: hick209

fbshipit-source-id: cfa075efd2096af3eea51f0a232260b049ab12d4
---
 .../facebook/ktfmt/format/TrailingCommas.kt   | 35 +++++++
 .../format/GoogleStyleFormatterKtTest.kt      | 99 +++++++++++++++++++
 2 files changed, 134 insertions(+)

diff --git a/core/src/main/java/com/facebook/ktfmt/format/TrailingCommas.kt b/core/src/main/java/com/facebook/ktfmt/format/TrailingCommas.kt
index 4f9abbec..4ee63750 100644
--- a/core/src/main/java/com/facebook/ktfmt/format/TrailingCommas.kt
+++ b/core/src/main/java/com/facebook/ktfmt/format/TrailingCommas.kt
@@ -1,10 +1,28 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
 package com.facebook.ktfmt.format
 
 import org.jetbrains.kotlin.com.intellij.psi.PsiComment
 import org.jetbrains.kotlin.com.intellij.psi.PsiElement
 import org.jetbrains.kotlin.com.intellij.psi.PsiWhiteSpace
+import org.jetbrains.kotlin.psi.KtClassBody
 import org.jetbrains.kotlin.psi.KtCollectionLiteralExpression
 import org.jetbrains.kotlin.psi.KtElement
+import org.jetbrains.kotlin.psi.KtEnumEntry
 import org.jetbrains.kotlin.psi.KtFunctionLiteral
 import org.jetbrains.kotlin.psi.KtLambdaExpression
 import org.jetbrains.kotlin.psi.KtParameterList
@@ -70,6 +88,13 @@ object TrailingCommas {
             return // Never add trailing commas to lambda param lists
           }
         }
+        is KtClassBody -> {
+          EnumEntryList.extractChildList(element)?.also {
+            if (it.terminatingSemicolon != null) {
+              return // Never add a trailing comma after there is already a terminating semicolon
+            }
+          }
+        }
       }
 
       val list = extractManagedList(element) ?: return
@@ -96,6 +121,16 @@ object TrailingCommas {
         ManagedList(element.getInnerExpressions(), element.trailingComma)
       }
       is KtWhenEntry -> ManagedList(element.conditions.toList(), element.trailingComma)
+      is KtEnumEntry -> {
+        EnumEntryList.extractParentList(element).let {
+          ManagedList(it.enumEntries, it.trailingComma)
+        }
+      }
+      is KtClassBody -> {
+        EnumEntryList.extractChildList(element)?.let {
+          ManagedList(it.enumEntries, it.trailingComma)
+        }
+      }
       else -> null
     }
   }
diff --git a/core/src/test/java/com/facebook/ktfmt/format/GoogleStyleFormatterKtTest.kt b/core/src/test/java/com/facebook/ktfmt/format/GoogleStyleFormatterKtTest.kt
index 7b6c6321..337259ba 100644
--- a/core/src/test/java/com/facebook/ktfmt/format/GoogleStyleFormatterKtTest.kt
+++ b/core/src/test/java/com/facebook/ktfmt/format/GoogleStyleFormatterKtTest.kt
@@ -1694,6 +1694,105 @@ class GoogleStyleFormatterKtTest {
           formattingOptions = Formatter.GOOGLE_FORMAT,
           deduceMaxWidth = true)
 
+  @Test
+  fun `trailing commas in enums`() {
+    val code =
+        """
+        |enum class A {}
+        |
+        |enum class B {
+        |  Z // Comment
+        |}
+        |
+        |enum class C {
+        |  Z, // Comment
+        |}
+        |
+        |enum class D {
+        |  Z,
+        |  Y // Comment
+        |}
+        |
+        |enum class E {
+        |  Z,
+        |  Y, // Comment
+        |}
+        |
+        |enum class F {
+        |  Z,
+        |  Y; // Comment
+        |
+        |  val x = 0
+        |}
+        |
+        |enum class G {
+        |  Z,
+        |  Y,; // Comment
+        |
+        |  val x = 0
+        |}
+        |
+        |enum class H {
+        |  Z,
+        |  Y() {} // Comment
+        |}
+        |
+        |enum class I {
+        |  Z,
+        |  Y() {}, // Comment
+        |}
+        |"""
+            .trimMargin()
+    val expected =
+        """
+        |enum class A {}
+        |
+        |enum class B {
+        |  Z // Comment
+        |}
+        |
+        |enum class C {
+        |  Z // Comment
+        |}
+        |
+        |enum class D {
+        |  Z,
+        |  Y, // Comment
+        |}
+        |
+        |enum class E {
+        |  Z,
+        |  Y, // Comment
+        |}
+        |
+        |enum class F {
+        |  Z,
+        |  Y; // Comment
+        |
+        |  val x = 0
+        |}
+        |
+        |enum class G {
+        |  Z,
+        |  Y; // Comment
+        |
+        |  val x = 0
+        |}
+        |
+        |enum class H {
+        |  Z,
+        |  Y() {}, // Comment
+        |}
+        |
+        |enum class I {
+        |  Z,
+        |  Y() {}, // Comment
+        |}
+        |"""
+            .trimMargin()
+    assertThatFormatting(code).withOptions(Formatter.GOOGLE_FORMAT).isEqualTo(expected)
+  }
+
   companion object {
     /** Triple quotes, useful to use within triple-quoted strings. */
     private const val TQ = "\"\"\""