Skip to content

Commit 4d31ef7

Browse files
koperagenkoperagen
committed
Render FrameColumn<T> as List<T> in generated code to enhance listOf(generated classes).toDataFrame() workflow
1 parent e6f744f commit 4d31ef7

File tree

9 files changed

+83
-31
lines changed

9 files changed

+83
-31
lines changed

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

Lines changed: 10 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -6,8 +6,12 @@ import org.jetbrains.kotlinx.dataframe.schema.ColumnSchema
66

77
public sealed interface FieldType {
88
public data class ValueFieldType(public val typeFqName: String) : FieldType
9-
public data class FrameFieldType(public val markerName: String, public val nullable: Boolean) : FieldType
10-
public data class GroupFieldType(public val markerName: String) : FieldType
9+
public data class FrameFieldType(
10+
public val markerName: String,
11+
public val nullable: Boolean,
12+
public val renderAsList: Boolean
13+
) : FieldType
14+
public data class GroupFieldType(public val markerName: String, public val renderAsObject: Boolean) : FieldType
1115
}
1216

1317
/**
@@ -36,8 +40,8 @@ private fun String.toNullable() = if (this.last() == '?' || this == "*") this el
3640
public fun FieldType.toNullable(): FieldType =
3741
if (isNotNullable()) {
3842
when (this) {
39-
is FieldType.FrameFieldType -> FieldType.FrameFieldType(markerName.toNullable(), nullable)
40-
is FieldType.GroupFieldType -> FieldType.GroupFieldType(markerName.toNullable())
43+
is FieldType.FrameFieldType -> FieldType.FrameFieldType(markerName.toNullable(), nullable, renderAsList)
44+
is FieldType.GroupFieldType -> FieldType.GroupFieldType(markerName.toNullable(), renderAsObject)
4145
is FieldType.ValueFieldType -> FieldType.ValueFieldType(typeFqName.toNullable())
4246
}
4347
} else this
@@ -55,13 +59,15 @@ public fun FieldType.toNotNullable(): FieldType =
5559
else it.removeSuffix("?")
5660
},
5761
nullable = nullable,
62+
renderAsList
5863
)
5964

6065
is FieldType.GroupFieldType -> FieldType.GroupFieldType(
6166
markerName = markerName.let {
6267
if (it == "*") "Any"
6368
else it.removeSuffix("?")
6469
},
70+
renderAsObject
6571
)
6672

6773
is FieldType.ValueFieldType -> FieldType.ValueFieldType(
@@ -92,7 +98,6 @@ public class ValidFieldName private constructor(private val identifier: String,
9298
return identifier
9399
}
94100

95-
96101
public companion object {
97102
public fun of(name: String): ValidFieldName {
98103
val needsQuote = name.needsQuoting()

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

Lines changed: 37 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -17,14 +17,33 @@ import kotlin.reflect.full.withNullability
1717
import kotlin.reflect.jvm.jvmErasure
1818
import kotlin.reflect.typeOf
1919

20-
internal fun KType.shouldBeConvertedToFrameColumn(): Boolean = when (jvmErasure) {
21-
DataFrame::class -> true
22-
List::class -> arguments[0].type?.jvmErasure?.hasAnnotation<DataSchema>() == true
23-
else -> false
20+
internal fun KType.getFieldKind(): FieldKind = when {
21+
jvmErasure == DataFrame::class -> Frame
22+
jvmErasure == List::class && (arguments[0].type?.jvmErasure?.hasAnnotation<DataSchema>() == true) -> ListToFrame
23+
jvmErasure == DataRow::class -> Group
24+
jvmErasure.hasAnnotation<DataSchema>() -> ObjectToGroup
25+
else -> Default
2426
}
2527

26-
internal fun KType.shouldBeConvertedToColumnGroup(): Boolean = jvmErasure.let {
27-
it == DataRow::class || it.hasAnnotation<DataSchema>()
28+
internal sealed interface FieldKind {
29+
val shouldBeConvertedToColumnGroup: Boolean get() = false
30+
val shouldBeConvertedToFrameColumn: Boolean get() = false
31+
}
32+
internal data object Frame : FieldKind {
33+
override val shouldBeConvertedToFrameColumn: Boolean = true
34+
}
35+
internal data object ListToFrame : FieldKind {
36+
override val shouldBeConvertedToFrameColumn: Boolean = true
37+
}
38+
39+
internal data object Default : FieldKind
40+
41+
internal data object Group : FieldKind {
42+
override val shouldBeConvertedToColumnGroup: Boolean = true
43+
}
44+
45+
internal data object ObjectToGroup : FieldKind {
46+
override val shouldBeConvertedToColumnGroup: Boolean = true
2847
}
2948

3049
private fun String.toNullable(): String = if (endsWith("?")) this else "$this?"
@@ -62,18 +81,26 @@ internal object MarkersExtractor {
6281
val type = it.returnType
6382
val fieldType: FieldType
6483
val clazz = type.jvmErasure
84+
val fieldKind = type.getFieldKind()
6585
val columnSchema = when {
66-
type.shouldBeConvertedToColumnGroup() -> {
86+
fieldKind.shouldBeConvertedToColumnGroup -> {
6787
val nestedType = if (clazz == DataRow::class) type.arguments[0].type ?: typeOf<Any?>() else type
6888
val marker = get(nestedType.jvmErasure, nullableProperties || type.isMarkedNullable)
69-
fieldType = FieldType.GroupFieldType(marker.name)
89+
fieldType = FieldType.GroupFieldType(
90+
marker.name,
91+
renderAsObject = fieldKind is ObjectToGroup
92+
)
7093
ColumnSchema.Group(marker.schema, nestedType)
7194
}
7295

73-
type.shouldBeConvertedToFrameColumn() -> {
96+
fieldKind.shouldBeConvertedToFrameColumn -> {
7497
val frameType = type.arguments[0].type ?: typeOf<Any?>()
7598
val marker = get(frameType.jvmErasure, nullableProperties || type.isMarkedNullable)
76-
fieldType = FieldType.FrameFieldType(marker.name, type.isMarkedNullable || nullableProperties)
99+
fieldType = FieldType.FrameFieldType(
100+
marker.name,
101+
type.isMarkedNullable || nullableProperties,
102+
renderAsList = fieldKind is ListToFrame
103+
)
77104
ColumnSchema.Frame(marker.schema, type.isMarkedNullable, frameType)
78105
}
79106

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

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -11,8 +11,7 @@ import org.jetbrains.kotlinx.dataframe.api.TraversePropertiesDsl
1111
import org.jetbrains.kotlinx.dataframe.api.concat
1212
import org.jetbrains.kotlinx.dataframe.api.dataFrameOf
1313
import org.jetbrains.kotlinx.dataframe.api.toDataFrameFromPairs
14-
import org.jetbrains.kotlinx.dataframe.codeGen.shouldBeConvertedToColumnGroup
15-
import org.jetbrains.kotlinx.dataframe.codeGen.shouldBeConvertedToFrameColumn
14+
import org.jetbrains.kotlinx.dataframe.codeGen.getFieldKind
1615
import org.jetbrains.kotlinx.dataframe.columns.ColumnPath
1716
import org.jetbrains.kotlinx.dataframe.impl.asList
1817
import org.jetbrains.kotlinx.dataframe.impl.columnName
@@ -248,11 +247,12 @@ internal fun convertToDataFrame(
248247
}
249248
}
250249
val kClass = returnType.classifier as KClass<*>
250+
val fieldKind = returnType.getFieldKind()
251251

252252
val shouldCreateValueCol = (
253253
maxDepth <= 0 &&
254-
!returnType.shouldBeConvertedToFrameColumn() &&
255-
!returnType.shouldBeConvertedToColumnGroup()
254+
!fieldKind.shouldBeConvertedToFrameColumn &&
255+
!fieldKind.shouldBeConvertedToColumnGroup
256256
) ||
257257
kClass == Any::class ||
258258
kClass in preserveClasses ||

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

Lines changed: 18 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -146,11 +146,17 @@ internal object FullyQualifiedNames : TypeRenderingStrategy {
146146
is FieldType.ValueFieldType ->
147147
fieldType.typeFqName
148148

149-
is FieldType.GroupFieldType ->
149+
is FieldType.GroupFieldType -> if (fieldType.renderAsObject) {
150150
fieldType.markerName
151+
} else {
152+
renderAccessorFieldType()
153+
}
151154

152-
is FieldType.FrameFieldType ->
153-
"$dataFrame<${fieldType.markerName}>${renderNullability(fieldType.nullable)}"
155+
is FieldType.FrameFieldType -> if (fieldType.renderAsList) {
156+
"List<${fieldType.markerName}>${renderNullability(fieldType.nullable)}"
157+
} else {
158+
renderAccessorFieldType()
159+
}
154160
}
155161
}
156162

@@ -196,11 +202,17 @@ internal object ShortNames : TypeRenderingStrategy {
196202
is FieldType.ValueFieldType ->
197203
fieldType.typeFqName.shorten()
198204

199-
is FieldType.GroupFieldType ->
205+
is FieldType.GroupFieldType -> if (fieldType.renderAsObject) {
200206
fieldType.markerName
207+
} else {
208+
renderAccessorFieldType()
209+
}
201210

202-
is FieldType.FrameFieldType ->
203-
"$dataFrame<${fieldType.markerName}>${renderNullability(fieldType.nullable)}"
211+
is FieldType.FrameFieldType -> if (fieldType.renderAsList) {
212+
"List<${fieldType.markerName}>${renderNullability(fieldType.nullable)}"
213+
} else {
214+
renderAccessorFieldType()
215+
}
204216
}
205217

206218
private fun String.shorten() = removeRedundantQualifier(this)

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

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -61,10 +61,11 @@ internal class SchemaProcessorImpl(
6161

6262
fun getFieldType(columnSchema: ColumnSchema): FieldType = when (columnSchema) {
6363
is ColumnSchema.Value -> FieldType.ValueFieldType(columnSchema.type.toString())
64-
is ColumnSchema.Group -> FieldType.GroupFieldType(process(columnSchema.schema, false, visibility).name)
64+
is ColumnSchema.Group -> FieldType.GroupFieldType(process(columnSchema.schema, false, visibility).name, renderAsObject = true)
6565
is ColumnSchema.Frame -> FieldType.FrameFieldType(
6666
process(columnSchema.schema, false, visibility).name,
67-
columnSchema.nullable
67+
columnSchema.nullable,
68+
renderAsList = true
6869
)
6970

7071
else -> throw NotImplementedError()

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

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -50,7 +50,7 @@ internal class ShortNamesRenderingTest : TypeRenderingStrategy by ShortNames {
5050
fun `list parametrized by data schema type`() {
5151
fields.keys.asClue {
5252
fields["c"]!!.renderAccessorFieldType() shouldBe "DataFrame<org.jetbrains.kotlinx.dataframe.internal.codeGen.ShortNamesRenderingTest.DataSchemaMarker>"
53-
fields["c"]!!.renderFieldType() shouldBe "DataFrame<org.jetbrains.kotlinx.dataframe.internal.codeGen.ShortNamesRenderingTest.DataSchemaMarker>"
53+
fields["c"]!!.renderFieldType() shouldBe "List<org.jetbrains.kotlinx.dataframe.internal.codeGen.ShortNamesRenderingTest.DataSchemaMarker>"
5454
}
5555
}
5656

@@ -66,7 +66,7 @@ internal class ShortNamesRenderingTest : TypeRenderingStrategy by ShortNames {
6666
fun `data row`() {
6767
fields.keys.asClue {
6868
fields["e"]!!.renderAccessorFieldType() shouldBe "DataRow<org.jetbrains.kotlinx.dataframe.internal.codeGen.ShortNamesRenderingTest.Marker>"
69-
fields["e"]!!.renderFieldType() shouldBe "org.jetbrains.kotlinx.dataframe.internal.codeGen.ShortNamesRenderingTest.Marker"
69+
fields["e"]!!.renderFieldType() shouldBe "DataRow<org.jetbrains.kotlinx.dataframe.internal.codeGen.ShortNamesRenderingTest.Marker>"
7070
}
7171
}
7272

‎dataframe-openapi/src/main/kotlin/org/jetbrains/kotlinx/dataframe/io/OpenApiMarker.kt

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -133,6 +133,7 @@ internal sealed class OpenApiMarker private constructor(
133133
override fun toFieldType(): FieldType =
134134
FieldType.GroupFieldType(
135135
markerName = name + if (nullable) "?" else "",
136+
renderAsObject = true
136137
)
137138

138139
override fun withName(name: String, prependTopInterfaceName: Boolean): Interface =
@@ -207,6 +208,7 @@ internal sealed class OpenApiMarker private constructor(
207208
FieldType.FrameFieldType(
208209
markerName = name + if (nullable) "?" else "",
209210
nullable = false,
211+
renderAsList = false
210212
)
211213

212214
override fun withName(name: String, prependTopInterfaceName: Boolean): AdditionalPropertyInterface =
@@ -331,6 +333,7 @@ internal sealed class OpenApiMarker private constructor(
331333
override fun toFieldType(): FieldType =
332334
FieldType.GroupFieldType(
333335
markerName = name + if (nullable) "?" else "",
336+
renderAsObject = true
334337
)
335338

336339
override fun withName(name: String, prependTopInterfaceName: Boolean): MarkerAlias =

‎dataframe-openapi/src/main/kotlin/org/jetbrains/kotlinx/dataframe/io/OpenApiType.kt

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -68,6 +68,7 @@ internal sealed class OpenApiType(val name: kotlin.String?) : IsObject {
6868
markerName = marker.name.let {
6969
if (nullable) it.toNullable() else it
7070
},
71+
renderAsObject = true
7172
)
7273
}
7374

@@ -77,6 +78,7 @@ internal sealed class OpenApiType(val name: kotlin.String?) : IsObject {
7778
fun getType(nullable: kotlin.Boolean): FieldType =
7879
FieldType.GroupFieldType(
7980
markerName = (if (nullable) typeOf<DataRow<kotlin.Any?>>() else typeOf<DataRow<kotlin.Any>>()).toString(),
81+
renderAsObject = true
8082
)
8183
}
8284

@@ -93,6 +95,7 @@ internal sealed class OpenApiType(val name: kotlin.String?) : IsObject {
9395
FieldType.FrameFieldType(
9496
markerName = markerName.let { if (nullable) it.toNullable() else it },
9597
nullable = false, // preferring DataFrame<Something?> over DataFrame<Something>?
98+
renderAsList = false
9699
)
97100

98101
/** used for list of AdditionalProperty objects (read as List<DataFrame<MyMarker>>) */

‎plugins/symbol-processor/src/main/kotlin/org/jetbrains/dataframe/ksp/PropertyRenderer.kt

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -50,10 +50,11 @@ internal fun renderExtensions(
5050
FieldType.FrameFieldType(
5151
markerName = type.renderTypeArguments(),
5252
nullable = type.isMarkedNullable,
53+
renderAsList = true
5354
)
5455

55-
type.declaration.isAnnotationPresent(DataSchema::class) -> FieldType.GroupFieldType(type.render())
56-
qualifiedTypeReference == DataFrameNames.DATA_ROW -> FieldType.GroupFieldType(type.renderTypeArguments())
56+
type.declaration.isAnnotationPresent(DataSchema::class) -> FieldType.GroupFieldType(type.render(), renderAsObject = true)
57+
qualifiedTypeReference == DataFrameNames.DATA_ROW -> FieldType.GroupFieldType(type.renderTypeArguments(), renderAsObject = false)
5758
else -> FieldType.ValueFieldType(type.render())
5859
}
5960

0 commit comments

Comments
 (0)