Merge pull request #193 from Kotlin/jupyter-inherit-from-dataclass-bug

fix for "Jupyter codegen: attempt to inherit from data class"

Author: Jolanrensen

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

@@ -37,7 +37,10 @@ internal object MarkersExtractor {
     fun get(markerClass: KClass<*>, nullableProperties: Boolean = false): Marker =
         cache.getOrPut(Pair(markerClass, nullableProperties)) {
             val fields = getFields(markerClass, nullableProperties)
-            val isOpen = markerClass.findAnnotation<DataSchema>()?.isOpen ?: false
+            val isOpen = !markerClass.isSealed &&
+                markerClass.java.isInterface &&
+                markerClass.findAnnotation<DataSchema>()?.isOpen == true
+
             val baseSchemas = markerClass.superclasses.filter { it != Any::class }.map { get(it, nullableProperties) }
             Marker(
                 name = markerClass.qualifiedName ?: markerClass.simpleName!!,

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

@@ -6,13 +6,22 @@ import org.jetbrains.kotlinx.dataframe.codeGen.CodeWithConverter
 import org.jetbrains.kotlinx.dataframe.impl.codeGen.ReplCodeGeneratorImpl
 import org.jetbrains.kotlinx.jupyter.api.Code
 import kotlin.reflect.KClass
+import kotlin.reflect.KMutableProperty
 import kotlin.reflect.KProperty
 
 internal interface ReplCodeGenerator {
 
-    fun process(df: AnyFrame, property: KProperty<*>? = null): CodeWithConverter
+    fun process(
+        df: AnyFrame,
+        property: KProperty<*>? = null,
+        isMutable: Boolean = property is KMutableProperty?,
+    ): CodeWithConverter
 
-    fun process(row: AnyRow, property: KProperty<*>? = null): CodeWithConverter
+    fun process(
+        row: AnyRow,
+        property: KProperty<*>? = null,
+        isMutable: Boolean = property is KMutableProperty?,
+    ): CodeWithConverter
 
     fun process(markerClass: KClass<*>): Code
 

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

@@ -4,7 +4,10 @@ import org.jetbrains.dataframe.impl.codeGen.CodeGenerator
 import org.jetbrains.kotlinx.dataframe.DataFrame
 import org.jetbrains.kotlinx.dataframe.api.schema
 
-public inline fun <reified T> DataFrame<T>.generateCode(fields: Boolean = true, extensionProperties: Boolean = true): String {
+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"
@@ -15,16 +18,16 @@ public fun <T> DataFrame<T>.generateCode(
     markerName: String,
     fields: Boolean = true,
     extensionProperties: Boolean = true,
-    visibility: MarkerVisibility = MarkerVisibility.IMPLICIT_PUBLIC
+    visibility: MarkerVisibility = MarkerVisibility.IMPLICIT_PUBLIC,
 ): String {
     val codeGen = CodeGenerator.create()
     return codeGen.generate(
-        schema(),
-        markerName,
+        schema = schema(),
+        name = markerName,
         fields = fields,
         extensionProperties = extensionProperties,
         isOpen = true,
-        visibility
+        visibility = visibility,
     ).code.declarations
 }
 
@@ -34,7 +37,7 @@ public inline fun <reified T> DataFrame<T>.generateInterfaces(): String = genera
 )
 
 public fun <T> DataFrame<T>.generateInterfaces(markerName: String): String = generateCode(
-    markerName,
+    markerName = markerName,
     fields = true,
     extensionProperties = false
 )

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

@@ -37,8 +37,8 @@ import org.jetbrains.kotlinx.jupyter.api.Code
 
 private fun renderNullability(nullable: Boolean) = if (nullable) "?" else ""
 
-internal fun getRequiredMarkers(schema: DataFrameSchema, markers: Iterable<Marker>) = markers
-    .filter { it.isOpen && it.schema.compare(schema).isSuperOrEqual() }
+internal fun Iterable<Marker>.filterRequiredForSchema(schema: DataFrameSchema) =
+    filter { it.isOpen && it.schema.compare(schema).isSuperOrEqual() }
 
 internal val charsToQuote = """[ `(){}\[\].<>'"/|\\!?@:;%^&*#$-]""".toRegex()
 

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

@@ -19,6 +19,7 @@ import kotlin.reflect.KClass
 import kotlin.reflect.KMutableProperty
 import kotlin.reflect.KProperty
 import kotlin.reflect.KType
+import kotlin.reflect.KVisibility
 import kotlin.reflect.full.findAnnotation
 import kotlin.reflect.full.superclasses
 import kotlin.reflect.jvm.jvmErasure
@@ -44,11 +45,12 @@ internal class ReplCodeGeneratorImpl : ReplCodeGenerator {
             else -> null
         }
 
-    override fun process(row: AnyRow, property: KProperty<*>?) = process(row.df(), property)
+    override fun process(row: AnyRow, property: KProperty<*>?, isMutable: Boolean) =
+        process(row.df(), property, isMutable)
 
-    override fun process(df: AnyFrame, property: KProperty<*>?): CodeWithConverter {
+    override fun process(df: AnyFrame, property: KProperty<*>?, isMutable: Boolean): CodeWithConverter {
         var targetSchema = df.schema()
-        var isMutable = false
+        var isMutable = isMutable
 
         if (property != null) {
             val wasProcessedBefore = property in registeredProperties
@@ -66,8 +68,7 @@ internal class ReplCodeGeneratorImpl : ReplCodeGenerator {
                     // for mutable properties we do strong typing only at the first processing, after that we allow its type to be more general than actual data frame type
                     if (wasProcessedBefore || columnSchema == targetSchema) {
                         // property scheme is valid for current data frame, but we should also check that all compatible open markers are implemented by it
-                        val requiredBaseMarkers =
-                            getRequiredMarkers(columnSchema, registeredMarkers.values)
+                        val requiredBaseMarkers = registeredMarkers.values.filterRequiredForSchema(columnSchema)
                         if (requiredBaseMarkers.any() && requiredBaseMarkers.all { currentMarker.implements(it) }) {
                             return CodeWithConverter.Empty
                         }
@@ -78,7 +79,7 @@ internal class ReplCodeGeneratorImpl : ReplCodeGenerator {
             }
         }
 
-        return generate(targetSchema, markerInterfacePrefix, isMutable)
+        return generate(schema = targetSchema, name = markerInterfacePrefix, isOpen = isMutable)
     }
 
     fun generate(
@@ -87,13 +88,15 @@ internal class ReplCodeGeneratorImpl : ReplCodeGenerator {
         isOpen: Boolean,
     ): CodeWithConverter {
         val result = generator.generate(
-            schema,
-            name,
+            schema = schema,
+            name = name,
             fields = false,
             extensionProperties = true,
-            isOpen,
-            MarkerVisibility.IMPLICIT_PUBLIC,
-            registeredMarkers.values
+            isOpen = isOpen,
+            visibility = MarkerVisibility.IMPLICIT_PUBLIC,
+            knownMarkers = registeredMarkers
+                .filterKeys { it.visibility != KVisibility.PRIVATE }
+                .values,
         )
 
         result.newMarkers.forEach {
@@ -121,12 +124,12 @@ internal class ReplCodeGeneratorImpl : ReplCodeGenerator {
                 if (baseClassNames == tempBaseClassNames) {
                     val newBaseMarkers = baseClasses.map { resolve(it) }
                     val newMarker = Marker(
-                        clazz.qualifiedName!!,
-                        temp.isOpen,
-                        temp.fields,
-                        newBaseMarkers,
-                        MarkerVisibility.IMPLICIT_PUBLIC,
-                        clazz
+                        name = clazz.qualifiedName!!,
+                        isOpen = temp.isOpen,
+                        fields = temp.fields,
+                        superMarkers = newBaseMarkers,
+                        visibility = MarkerVisibility.IMPLICIT_PUBLIC,
+                        klass = clazz,
                     )
                     registeredMarkers[markerClass] = newMarker
                     generatedMarkers.remove(temp.name)

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

@@ -22,7 +22,7 @@ internal class SchemaProcessorImpl(
     override val generatedMarkers = mutableListOf<Marker>()
 
     private fun DataFrameSchema.getAllSuperMarkers() = registeredMarkers
-        .filter { it.schema.compare(this).isSuperOrEqual() }
+        .filter { it.isOpen && it.schema.compare(this).isSuperOrEqual() }
 
     private fun List<Marker>.onlyLeafs(): List<Marker> {
         val skip = flatMap { it.allSuperMarkers.keys }.toSet()
@@ -133,7 +133,7 @@ internal class SchemaProcessorImpl(
         return Marker(name, isOpen, fields, baseMarkers.onlyLeafs(), visibility, emptyList(), emptyList())
     }
 
-    private fun DataFrameSchema.getRequiredMarkers() = getRequiredMarkers(this, registeredMarkers)
+    private fun DataFrameSchema.getRequiredMarkers() = registeredMarkers.filterRequiredForSchema(this)
 
     override fun process(
         schema: DataFrameSchema,

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

@@ -26,7 +26,7 @@ class ReplCodeGenTests : BaseTest() {
     val stringName = String::class.simpleName!!
 
     class Test1 {
-        @DataSchema(isOpen = false)
+        @DataSchema
         interface _DataFrameType
 
         @DataSchema(isOpen = false)
@@ -37,10 +37,10 @@ class ReplCodeGenTests : BaseTest() {
     }
 
     class Test2 {
-        @DataSchema(isOpen = false)
+        @DataSchema
         interface _DataFrameType
 
-        @DataSchema(isOpen = false)
+        @DataSchema
         interface _DataFrameType1
 
         @DataSchema(isOpen = false)
@@ -68,13 +68,13 @@ class ReplCodeGenTests : BaseTest() {
     @Test
     fun `process derived markers`() {
         val repl = ReplCodeGenerator.create()
-        val code = repl.process(df).declarations
+        val code = repl.process(df, isMutable = true).declarations
 
         val marker = ReplCodeGeneratorImpl.markerInterfacePrefix
         val markerFull = Test1._DataFrameType::class.qualifiedName!!
 
         val expected = """
-            @DataSchema(isOpen = false)
+            @DataSchema
             interface $marker
             
             val $dfName<$marker>.age: $dataCol<$intName> @JvmName("${marker}_age") get() = this["age"] as $dataCol<$intName>
@@ -100,7 +100,7 @@ class ReplCodeGenTests : BaseTest() {
         code2 shouldBe ""
 
         val df3 = typed.filter { city != null }
-        val code3 = repl.process(df3).declarations
+        val code3 = repl.process(df3, isMutable = false).declarations
         val marker3 = marker + "1"
         val expected3 = """
             @DataSchema(isOpen = false)
@@ -118,7 +118,7 @@ class ReplCodeGenTests : BaseTest() {
         code4 shouldBe ""
 
         val df5 = typed.filter { weight != null }
-        val code5 = repl.process(df5).declarations
+        val code5 = repl.process(df5, isMutable = false).declarations
         val marker5 = marker + "2"
         val expected5 = """
             @DataSchema(isOpen = false)
@@ -149,7 +149,7 @@ class ReplCodeGenTests : BaseTest() {
             
         """.trimIndent()
 
-        val code = repl.process(typed).declarations.trimIndent()
+        val code = repl.process(typed, isMutable = false).declarations.trimIndent()
         code shouldBe expected
     }
 
@@ -176,7 +176,7 @@ class ReplCodeGenTests : BaseTest() {
             val $dfRowName<$marker?>.weight: $intName? @JvmName("Nullable${marker}_weight") get() = this["weight"] as $intName?
         """.trimIndent()
 
-        val code = repl.process(typed).declarations.trimIndent()
+        val code = repl.process(typed, isMutable = false).declarations.trimIndent()
         code shouldBe expected
     }
 

core/src/test/kotlin/org/jetbrains/kotlinx/dataframe/jupyter/JupyterCodegenTests.kt

@@ -1,10 +1,12 @@
 package org.jetbrains.kotlinx.dataframe.jupyter
 
 import io.kotest.assertions.throwables.shouldNotThrowAny
+import io.kotest.matchers.should
 import io.kotest.matchers.shouldBe
 import io.kotest.matchers.types.shouldBeInstanceOf
 import org.intellij.lang.annotations.Language
 import org.jetbrains.kotlinx.dataframe.AnyFrame
+import org.jetbrains.kotlinx.dataframe.api.isNotEmpty
 import org.jetbrains.kotlinx.dataframe.columns.ValueColumn
 import org.jetbrains.kotlinx.dataframe.type
 import org.jetbrains.kotlinx.jupyter.api.MimeTypedResult
@@ -29,6 +31,90 @@ class JupyterCodegenTests : JupyterReplTestCase() {
         res2["value"].type shouldBe typeOf<List<Any?>>()
     }
 
+    @Test
+    fun `Don't inherit from data class`() {
+        @Language("kts")
+        val res1 = exec(
+            """
+            @DataSchema
+            data class A(val a: Int)
+            """.trimIndent()
+        )
+
+        @Language("kts")
+        val res2 = execRaw(
+            """
+            val df = dataFrameOf("a", "b")(1, 2)
+            df
+            """.trimIndent()
+        )
+
+        (res2 as AnyFrame).should { it.isNotEmpty() }
+    }
+
+    @Test
+    fun `Don't inherit from non open class`() {
+        @Language("kts")
+        val res1 = exec(
+            """
+            @DataSchema
+            class A(val a: Int)
+            """.trimIndent()
+        )
+
+        @Language("kts")
+        val res2 = execRaw(
+            """
+            val df = dataFrameOf("a", "b")(1, 2)
+            df
+            """.trimIndent()
+        )
+
+        (res2 as AnyFrame).should { it.isNotEmpty() }
+    }
+
+    @Test
+    fun `Don't inherit from open class`() {
+        @Language("kts")
+        val res1 = exec(
+            """
+            @DataSchema
+            open class A(val a: Int)
+            """.trimIndent()
+        )
+
+        @Language("kts")
+        val res2 = execRaw(
+            """
+            val df = dataFrameOf("a", "b")(1, 2)
+            df
+            """.trimIndent()
+        )
+
+        (res2 as AnyFrame).should { it.isNotEmpty() }
+    }
+
+    @Test
+    fun `Do inherit from open interface`() {
+        @Language("kts")
+        val res1 = exec(
+            """
+            @DataSchema
+            interface A { val a: Int }
+            """.trimIndent()
+        )
+
+        @Language("kts")
+        val res2 = execRaw(
+            """
+            val df = dataFrameOf("a", "b")(1, 2)
+            df
+            """.trimIndent()
+        )
+
+        (res2 as AnyFrame).should { it.isNotEmpty() }
+    }
+
     @Test
     fun `codegen for enumerated frames`() {
         @Language("kts")