generate a new marker when dataframe marker is not a data schema

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

@@ -7,6 +7,7 @@ import org.jetbrains.kotlinx.dataframe.AnyFrame
 import org.jetbrains.kotlinx.dataframe.AnyRow
 import org.jetbrains.kotlinx.dataframe.DataFrame
 import org.jetbrains.kotlinx.dataframe.DataRow
+import org.jetbrains.kotlinx.dataframe.annotations.DataSchema
 import org.jetbrains.kotlinx.dataframe.api.schema
 import org.jetbrains.kotlinx.dataframe.codeGen.CodeWithConverter
 import org.jetbrains.kotlinx.dataframe.codeGen.Marker
@@ -18,6 +19,7 @@ import kotlin.reflect.KClass
 import kotlin.reflect.KMutableProperty
 import kotlin.reflect.KProperty
 import kotlin.reflect.KType
+import kotlin.reflect.full.findAnnotation
 import kotlin.reflect.full.superclasses
 import kotlin.reflect.jvm.jvmErasure
 
@@ -54,7 +56,9 @@ internal class ReplCodeGeneratorImpl : ReplCodeGenerator {
             isMutable = property is KMutableProperty
 
             // maybe property is already properly typed, let's do some checks
-            val currentMarker = getMarkerClass(property.returnType)?.let { registeredMarkers[it] ?: MarkersExtractor[it] }
+            val currentMarker = getMarkerClass(property.returnType)
+                ?.takeIf { it.findAnnotation<DataSchema>() != null }
+                ?.let { registeredMarkers[it] ?: MarkersExtractor[it] }
             if (currentMarker != null) {
                 // if property is mutable, we need to make sure that its marker type is open in order to let derived data frames be assignable to it
                 if (!isMutable || currentMarker.isOpen) {

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

@@ -1,5 +1,6 @@
 package org.jetbrains.kotlinx.dataframe.jupyter
 
+import io.kotest.assertions.throwables.shouldNotThrowAny
 import io.kotest.matchers.shouldBe
 import io.kotest.matchers.types.shouldBeInstanceOf
 import org.intellij.lang.annotations.Language
@@ -171,4 +172,39 @@ class JupyterCodegenTests : JupyterReplTestCase() {
         )
         res2.shouldBeInstanceOf<Unit>()
     }
+
+    @Test
+    fun `generate a new marker when dataframe marker is not a data schema so that columns are accessible with extensions`() {
+        exec(
+            """
+            enum class State {
+                Idle, Productive, Maintenance
+            }
+
+            class Event(val toolId: String, val state: State, val timestamp: Long)
+
+            val tool1 = "tool_1"
+            val tool2 = "tool_2"
+            val tool3 = "tool_3"
+            val events = listOf(
+                Event(tool1, State.Idle, 0),
+                Event(tool1, State.Productive, 5),
+                Event(tool2, State.Idle, 0),
+                Event(tool2, State.Maintenance, 10),
+                Event(tool2, State.Idle, 20),
+                Event(tool3, State.Idle, 0),
+                Event(tool3, State.Productive, 25),
+            ).toDataFrame()
+            """.trimIndent()
+        )
+        shouldNotThrowAny {
+            exec(
+                """
+                events.toolId
+                events.state
+                events.timestamp
+                """.trimIndent()
+            )
+        }
+    }
 }