Add generateDataClasses for df.toListOf<generated classes> workflow

Author: koperagen

core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/codeGen/CodeGenerator.kt

@@ -36,6 +36,7 @@ public interface CodeGenerator : ExtensionsCodeGenerator {
         knownMarkers: Iterable<Marker> = emptyList(),
         readDfMethod: DefaultReadDfMethod? = null,
         fieldNameNormalizer: NameNormalizer = NameNormalizer.id(),
+        asDataClass: Boolean = false
     ): CodeGenResult
 
     public fun generate(

core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/codeGen/generateCode.kt

@@ -8,9 +8,7 @@ public inline fun <reified T> DataFrame<T>.generateCode(
     fields: Boolean = true,
     extensionProperties: Boolean = true,
 ): String {
-    val name = if (T::class.isAbstract) {
-        T::class.simpleName!!
-    } else "DataEntry"
+    val name = markerName<T>()
     return generateCode(name, fields, extensionProperties)
 }
 
@@ -36,6 +34,30 @@ public inline fun <reified T> DataFrame<T>.generateInterfaces(): String = genera
     extensionProperties = false
 )
 
+public inline fun <reified T> DataFrame<T>.generateDataClasses(
+    markerName: String? = null,
+    extensionProperties: Boolean = true,
+    visibility: MarkerVisibility = MarkerVisibility.IMPLICIT_PUBLIC,
+    useFqNames: Boolean = false
+): String {
+    val name = markerName ?: markerName<T>()
+    val codeGen = CodeGenerator.create(useFqNames)
+    return codeGen.generate(
+        schema = schema(),
+        name = name,
+        fields = true,
+        extensionProperties = extensionProperties,
+        isOpen = false,
+        visibility = visibility,
+        asDataClass = true
+    ).code.declarations
+}
+
+@PublishedApi
+internal inline fun <reified T> markerName(): String = if (T::class.isAbstract) {
+    T::class.simpleName!!
+} else "DataEntry"
+
 public fun <T> DataFrame<T>.generateInterfaces(markerName: String): String = generateCode(
     markerName = markerName,
     fields = true,

core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/codeGen/CodeGeneratorImpl.kt

@@ -448,12 +448,18 @@ internal class CodeGeneratorImpl(typeRendering: TypeRenderingStrategy = FullyQua
         knownMarkers: Iterable<Marker>,
         readDfMethod: DefaultReadDfMethod?,
         fieldNameNormalizer: NameNormalizer,
+        asDataClass: Boolean
     ): CodeGenResult {
-        val context = SchemaProcessor.create(name, knownMarkers, fieldNameNormalizer)
+        val context = SchemaProcessor.create(name, if (asDataClass) emptyList() else knownMarkers, fieldNameNormalizer)
         val marker = context.process(schema, isOpen, visibility)
         val declarations = mutableListOf<Code>()
         context.generatedMarkers.forEach { itMarker ->
-            declarations.add(generateInterface(itMarker, fields, readDfMethod.takeIf { marker == itMarker }))
+            val declaration = if (asDataClass) {
+                generateClasses(itMarker)
+            } else {
+                generateInterface(itMarker, fields, readDfMethod.takeIf { marker == itMarker })
+            }
+            declarations.add(declaration)
             if (extensionProperties) {
                 declarations.add(generateExtensionProperties(itMarker, withNullable = false))
             }
@@ -479,17 +485,7 @@ internal class CodeGeneratorImpl(typeRendering: TypeRenderingStrategy = FullyQua
                 .joinToString() else ""
         val resultDeclarations = mutableListOf<String>()
 
-        val fieldsDeclaration = if (fields) marker.fields.map {
-            val override = if (it.overrides) "override " else ""
-            val columnNameAnnotation = if (it.columnName != it.fieldName.quotedIfNeeded) {
-                "    @ColumnName(\"${renderStringLiteral(it.columnName)}\")\n"
-            } else {
-                ""
-            }
-
-            val fieldType = it.renderFieldType()
-            "$columnNameAnnotation    ${propertyVisibility}${override}val ${it.fieldName.quotedIfNeeded}: $fieldType"
-        }.join() else ""
+        val fieldsDeclaration = if (fields) renderFields(marker, propertyVisibility).join() else ""
 
         val readDfMethodDeclaration = readDfMethod?.toDeclaration(marker, propertyVisibility)
 
@@ -515,6 +511,36 @@ internal class CodeGeneratorImpl(typeRendering: TypeRenderingStrategy = FullyQua
         resultDeclarations.add(header + baseInterfacesDeclaration + body)
         return resultDeclarations.join()
     }
+
+    private fun generateClasses(marker: Marker): Code {
+        val annotationName = DataSchema::class.simpleName
+
+        val visibility = renderTopLevelDeclarationVisibility(marker)
+        val propertyVisibility = renderInternalDeclarationVisibility(marker)
+        val header =
+            "@$annotationName\n${visibility}data class ${marker.name}("
+
+        val fieldsDeclaration = renderFields(marker, propertyVisibility).joinToString(",\n")
+        return buildString {
+            appendLine(header)
+            appendLine(fieldsDeclaration)
+            append(")")
+        }
+    }
+
+    private fun renderFields(marker: Marker, propertyVisibility: String): List<String> {
+        return marker.fields.map {
+            val override = if (it.overrides) "override " else ""
+            val columnNameAnnotation = if (it.columnName != it.fieldName.quotedIfNeeded) {
+                "    @ColumnName(\"${renderStringLiteral(it.columnName)}\")\n"
+            } else {
+                ""
+            }
+
+            val fieldType = it.renderFieldType()
+            "$columnNameAnnotation    ${propertyVisibility}${override}val ${it.fieldName.quotedIfNeeded}: $fieldType"
+        }
+    }
 }
 
 public fun CodeWithConverter.toStandaloneSnippet(packageName: String, additionalImports: List<String>): String =

core/src/test/kotlin/org/jetbrains/kotlinx/dataframe/codeGen/CodeGenerationTests.kt

@@ -11,8 +11,10 @@ import org.jetbrains.kotlinx.dataframe.DataColumn
 import org.jetbrains.kotlinx.dataframe.DataRow
 import org.jetbrains.kotlinx.dataframe.api.dataFrameOf
 import org.jetbrains.kotlinx.dataframe.api.dropNulls
+import org.jetbrains.kotlinx.dataframe.api.groupBy
 import org.jetbrains.kotlinx.dataframe.api.move
 import org.jetbrains.kotlinx.dataframe.api.schema
+import org.jetbrains.kotlinx.dataframe.api.toDataFrame
 import org.jetbrains.kotlinx.dataframe.api.under
 import org.jetbrains.kotlinx.dataframe.columns.ColumnGroup
 import org.jetbrains.kotlinx.dataframe.impl.codeGen.ReplCodeGeneratorImpl
@@ -305,6 +307,27 @@ class CodeGenerationTests : BaseTest() {
         }
     }
 
+    @Test
+    fun `check method generateDataClasses`() {
+        val code = typed.groupBy { name }.toDataFrame().generateDataClasses(extensionProperties = false)
+
+        code shouldBe """
+            @DataSchema
+            data class Person1(
+                val age: Int,
+                val city: String?,
+                val name: String,
+                val weight: Int?
+            )
+
+            @DataSchema
+            data class Person(
+                val group: List<Person1>,
+                val name: String
+            )
+        """.trimIndent()
+    }
+
     @Test
     fun patterns() {
         """^[\d]""".toRegex().matches("3fds")