@@ -2,12 +2,17 @@ package org.jetbrains.kotlinx.dataframe.codeGen
22
33import io.kotest.matchers.shouldBe
44import io.kotest.matchers.string.shouldNotBeEmpty
5+ import org.jetbrains.kotlinx.dataframe.AnyRow
56import org.jetbrains.kotlinx.dataframe.ColumnsScope
67import org.jetbrains.kotlinx.dataframe.DataColumn
78import org.jetbrains.kotlinx.dataframe.DataRow
89import org.jetbrains.kotlinx.dataframe.annotations.DataSchema
10+ import org.jetbrains.kotlinx.dataframe.api.add
11+ import org.jetbrains.kotlinx.dataframe.api.asFrame
12+ import org.jetbrains.kotlinx.dataframe.api.convert
913import org.jetbrains.kotlinx.dataframe.api.dataFrameOf
1014import org.jetbrains.kotlinx.dataframe.api.filter
15+ import org.jetbrains.kotlinx.dataframe.api.first
1116import org.jetbrains.kotlinx.dataframe.api.select
1217import org.jetbrains.kotlinx.dataframe.impl.codeGen.ReplCodeGenerator
1318import org.jetbrains.kotlinx.dataframe.impl.codeGen.ReplCodeGeneratorImpl
@@ -81,7 +86,12 @@ class ReplCodeGenTests : BaseTest() {
8186 val expected =
8287 """
8388 @DataSchema
84- interface $marker { }
89+ interface $marker {
90+ val age: Int
91+ val city: String?
92+ val name: String
93+ val weight: Int?
94+ }
8595
8696 val $dfName <$marker >.age: $dataCol <$intName > @JvmName("${marker} _age") get() = this["age"] as $dataCol <$intName >
8797 val $dfRowName <$marker >.age: $intName @JvmName("${marker} _age") get() = this["age"] as $intName
@@ -103,7 +113,9 @@ class ReplCodeGenTests : BaseTest() {
103113 val expected3 =
104114 """
105115 @DataSchema
106- interface $marker3 : $markerFull { }
116+ interface $marker3 : $markerFull {
117+ override val city: String
118+ }
107119
108120 val $dfName <$marker3 >.city: $dataCol <$stringName > @JvmName("${marker3} _city") get() = this["city"] as $dataCol <$stringName >
109121 val $dfRowName <$marker3 >.city: $stringName @JvmName("${marker3} _city") get() = this["city"] as $stringName
@@ -120,7 +132,9 @@ class ReplCodeGenTests : BaseTest() {
120132 val expected5 =
121133 """
122134 @DataSchema
123- interface $marker5 : $markerFull { }
135+ interface $marker5 : $markerFull {
136+ override val weight: Int
137+ }
124138
125139 val $dfName <$marker5 >.weight: $dataCol <$intName > @JvmName("${marker5} _weight") get() = this["weight"] as $dataCol <$intName >
126140 val $dfRowName <$marker5 >.weight: $intName @JvmName("${marker5} _weight") get() = this["weight"] as $intName
@@ -163,7 +177,10 @@ class ReplCodeGenTests : BaseTest() {
163177 val expected =
164178 """
165179 @DataSchema
166- interface $marker : ${Test2 ._DataFrameType ::class .qualifiedName} { }
180+ interface $marker : ${Test2 ._DataFrameType ::class .qualifiedName} {
181+ val city: String?
182+ val weight: Int?
183+ }
167184
168185 val $dfName <$marker >.city: $dataCol <$stringName ?> @JvmName("${marker} _city") get() = this["city"] as $dataCol <$stringName ?>
169186 val $dfRowName <$marker >.city: $stringName ? @JvmName("${marker} _city") get() = this["city"] as $stringName ?
@@ -218,4 +235,69 @@ class ReplCodeGenTests : BaseTest() {
218235 val c = repl.process(Test4 .df, Test4 ::df)
219236 """ val .*ColumnsScope<\w*>.a:""" .toRegex().findAll(c.declarations).count() shouldBe 1
220237 }
238+
239+ object Test5 {
240+ @DataSchema(isOpen = false )
241+ interface _DataFrameType1 {
242+ val a: Int
243+ val b: Int
244+ }
245+
246+ val ColumnsScope <_DataFrameType1 >.a: DataColumn <Int >
247+ @JvmName(" _DataFrameType1_a" )
248+ get() = this [" a" ] as DataColumn <Int >
249+ val DataRow <_DataFrameType1 >.a: Int
250+ @JvmName(" _DataFrameType1_a" )
251+ get() = this [" a" ] as Int
252+ val ColumnsScope <_DataFrameType1 >.b: DataColumn <Int >
253+ @JvmName(" _DataFrameType1_b" )
254+ get() = this [" b" ] as DataColumn <Int >
255+ val DataRow <_DataFrameType1 >.b: Int
256+ @JvmName(" _DataFrameType1_b" )
257+ get() = this [" b" ] as Int
258+
259+ @DataSchema
260+ interface _DataFrameType {
261+ val col: String
262+ val leaf: _DataFrameType1
263+ }
264+
265+ val df = dataFrameOf(" col" to listOf (" a" ), " leaf" to listOf (dataFrameOf(" a" )(1 ).first()))
266+ .convert(" leaf" ).cast<AnyRow >().asFrame { it.add(" c" ) { 3 } }
267+ }
268+
269+ @Test
270+ fun `process closed inheritance override` () {
271+ // if ReplCodeGenerator would generate schemas with isOpen = true or with fields = false, _DataFrameType2 could implement _DataFrameType
272+ // but with isOpen = false and fields = true _DataFrameType2 : _DataFrameType produces incorrect override that couldn't be compiled
273+ // so we avoid this relation
274+ val repl = ReplCodeGenerator .create()
275+ repl.process<Test5 ._DataFrameType >()
276+ repl.process<Test5 ._DataFrameType1 >()
277+ val c = repl.process(Test5 .df, Test5 ::df)
278+ c.declarations shouldBe
279+ """
280+ @DataSchema(isOpen = false)
281+ interface _DataFrameType3 {
282+ val a: Int
283+ val c: Int
284+ }
285+
286+ val $dfName <_DataFrameType3>.a: $dataCol <Int> @JvmName("_DataFrameType3_a") get() = this["a"] as $dataCol <Int>
287+ val $dfRowName <_DataFrameType3>.a: Int @JvmName("_DataFrameType3_a") get() = this["a"] as Int
288+ val $dfName <_DataFrameType3>.c: $dataCol <Int> @JvmName("_DataFrameType3_c") get() = this["c"] as $dataCol <Int>
289+ val $dfRowName <_DataFrameType3>.c: Int @JvmName("_DataFrameType3_c") get() = this["c"] as Int
290+
291+ @DataSchema
292+ interface _DataFrameType2 {
293+ val col: String
294+ val leaf: _DataFrameType3
295+ }
296+
297+ val $dfName <_DataFrameType2>.col: $dataCol <String> @JvmName("_DataFrameType2_col") get() = this["col"] as $dataCol <String>
298+ val $dfRowName <_DataFrameType2>.col: String @JvmName("_DataFrameType2_col") get() = this["col"] as String
299+ val $dfName <_DataFrameType2>.leaf: ColumnGroup<_DataFrameType3> @JvmName("_DataFrameType2_leaf") get() = this["leaf"] as ColumnGroup<_DataFrameType3>
300+ val $dfRowName <_DataFrameType2>.leaf: $dfRowName <_DataFrameType3> @JvmName("_DataFrameType2_leaf") get() = this["leaf"] as $dfRowName <_DataFrameType3>
301+ """ .trimIndent()
302+ }
221303}
0 commit comments