@@ -2,12 +2,17 @@ package org.jetbrains.kotlinx.dataframe.codeGen
2
2
3
3
import io.kotest.matchers.shouldBe
4
4
import io.kotest.matchers.string.shouldNotBeEmpty
5
+ import org.jetbrains.kotlinx.dataframe.AnyRow
5
6
import org.jetbrains.kotlinx.dataframe.ColumnsScope
6
7
import org.jetbrains.kotlinx.dataframe.DataColumn
7
8
import org.jetbrains.kotlinx.dataframe.DataRow
8
9
import 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
9
13
import org.jetbrains.kotlinx.dataframe.api.dataFrameOf
10
14
import org.jetbrains.kotlinx.dataframe.api.filter
15
+ import org.jetbrains.kotlinx.dataframe.api.first
11
16
import org.jetbrains.kotlinx.dataframe.api.select
12
17
import org.jetbrains.kotlinx.dataframe.impl.codeGen.ReplCodeGenerator
13
18
import org.jetbrains.kotlinx.dataframe.impl.codeGen.ReplCodeGeneratorImpl
@@ -81,7 +86,12 @@ class ReplCodeGenTests : BaseTest() {
81
86
val expected =
82
87
"""
83
88
@DataSchema
84
- interface $marker { }
89
+ interface $marker {
90
+ val age: Int
91
+ val city: String?
92
+ val name: String
93
+ val weight: Int?
94
+ }
85
95
86
96
val $dfName <$marker >.age: $dataCol <$intName > @JvmName("${marker} _age") get() = this["age"] as $dataCol <$intName >
87
97
val $dfRowName <$marker >.age: $intName @JvmName("${marker} _age") get() = this["age"] as $intName
@@ -103,7 +113,9 @@ class ReplCodeGenTests : BaseTest() {
103
113
val expected3 =
104
114
"""
105
115
@DataSchema
106
- interface $marker3 : $markerFull { }
116
+ interface $marker3 : $markerFull {
117
+ override val city: String
118
+ }
107
119
108
120
val $dfName <$marker3 >.city: $dataCol <$stringName > @JvmName("${marker3} _city") get() = this["city"] as $dataCol <$stringName >
109
121
val $dfRowName <$marker3 >.city: $stringName @JvmName("${marker3} _city") get() = this["city"] as $stringName
@@ -120,7 +132,9 @@ class ReplCodeGenTests : BaseTest() {
120
132
val expected5 =
121
133
"""
122
134
@DataSchema
123
- interface $marker5 : $markerFull { }
135
+ interface $marker5 : $markerFull {
136
+ override val weight: Int
137
+ }
124
138
125
139
val $dfName <$marker5 >.weight: $dataCol <$intName > @JvmName("${marker5} _weight") get() = this["weight"] as $dataCol <$intName >
126
140
val $dfRowName <$marker5 >.weight: $intName @JvmName("${marker5} _weight") get() = this["weight"] as $intName
@@ -163,7 +177,10 @@ class ReplCodeGenTests : BaseTest() {
163
177
val expected =
164
178
"""
165
179
@DataSchema
166
- interface $marker : ${Test2 ._DataFrameType ::class .qualifiedName} { }
180
+ interface $marker : ${Test2 ._DataFrameType ::class .qualifiedName} {
181
+ val city: String?
182
+ val weight: Int?
183
+ }
167
184
168
185
val $dfName <$marker >.city: $dataCol <$stringName ?> @JvmName("${marker} _city") get() = this["city"] as $dataCol <$stringName ?>
169
186
val $dfRowName <$marker >.city: $stringName ? @JvmName("${marker} _city") get() = this["city"] as $stringName ?
@@ -218,4 +235,69 @@ class ReplCodeGenTests : BaseTest() {
218
235
val c = repl.process(Test4 .df, Test4 ::df)
219
236
""" val .*ColumnsScope<\w*>.a:""" .toRegex().findAll(c.declarations).count() shouldBe 1
220
237
}
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
+ }
221
303
}
0 commit comments