Working exclude classes in toDataFrame

Author: koperagen

core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/impl/api/toDataFrame.kt

@@ -186,6 +186,7 @@ internal class CreateDataFrameDslImpl<T>(
             excludes = dsl.excludeProperties,
             preserveClasses = dsl.preserveClasses,
             preserveProperties = dsl.preserveProperties,
+            excludeClasses = dsl.excludeClasses,
             maxDepth = maxDepth,
         )
         df.columns().forEach {
@@ -212,6 +213,7 @@ internal fun convertToDataFrame(
     excludes: Set<KCallable<*>>,
     preserveClasses: Set<KClass<*>>,
     preserveProperties: Set<KCallable<*>>,
+    excludeClasses: Set<KClass<*>>,
     maxDepth: Int,
 ): AnyFrame {
     val clazz = type.classifierOrAny()
@@ -231,10 +233,12 @@ internal fun convertToDataFrame(
     val columns = properties.mapNotNull {
         val property = it
         if (excludes.contains(property)) return@mapNotNull null
+        val klass = it.returnType.classifier as? KClass<*>
+        if (excludeClasses.contains(klass)) return@mapNotNull null
 
         class ValueClassConverter(val unbox: Method, val box: Method)
 
-        val valueClassConverter = (it.returnType.classifier as? KClass<*>)?.let { kClass ->
+        val valueClassConverter = klass?.let { kClass ->
             if (!kClass.isValue) return@let null
 
             val constructor = requireNotNull(kClass.primaryConstructor) {
@@ -374,6 +378,7 @@ internal fun convertToDataFrame(
                                             excludes = excludes,
                                             preserveClasses = preserveClasses,
                                             preserveProperties = preserveProperties,
+                                            excludeClasses = excludeClasses,
                                             maxDepth = maxDepth - 1,
                                         )
                                     }
@@ -392,6 +397,7 @@ internal fun convertToDataFrame(
                     excludes = excludes,
                     preserveClasses = preserveClasses,
                     preserveProperties = preserveProperties,
+                    excludeClasses = excludeClasses,
                     maxDepth = maxDepth - 1,
                 )
                 DataColumn.createColumnGroup(name = it.columnName, df = df)

core/src/test/kotlin/org/jetbrains/kotlinx/dataframe/api/toDataFrame.kt

@@ -156,8 +156,11 @@ class CreateDataFrameTests {
         df.a[0].v shouldBe 7
 
         val df2 = data.toDataFrame {
-            preserve(B::row)
-            properties { preserve(DataFrame::class) }
+            preserve(B::row) // this@toDataFrame: TraversePropertiesDsl - works
+            properties {
+                preserve(DataFrame::class) // this@properties: TraversePropertiesDsl - always works
+            }
+            preserve(B::row) // this@toDataFrame: TraversePropertiesDsl - doesn't
         }
         df2.frame.kind shouldBe ColumnKind.Value
         df2.frame.type shouldBe typeOf<DataFrame<A>>()
@@ -186,6 +189,24 @@ class CreateDataFrameTests {
         res.schema() shouldBe data.toDataFrame(maxDepth = 0).schema()
     }
 
+    class NestedExcludeClasses(val s: String, val list1: List<String>)
+
+    class ExcludeClasses(val i: Int, val list: List<Int>, val nested: NestedExcludeClasses)
+
+    @Test
+    fun `exclude classes`() {
+        val list = listOf(
+            ExcludeClasses(1, listOf(1, 2, 3), NestedExcludeClasses("str", listOf("foo", "bar"))),
+        )
+        val df = list.toDataFrame {
+            properties(maxDepth = 2) {
+                exclude(List::class)
+            }
+        }
+
+        df shouldBe list.toDataFrame(maxDepth = 2).remove { "list" and "nested"["list1"] }
+    }
+
     enum class DummyEnum { A }
 
     @Test