[Compiler plugin] Support group() in AddDsl

Kotlin · Jan 28, 2025 · 8f8d368 · 8f8d368
1 parent 1349a8d
commit 8f8d368
Show file tree

Hide file tree

Showing 4 changed files with 19 additions and 0 deletions.
diff --git a/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/api/add.kt b/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/api/add.kt
@@ -225,6 +225,7 @@ public class AddDsl<T>(
     @AccessApiOverload
     public fun group(column: AnyColumnGroupAccessor, body: AddDsl<T>.() -> Unit): Unit = group(column.name(), body)
 
+    @Interpretable("AddDslNamedGroup")
     public fun group(name: String, body: AddDsl<T>.() -> Unit) {
         val dsl = AddDsl(df)
         body(dsl)

diff --git a/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/impl/api/add.kt b/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/impl/api/add.kt
@@ -68,3 +68,15 @@ class AddDslStringInvoke : AbstractInterpreter<Unit>() {
         dsl.columns.add(SimpleColumnGroup(receiver, addDsl.columns))
     }
 }
+
+class AddDslNamedGroup : AbstractInterpreter<Unit>() {
+    val Arguments.dsl: AddDslApproximation by arg()
+    val Arguments.name: String by arg()
+    val Arguments.body by dsl()
+
+    override fun Arguments.interpret() {
+        val addDsl = AddDslApproximation(mutableListOf())
+        body(addDsl, emptyMap())
+        dsl.columns.add(SimpleColumnGroup(name, addDsl.columns))
+    }
+}
diff --git a/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/loadInterpreter.kt b/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/loadInterpreter.kt
@@ -71,6 +71,7 @@ import org.jetbrains.kotlin.name.CallableId
 import org.jetbrains.kotlin.name.ClassId
 import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.name.StandardClassIds
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.AddDslNamedGroup
 import org.jetbrains.kotlinx.dataframe.plugin.impl.api.AddDslStringInvoke
 import org.jetbrains.kotlinx.dataframe.plugin.impl.api.AddId
 import org.jetbrains.kotlinx.dataframe.plugin.impl.api.Aggregate
@@ -255,6 +256,7 @@ internal inline fun <reified T> String.load(): T {
         "FlattenDefault" -> FlattenDefault()
         "AddId" -> AddId()
         "AddDslStringInvoke" -> AddDslStringInvoke()
+        "AddDslNamedGroup" -> AddDslNamedGroup()
         "MapToFrame" -> MapToFrame()
         "Move0" -> Move0()
         "ToTop" -> ToTop()

diff --git a/plugins/kotlin-dataframe/testData/box/addDsl.kt b/plugins/kotlin-dataframe/testData/box/addDsl.kt
@@ -9,8 +9,12 @@ fun box(): String {
         "group" {
             "a" from { it }
         }
+        group("group1") {
+            "b" from { it }
+        }
     }
 
     df.group.a
+    df.group1.b
     return "OK"
 }