Custom DataFrame KtLint ruleset for our users?

[Jolanrensen]

As I learned in the talk from @paul-dingemans at Kotlin Dev Day, it's possible to create a custom KtLint ruleset which would allow us to define a recommended style for using DataFrame in cases where it deviates from "ordinary" Kotlin.

There are a couple of places where this would make sense. If you look at our project, you can see we've suppressed KtLint here as well, or disabled a rule in our .editorconfig entirely.

There might be more, but these two seem like good candidates:

DataFrame operation chaining

DataFrame operations often use "intermediate classes" in its DSL to form a sort-of sentence to describe what needs to be done with the data. It's okay to put these intermediate steps on separate lines if it's just a single operation being called on a dataframe:

dataFrame
    .update { colA and colB }
    .where { it > 10 }
    .with { 100 * it }

However, when multiple operations need to happen after another, it's hard to tell where the instruction for one ends and the next begins:

dataFrame
    .update { colA and colB }
    .where { it > 10 }
    .with { 100 * it }
    .split { colC }
    .by(",")
    .into { "colC$it" }

In our examples and tests, we suppress ktlint_standard_chain-method-continuation and write it like this:

dataFrame
    .update { colA and colB }.where { it > 10 }.with { 100 * it }
    .split { colC }.by(",").into { "colC$it" }

If the contents of an invocation are too long, we could still recommend putting continuation on the same line, while requiring a full line break in between operations:

dataFrame
    .update {
        colA and colB
    }.where { 
        it > 10
    }.with { 
        100 * it 
    }
    .split { 
        colC
    }.by(
        ",",
    ).into { 
        "colC$it"
    }

or a mix like:

dataFrame.update { colA and colB }.where { it > 10 }.with { 
        100 * it 
    }
    .split { colC }.by(",").into { 
        "colC$it"
    }

dataFrameOf(header)(values)

Take a look at:

dataFrameOf("firstName", "lastName", "age", "city", "weight", "isHappy")(
    "Alice", "Cooper", 15, "London", 54, true,
    "Bob", "Dylan", 45, "Dubai", 87, true,
    "Charlie", "Daniels", 20, "Moscow", null, false,
    "Charlie", "Chaplin", 40, "Milan", null, true,
    "Bob", "Marley", 30, "Tokyo", 68, true,
    "Alice", "Wolf", 20, null, 55, false,
    "Charlie", "Byrd", 30, "Moscow", 90, true,
)

Writing out all values like this is only possible if we suppress ktlint:standard:argument-list-wrapping. Otherwise, it will be formatted like:

dataFrameOf("firstName", "lastName", "age", "city", "weight", "isHappy")(
    "Alice",
    "Cooper",
    15,
    "London",
    54,
    true,
    "Bob",
    "Dylan",
    45,
    "Dubai",
    87,
    true,
    ...
)

This loses all readability of the function and makes it harder to use. Maybe we could create a custom rule for DataFrameBuilderInvoke0 that could ignore this function or format it like this:

dataFrameOf("firstName", "lastName", "age", "city", "weight", "isHappy")(
    "Alice",   "Cooper",  15, "London", 54,   true,
    "Bob",     "Dylan",   45, "Dubai",  87,   true,
    "Charlie", "Daniels", 20, "Moscow", null, false,
    "Charlie", "Chaplin", 40, "Milan",  null, true,
    "Bob",     "Marley",  30, "Tokyo",  68,   true,
    "Alice",   "Wolf",    20, null,     55,   false,
    "Charlie", "Byrd",    30, "Moscow", 90,   true,
)

Maybe even like this!

dataFrameOf(
    "firstName", "lastName", "age", "city",   "weight", "isHappy",
)(
    "Alice",     "Cooper",   15,    "London", 54,       true,
    "Bob",       "Dylan",    45,    "Dubai",  87,       true,
    "Charlie",   "Daniels",  20,    "Moscow", null,     false,
    "Charlie",   "Chaplin",  40,    "Milan",  null,     true,
    "Bob",       "Marley",   30,    "Tokyo",  68,       true,
    "Alice",     "Wolf",     20,    null,     55,       false,
    "Charlie",   "Byrd",     30,    "Moscow", 90,       true,
)

(Yes, if we recommend people to use dataFrameOf(vararg Pair<String, Column>), this won't be necessary anymore.)

[paul-dingemans]

Hi @Jolanrensen

Great idea that you want to investigate how a custom Ktlint ruleset can be created for DataFrame. The example that you mentioned about the chain-method-continuation is unwanted for dataFrame expressions made my wonder how you can easily work around that without disabling that entire rule for non-dataFrame chains.

I came up with following proof of concept, that will allow you to create a custom ruleset for method chains that allows different handling for dataframe, and keep default Ktlint behavior for other method chains. For this proof of concept, I only looked for method chains starting with an identifier with name dataFrame.

The custom rule for method chaining could look like following:

package com.pinterest.ktlint.ruleset.standard.rules

import com.pinterest.ktlint.rule.engine.core.api.AutocorrectDecision
import com.pinterest.ktlint.rule.engine.core.api.ElementType.DOT
import com.pinterest.ktlint.rule.engine.core.api.ElementType.DOT_QUALIFIED_EXPRESSION
import com.pinterest.ktlint.rule.engine.core.api.ElementType.IDENTIFIER
import com.pinterest.ktlint.rule.engine.core.api.ElementType.REFERENCE_EXPRESSION
import com.pinterest.ktlint.rule.engine.core.api.ElementType.SAFE_ACCESS
import com.pinterest.ktlint.rule.engine.core.api.Rule
import com.pinterest.ktlint.rule.engine.core.api.Rule.VisitorModifier.RunAfterRule
import com.pinterest.ktlint.rule.engine.core.api.Rule.VisitorModifier.RunAfterRule.Mode.REGARDLESS_WHETHER_RUN_AFTER_RULE_IS_LOADED_OR_DISABLED
import com.pinterest.ktlint.rule.engine.core.api.RuleAutocorrectApproveHandler
import com.pinterest.ktlint.rule.engine.core.api.RuleId
import com.pinterest.ktlint.rule.engine.core.api.editorconfig.CODE_STYLE_PROPERTY
import com.pinterest.ktlint.rule.engine.core.api.editorconfig.EditorConfig
import com.pinterest.ktlint.rule.engine.core.api.editorconfig.INDENT_SIZE_PROPERTY
import com.pinterest.ktlint.rule.engine.core.api.editorconfig.INDENT_STYLE_PROPERTY
import com.pinterest.ktlint.rule.engine.core.api.editorconfig.MAX_LINE_LENGTH_PROPERTY
import com.pinterest.ktlint.rule.engine.core.api.parent
import org.jetbrains.kotlin.com.intellij.lang.ASTNode
import org.jetbrains.kotlin.com.intellij.psi.tree.TokenSet

public class DataFrameChainMethodContinuationRule :
    Rule(
        ruleId = RuleId("dataframe:chain-method-continuation"),
        about =
            About(
                maintainer = "Dataframe",
                repositoryUrl = "https://github.com/Kotlin/dataframe",
                issueTrackerUrl = "https://github.com/Kotlin/dataframe/issues",
            ),
        visitorModifiers = setOf(RunAfterRule(ARGUMENT_LIST_WRAPPING_RULE_ID, REGARDLESS_WHETHER_RUN_AFTER_RULE_IS_LOADED_OR_DISABLED)),
        setOf(
            CODE_STYLE_PROPERTY,
            ChainMethodContinuationRule.FORCE_MULTILINE_WHEN_CHAIN_OPERATOR_COUNT_GREATER_OR_EQUAL_THAN_PROPERTY,
            INDENT_SIZE_PROPERTY,
            INDENT_STYLE_PROPERTY,
            MAX_LINE_LENGTH_PROPERTY,
        ),
    ),
    RuleAutocorrectApproveHandler {
    private val chainMethodContinuationRule = ChainMethodContinuationRule()

    override fun beforeFirstNode(editorConfig: EditorConfig) {
        chainMethodContinuationRule.beforeFirstNode(editorConfig)
    }

    override fun beforeVisitChildNodes(
        node: ASTNode,
        emit: (offset: Int, errorMessage: String, canBeAutoCorrected: Boolean) -> AutocorrectDecision,
    ) {
        if (node.isDataFrameDotQualifiedExpression() || node.isChainOperatorInDataFrameDotQualifiedExpression()) {
            // Add logic for lint/format for dataFrames
        } else {
            chainMethodContinuationRule.beforeVisitChildNodes(node, emit)
        }
    }

    private fun ASTNode.isChainOperatorInDataFrameDotQualifiedExpression(): Boolean =
        elementType in chainOperatorTokenSet && parent?.isDataFrameDotQualifiedExpression() == true

    private fun ASTNode.isDataFrameDotQualifiedExpression(): Boolean =
        takeIf { it.elementType == DOT_QUALIFIED_EXPRESSION }
            ?.firstChildNode
            ?.let {
                when (it.elementType) {
                    DOT_QUALIFIED_EXPRESSION -> {
                        it.isDataFrameDotQualifiedExpression()
                    }

                    REFERENCE_EXPRESSION -> {
                        with(it.firstChildNode) {
                            elementType == IDENTIFIER && text == "dataFrame"
                        }
                    }

                    else -> {
                        false
                    }
                }
            }
            ?: false

    override fun afterVisitChildNodes(
        node: ASTNode,
        emit: (offset: Int, errorMessage: String, canBeAutoCorrected: Boolean) -> AutocorrectDecision,
    ) {
        chainMethodContinuationRule.afterVisitChildNodes(node, emit)
    }

    override fun afterLastNode() {
        chainMethodContinuationRule.afterLastNode()
    }

    private val chainOperatorTokenSet = TokenSet.create(DOT, SAFE_ACCESS)
}

The unit test class:

package com.pinterest.ktlint.ruleset.standard.rules

import com.pinterest.ktlint.ruleset.standard.StandardRuleSetProvider
import com.pinterest.ktlint.ruleset.standard.rules.ChainMethodContinuationRule.Companion.FORCE_MULTILINE_WHEN_CHAIN_OPERATOR_COUNT_GREATER_OR_EQUAL_THAN_PROPERTY
import com.pinterest.ktlint.test.KtLintAssertThat.Companion.assertThatRuleBuilder
import com.pinterest.ktlint.test.LintViolation
import org.junit.jupiter.api.Nested
import org.junit.jupiter.api.Test

class DataFrameChainMethodContinuationRuleTest {
    private val dataFrameChainMethodContinuationRuleAssertThat =
        assertThatRuleBuilder { DataFrameChainMethodContinuationRule() }
            .addAdditionalRuleProvider { MaxLineLengthRule() }
            .addRequiredRuleProviderDependenciesFrom(StandardRuleSetProvider())
            .assertThat()

    @Nested
    inner class `DataFrame chains are processed by own rule` {
        @Test
        fun `Ignore method chain when starting with dataFrame`() {
            val code =
                """
                fun foo(dataFrame: DataFrame) {
                    dataFrame.update { colA and colB }.where { it > 10 }.with {
                        100 * it
                    }
                    .split { colC }.by(",").into {
                        "colC${'$'}it"
                    }
                }
                """.trimIndent()
            dataFrameChainMethodContinuationRuleAssertThat(code).hasNoLintViolations()
        }
    }

    @Nested
    inner class `Non dataFrame chains are processed by base rule` {
        @Test
        fun `Given that no maximum line length is set, and a single line method chain does exceed the maximum number of chain operators then wrap`() {
            val code =
                """
                val foo = listOf(1, 2, 3).filter { it > 2 }!!.takeIf { it.count() > 100 }?.sum()!!
                """.trimIndent()
            val formattedCode =
                """
                val foo = listOf(1, 2, 3)
                    .filter { it > 2 }!!
                    .takeIf { it.count() > 100 }
                    ?.sum()!!
                """.trimIndent()
            dataFrameChainMethodContinuationRuleAssertThat(code)
                .withEditorConfigOverride(FORCE_MULTILINE_WHEN_CHAIN_OPERATOR_COUNT_GREATER_OR_EQUAL_THAN_PROPERTY to 3)
                .hasLintViolations(
                    LintViolation(1, 26, "Expected newline before '.'"),
                    LintViolation(1, 46, "Expected newline before '.'"),
                    LintViolation(1, 74, "Expected newline before '?.'"),
                ).isFormattedAs(formattedCode)
        }
    }
}