Trailing commas in enums

ktlint-ruleset-standard/src/main/kotlin/com/pinterest/ktlint/ruleset/standard/TrailingCommaRule.kt

@@ -15,8 +15,11 @@ import org.jetbrains.kotlin.com.intellij.lang.ASTNode
 import org.jetbrains.kotlin.com.intellij.psi.PsiElement
 import org.jetbrains.kotlin.com.intellij.psi.PsiWhiteSpace
 import org.jetbrains.kotlin.com.intellij.psi.tree.TokenSet
+import org.jetbrains.kotlin.psi.KtClass
+import org.jetbrains.kotlin.psi.KtClassBody
 import org.jetbrains.kotlin.psi.KtCollectionLiteralExpression
 import org.jetbrains.kotlin.psi.KtDestructuringDeclaration
+import org.jetbrains.kotlin.psi.KtEnumEntry
 import org.jetbrains.kotlin.psi.KtFunctionLiteral
 import org.jetbrains.kotlin.psi.KtPsiFactory
 import org.jetbrains.kotlin.psi.KtValueArgumentList
@@ -91,9 +94,7 @@ public class TrailingCommaRule :
             ElementType.TYPE_PARAMETER_LIST -> visitTypeList(node, emit, autoCorrect)
             ElementType.VALUE_PARAMETER_LIST -> visitValueList(node, emit, autoCorrect)
             ElementType.WHEN_ENTRY -> visitWhenEntry(node, emit, autoCorrect)
-            else -> Unit
-        }
-        when (node.elementType) {
+            ElementType.CLASS -> visitClassEntry(node, emit, autoCorrect)
             ElementType.COLLECTION_LITERAL_EXPRESSION -> visitCollectionLiteralExpression(node, emit, autoCorrect)
             ElementType.INDICES -> visitIndices(node, emit, autoCorrect)
             ElementType.TYPE_ARGUMENT_LIST -> visitTypeList(node, emit, autoCorrect)
@@ -192,6 +193,34 @@ public class TrailingCommaRule :
         node.reportAndCorrectTrailingCommaNodeBefore(inspectNode, emit, autoCorrect)
     }
 
+    private fun visitClassEntry(
+        node: ASTNode,
+        emit: (offset: Int, errorMessage: String, canBeAutoCorrected: Boolean) -> Unit,
+        autoCorrect: Boolean
+    ) {
+        val psi = node.psi
+        require(psi is KtClass)
+
+        if (psi.isEnum()) {
+            val classBody = node.children().singleOrNull { it.psi is KtClassBody }
+                ?: return // Nothing to do for empty enum class
+
+            val lastRelevantChild = classBody.psi.children.toList().asReversed()
+                .first()
+
+            val fullAST = psi
+            val fullPsi = classBody.psi
+
+            if (lastRelevantChild !is KtEnumEntry) {
+                // Aborting processing since Enum contains more than just constants.
+                return
+            }
+
+            val inspectNode = classBody.lastChildNode
+            node.reportAndCorrectTrailingCommaNodeBefore(inspectNode, emit, autoCorrect)
+        }
+    }
+
     private fun ASTNode.reportAndCorrectTrailingCommaNodeBefore(
         inspectNode: ASTNode,
         emit: (offset: Int, errorMessage: String, canBeAutoCorrected: Boolean) -> Unit,
@@ -325,7 +354,8 @@ public class TrailingCommaRule :
             ElementType.FUNCTION_TYPE,
             ElementType.TYPE_PARAMETER_LIST,
             ElementType.VALUE_PARAMETER_LIST,
-            ElementType.WHEN_ENTRY
+            ElementType.WHEN_ENTRY,
+            ElementType.CLASS
         )
 
         private val TYPES_ON_CALL_SITE = TokenSet.create(
@@ -335,6 +365,11 @@ public class TrailingCommaRule :
             ElementType.VALUE_ARGUMENT_LIST
         )
 
+        private val ENUM_IGNORED_TYPES = TokenSet.create(
+            ElementType.WHITE_SPACE,
+            ElementType.RBRACE
+        )
+
         internal const val ALLOW_TRAILING_COMMA_NAME = "ij_kotlin_allow_trailing_comma"
         private const val ALLOW_TRAILING_COMMA_DESCRIPTION = "Defines whether a trailing comma (or no trailing comma)" +
             "should be enforced on the defining side," +

ktlint-ruleset-standard/src/test/kotlin/com/pinterest/ktlint/ruleset/standard/TrailingCommaRuleTest.kt

@@ -1021,4 +1021,157 @@ class TrailingCommaRuleTest {
                 LintViolation(7, 6, "Missing trailing comma before \")\"")
             ).isFormattedAs(formattedCode)
     }
+
+    @Test
+    fun `Given that a trailing comma is required then add trailing comma after last enum member`() {
+        val code =
+            """
+            enum class Shape {
+                SQUARE,
+                TRIANGLE
+            }
+            """.trimIndent()
+        val formattedCode =
+            """
+            enum class Shape {
+                SQUARE,
+                TRIANGLE,
+            }
+            """.trimIndent()
+
+        trailingCommaRuleAssertThat(code)
+            .withEditorConfigOverride(allowTrailingCommaProperty to true)
+            .hasLintViolation(3, 13, "Missing trailing comma before \"}\"")
+            .isFormattedAs(formattedCode)
+    }
+
+    @Test
+    fun `Given that a trailing comma is is not allowed then remove comma after last enum member`() {
+        val code =
+            """
+            enum class Shape {
+                SQUARE,
+                TRIANGLE,
+            }
+            """.trimIndent()
+        val formattedCode =
+            """
+            enum class Shape {
+                SQUARE,
+                TRIANGLE
+            }
+            """.trimIndent()
+
+        trailingCommaRuleAssertThat(code)
+            .withEditorConfigOverride(allowTrailingCommaProperty to false)
+            .hasLintViolation(3, 13, "Unnecessary trailing comma before \"}\"")
+            .isFormattedAs(formattedCode)
+    }
+
+    @Test
+    fun `Given that a trailing comma is not allowed then nothing happens for enums terminated with semicolon`() {
+        val code =
+            """
+            enum class Shape {
+                SQUARE,
+                TRIANGLE;
+
+                fun print() = name()
+            }
+            """.trimIndent()
+
+        trailingCommaRuleAssertThat(code)
+            .withEditorConfigOverride(allowTrailingCommaProperty to false)
+            .hasNoLintViolations()
+    }
+
+    @Test
+    fun `Given that a trailing comma is allowed then nothing happens for enums terminated with semicolon`() {
+        val code =
+            """
+            enum class Shape {
+                SQUARE,
+                TRIANGLE;
+
+                fun print() = name()
+            }
+            """.trimIndent()
+
+        trailingCommaRuleAssertThat(code)
+            .withEditorConfigOverride(allowTrailingCommaProperty to true)
+            .hasNoLintViolations()
+    }
+
+    @Test
+    fun `Given that a trailing comma is allowed then it is added for complicated enums`() {
+        val code =
+            """
+            interface Printable {
+                fun print(): String
+            }
+
+            enum class Shape : Printable {
+                /*
+                 * Block comment is OK!
+                 */
+                Square {
+                    override fun print() = "■" // EOL Comment should be OK
+                },
+
+                Triangle {
+                    override fun print() = "▲"
+                }
+            }
+            """.trimIndent()
+
+        val formattedCode =
+            """
+            interface Printable {
+                fun print(): String
+            }
+
+            enum class Shape : Printable {
+                /*
+                 * Block comment is OK!
+                 */
+                Square {
+                    override fun print() = "■" // EOL Comment should be OK
+                },
+
+                Triangle {
+                    override fun print() = "▲"
+                },
+            }
+            """.trimIndent()
+
+        trailingCommaRuleAssertThat(code)
+            .withEditorConfigOverride(allowTrailingCommaProperty to true)
+            .hasLintViolation(15, 6, "Missing trailing comma before \"}\"")
+            .isFormattedAs(formattedCode)
+    }
+
+    @Test
+    fun `Given that a trailing comma is allowed then it is added if enum is already terminated with semi-colon`() {
+        val code =
+            """
+            enum class Shape {
+                SQUARE,
+                TRIANGLE;
+            }
+            """.trimIndent()
+
+        val formattedCode =
+            """
+            enum class Shape {
+                SQUARE,
+                TRIANGLE,
+                ;
+            }
+            """.trimIndent()
+
+        trailingCommaRuleAssertThat(code)
+            .withEditorConfigOverride(allowTrailingCommaProperty to true)
+            .hasLintViolation(3, 14, "Missing trailing comma before \"}\"")
+            .isFormattedAs(formattedCode)
+    }
 }