Skip to content

Commit 54f3d18

Browse files
JolanrensenJolanrensen
committed
Merge branch 'refs/heads/master' into new-csv-implementation
2 parents 920c1a0 + b3b1f64 commit 54f3d18

14 files changed

+137
-39
lines changed

‎plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/FirDataFrameComponentRegistrar.kt

+2-2
Original file line numberDiff line numberDiff line change
@@ -19,7 +19,7 @@ import org.jetbrains.kotlin.fir.caches.FirCache
1919
import org.jetbrains.kotlin.fir.caches.firCachesFactory
2020
import org.jetbrains.kotlinx.dataframe.plugin.extensions.DataRowSchemaSupertype
2121
import org.jetbrains.kotlinx.dataframe.plugin.extensions.ExpressionAnalysisAdditionalChecker
22-
import org.jetbrains.kotlinx.dataframe.plugin.extensions.ExtensionsGenerator
22+
import org.jetbrains.kotlinx.dataframe.plugin.extensions.TopLevelExtensionsGenerator
2323
import org.jetbrains.kotlinx.dataframe.plugin.extensions.FunctionCallTransformer
2424
import org.jetbrains.kotlinx.dataframe.plugin.extensions.IrBodyFiller
2525
import org.jetbrains.kotlinx.dataframe.plugin.extensions.KotlinTypeFacade
@@ -68,7 +68,7 @@ class FirDataFrameExtensionRegistrar(
6868
) : FirExtensionRegistrar() {
6969
@OptIn(FirExtensionApiInternals::class)
7070
override fun ExtensionRegistrarContext.configurePlugin() {
71-
+::ExtensionsGenerator
71+
+::TopLevelExtensionsGenerator
7272
+::ReturnTypeBasedReceiverInjector
7373
+{ it: FirSession ->
7474
FunctionCallTransformer(path, it, jsonCache(it), schemasDirectory, isTest)

‎plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/analyzeRefinedCallShape.kt

-10
Original file line numberDiff line numberDiff line change
@@ -8,8 +8,6 @@ package org.jetbrains.kotlinx.dataframe.plugin
88
import org.jetbrains.kotlin.fir.expressions.FirExpression
99
import org.jetbrains.kotlin.fir.expressions.FirFunctionCall
1010
import org.jetbrains.kotlin.fir.types.ConeClassLikeType
11-
import org.jetbrains.kotlin.fir.types.ConeKotlinType
12-
import org.jetbrains.kotlin.fir.types.ConeTypeProjection
1311
import org.jetbrains.kotlin.fir.types.classId
1412
import org.jetbrains.kotlin.fir.types.resolvedType
1513
import org.jetbrains.kotlin.name.ClassId
@@ -63,11 +61,3 @@ data class RefinedArgument(val name: Name, val expression: FirExpression) {
6361
return "RefinedArgument(name=$name, expression=${expression})"
6462
}
6563
}
66-
67-
data class SchemaProperty(
68-
val marker: ConeTypeProjection,
69-
val name: String,
70-
val dataRowReturnType: ConeKotlinType,
71-
val columnContainerReturnType: ConeKotlinType,
72-
val override: Boolean = false
73-
)

‎plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/extensions/CallShapeAttribute.kt

+1-1
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
package org.jetbrains.kotlinx.dataframe.plugin.extensions
22

3-
import org.jetbrains.kotlinx.dataframe.plugin.SchemaProperty
3+
import org.jetbrains.kotlinx.dataframe.plugin.extensions.impl.SchemaProperty
44
import org.jetbrains.kotlin.fir.declarations.FirClass
55
import org.jetbrains.kotlin.fir.declarations.FirDeclarationDataKey
66
import org.jetbrains.kotlin.fir.declarations.FirDeclarationDataRegistry

‎plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/extensions/DataFramePlugin.kt

+1-1
Original file line numberDiff line numberDiff line change
@@ -2,4 +2,4 @@ package org.jetbrains.kotlinx.dataframe.plugin.extensions
22

33
import org.jetbrains.kotlin.GeneratedDeclarationKey
44

5-
class DataFramePlugin(val columnName: String?) : GeneratedDeclarationKey()
5+
data object DataFramePlugin : GeneratedDeclarationKey()

‎plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/extensions/FunctionCallTransformer.kt

+5-4
Original file line numberDiff line numberDiff line change
@@ -13,7 +13,7 @@ import org.jetbrains.kotlin.fir.FirSession
1313
import org.jetbrains.kotlin.fir.analysis.checkers.fullyExpandedClassId
1414
import org.jetbrains.kotlin.fir.caches.FirCache
1515
import org.jetbrains.kotlinx.dataframe.plugin.InterpretationErrorReporter
16-
import org.jetbrains.kotlinx.dataframe.plugin.SchemaProperty
16+
import org.jetbrains.kotlinx.dataframe.plugin.extensions.impl.SchemaProperty
1717
import org.jetbrains.kotlinx.dataframe.plugin.analyzeRefinedCallShape
1818
import org.jetbrains.kotlinx.dataframe.plugin.utils.Names
1919
import org.jetbrains.kotlinx.dataframe.plugin.utils.projectOverDataColumnType
@@ -74,6 +74,7 @@ import org.jetbrains.kotlin.name.FqName
7474
import org.jetbrains.kotlin.name.Name
7575
import org.jetbrains.kotlin.text
7676
import org.jetbrains.kotlin.types.Variance
77+
import org.jetbrains.kotlinx.dataframe.plugin.extensions.impl.PropertyName
7778
import org.jetbrains.kotlinx.dataframe.plugin.impl.PluginDataFrameSchema
7879
import org.jetbrains.kotlinx.dataframe.plugin.impl.SimpleCol
7980
import org.jetbrains.kotlinx.dataframe.plugin.impl.SimpleDataColumn
@@ -520,7 +521,7 @@ class FunctionCallTransformer(
520521
isNullable = false
521522
)
522523

523-
SchemaProperty(schema.defaultType(), it.name, dataRowReturnType, columnsContainerReturnType)
524+
SchemaProperty(schema.defaultType(), PropertyName.of(it.name), dataRowReturnType, columnsContainerReturnType)
524525
}
525526

526527
is SimpleFrameColumn -> {
@@ -534,7 +535,7 @@ class FunctionCallTransformer(
534535

535536
SchemaProperty(
536537
marker = schema.defaultType(),
537-
name = it.name,
538+
propertyName = PropertyName.of(it.name),
538539
dataRowReturnType = frameColumnReturnType,
539540
columnContainerReturnType = frameColumnReturnType.toFirResolvedTypeRef()
540541
.projectOverDataColumnType()
@@ -543,7 +544,7 @@ class FunctionCallTransformer(
543544

544545
is SimpleDataColumn -> SchemaProperty(
545546
marker = schema.defaultType(),
546-
name = it.name,
547+
propertyName = PropertyName.of(it.name),
547548
dataRowReturnType = it.type.type(),
548549
columnContainerReturnType = it.type.type().toFirResolvedTypeRef().projectOverDataColumnType()
549550
)

‎plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/extensions/IrBodyFiller.kt

+5-1
Original file line numberDiff line numberDiff line change
@@ -48,6 +48,7 @@ import org.jetbrains.kotlin.ir.types.classOrFail
4848
import org.jetbrains.kotlin.ir.types.classifierOrNull
4949
import org.jetbrains.kotlin.ir.types.getClass
5050
import org.jetbrains.kotlin.ir.util.constructors
51+
import org.jetbrains.kotlin.ir.util.findAnnotation
5152
import org.jetbrains.kotlin.ir.util.parentAsClass
5253
import org.jetbrains.kotlin.ir.util.primaryConstructor
5354
import org.jetbrains.kotlin.ir.util.superTypes
@@ -63,6 +64,7 @@ import org.jetbrains.kotlinx.dataframe.api.schema
6364
import org.jetbrains.kotlinx.dataframe.columns.ColumnGroup
6465
import org.jetbrains.kotlinx.dataframe.plugin.impl.data.IoSchema
6566
import org.jetbrains.kotlinx.dataframe.plugin.impl.data.serialize
67+
import org.jetbrains.kotlinx.dataframe.plugin.utils.Names
6668
import java.io.File
6769

6870
class IrBodyFiller(
@@ -209,7 +211,9 @@ private class DataFrameFileLowering(val context: IrPluginContext) : FileLowering
209211
val call = IrCallImpl(-1, -1, context.irBuiltIns.anyNType, get, 0, 1).also {
210212
val thisSymbol: IrValueSymbol = getter.extensionReceiverParameter?.symbol!!
211213
it.dispatchReceiver = IrGetValueImpl(-1, -1, thisSymbol)
212-
val columName = pluginKey.columnName ?: declaration.name.identifier
214+
val annotation = declaration.annotations.findAnnotation(Names.COLUMN_NAME_ANNOTATION.asSingleFqName())
215+
val columnName = (annotation?.valueArguments?.get(0) as? IrConst<*>)?.value as? String
216+
val columName = columnName ?: declaration.name.identifier
213217
it.putValueArgument(0, IrConstImpl.string(-1, -1, context.irBuiltIns.stringType, columName))
214218
}
215219

‎plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/extensions/TokenGenerator.kt

+11-7
Original file line numberDiff line numberDiff line change
@@ -47,16 +47,16 @@ class TokenGenerator(session: FirSession) : FirDeclarationGenerationExtension(se
4747
val resolvedTypeRef = buildResolvedTypeRef {
4848
type = property.dataRowReturnType
4949
}
50-
val propertyName = Name.identifier(property.name)
51-
propertyName to listOf(buildProperty(resolvedTypeRef, propertyName, k, order = index))
50+
val identifier = property.propertyName.identifier
51+
identifier to listOf(buildProperty(resolvedTypeRef, identifier, k, property.propertyName.columnNameAnnotation, order = index))
5252
}
5353
is CallShapeData.RefinedType -> callShapeData.scopes.associate {
54-
val propertyName = Name.identifier(it.name.identifier.replaceFirstChar { it.lowercaseChar() })
55-
propertyName to listOf(buildProperty(it.defaultType().toFirResolvedTypeRef(), propertyName, k, isScopeProperty = true))
54+
val identifier = Name.identifier(it.name.identifier.replaceFirstChar { it.lowercaseChar() })
55+
identifier to listOf(buildProperty(it.defaultType().toFirResolvedTypeRef(), identifier, k, isScopeProperty = true))
5656
}
5757
is CallShapeData.Scope -> callShapeData.columns.associate { schemaProperty ->
58-
val propertyName = Name.identifier(schemaProperty.name)
59-
val callableId = CallableId(k.classId, propertyName)
58+
val propertyName = schemaProperty.propertyName
59+
val callableId = CallableId(k.classId, propertyName.identifier)
6060
val dataRowExtension = generateExtensionProperty(
6161
callableId = callableId,
6262
symbol = k,
@@ -82,7 +82,7 @@ class TokenGenerator(session: FirSession) : FirDeclarationGenerationExtension(se
8282
symbol = k,
8383
effectiveVisibility = EffectiveVisibility.Local
8484
)
85-
propertyName to listOf(dataRowExtension, columnContainerExtension)
85+
propertyName.identifier to listOf(dataRowExtension, columnContainerExtension)
8686
}
8787
}
8888
}
@@ -109,6 +109,7 @@ class TokenGenerator(session: FirSession) : FirDeclarationGenerationExtension(se
109109
resolvedTypeRef: FirResolvedTypeRef,
110110
propertyName: Name,
111111
k: FirClassSymbol<*>,
112+
columnNameAnnotation: FirAnnotation? = null,
112113
isScopeProperty: Boolean = false,
113114
order: Int? = null,
114115
): FirProperty {
@@ -135,6 +136,9 @@ class TokenGenerator(session: FirSession) : FirDeclarationGenerationExtension(se
135136
argumentMapping = buildAnnotationArgumentMapping()
136137
}
137138
}
139+
columnNameAnnotation?.let {
140+
annotations += it
141+
}
138142
replaceAnnotations(annotations)
139143
}
140144
}

‎plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/extensions/ExtensionsGenerator.kt renamed to ‎plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/extensions/TopLevelExtensionsGenerator.kt

+9-7
Original file line numberDiff line numberDiff line change
@@ -29,11 +29,15 @@ import org.jetbrains.kotlin.fir.types.toSymbol
2929
import org.jetbrains.kotlin.fir.types.toTypeProjection
3030
import org.jetbrains.kotlin.name.CallableId
3131
import org.jetbrains.kotlin.name.FqName
32-
import org.jetbrains.kotlin.name.Name
3332
import org.jetbrains.kotlin.types.Variance
3433
import org.jetbrains.kotlinx.dataframe.annotations.DataSchema
34+
import org.jetbrains.kotlinx.dataframe.plugin.extensions.impl.PropertyName
3535

36-
class ExtensionsGenerator(session: FirSession) : FirDeclarationGenerationExtension(session) {
36+
/**
37+
* extensions inside scope classes are generated here:
38+
* @see org.jetbrains.kotlinx.dataframe.plugin.extensions.TokenGenerator
39+
*/
40+
class TopLevelExtensionsGenerator(session: FirSession) : FirDeclarationGenerationExtension(session) {
3741
private companion object {
3842
val dataSchema = FqName(DataSchema::class.qualifiedName!!)
3943
}
@@ -83,7 +87,7 @@ class ExtensionsGenerator(session: FirSession) : FirDeclarationGenerationExtensi
8387
val columnName = property.getAnnotationByClassId(Names.COLUMN_NAME_ANNOTATION, session)?.let { annotation ->
8488
(annotation.argumentMapping.mapping[Names.COLUMN_NAME_ARGUMENT] as? FirLiteralExpression)?.value as? String?
8589
}
86-
val propertyName = property.name
90+
val name = property.name
8791
val marker = owner.constructType(arrayOf(), isNullable = false).toTypeProjection(Variance.INVARIANT)
8892

8993
val columnGroupProjection: ConeTypeProjection? = if (resolvedReturnTypeRef.coneType.classId?.equals(
@@ -115,8 +119,7 @@ class ExtensionsGenerator(session: FirSession) : FirDeclarationGenerationExtensi
115119
typeArguments = arrayOf(marker),
116120
isNullable = false
117121
),
118-
propertyName = propertyName,
119-
columnName = columnName,
122+
propertyName = PropertyName.of(name, columnName?.let { PropertyName.buildAnnotation(it) }),
120123
returnTypeRef = resolvedReturnTypeRef
121124
)
122125

@@ -138,8 +141,7 @@ class ExtensionsGenerator(session: FirSession) : FirDeclarationGenerationExtensi
138141
typeArguments = arrayOf(marker),
139142
isNullable = false
140143
),
141-
propertyName = propertyName,
142-
columnName = columnName,
144+
propertyName = PropertyName.of(name, columnName?.let { PropertyName.buildAnnotation(it) }),
143145
returnTypeRef = columnReturnType
144146
)
145147
listOf(rowExtension.symbol, columnsContainerExtension.symbol)

‎plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/extensions/impl/PropertyName.kt

+52
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,52 @@
1+
package org.jetbrains.kotlinx.dataframe.plugin.extensions.impl
2+
3+
import org.jetbrains.kotlin.fir.expressions.FirAnnotation
4+
import org.jetbrains.kotlin.fir.expressions.builder.buildAnnotationArgumentMapping
5+
import org.jetbrains.kotlin.fir.expressions.builder.buildLiteralExpression
6+
import org.jetbrains.kotlin.fir.resolve.defaultType
7+
import org.jetbrains.kotlin.fir.types.builder.buildResolvedTypeRef
8+
import org.jetbrains.kotlin.name.Name
9+
import org.jetbrains.kotlin.types.ConstantValueKind
10+
import org.jetbrains.kotlinx.dataframe.codeGen.ValidFieldName
11+
import org.jetbrains.kotlinx.dataframe.plugin.utils.Names
12+
13+
data class PropertyName(val identifier: Name, val columnNameAnnotation: FirAnnotation?) {
14+
companion object {
15+
fun of(name: String): PropertyName {
16+
val valid = ValidFieldName.of(name)
17+
var columnName = false
18+
val identifier = if (valid.unquoted != name) {
19+
columnName = true
20+
Name.identifier(valid.unquoted)
21+
} else {
22+
Name.identifier(name)
23+
}
24+
val columnNameAnnotation: FirAnnotation? = if (columnName) {
25+
buildAnnotation(name)
26+
} else {
27+
null
28+
}
29+
return PropertyName(identifier, columnNameAnnotation)
30+
}
31+
32+
fun buildAnnotation(name: String): FirAnnotation {
33+
return org.jetbrains.kotlin.fir.expressions.builder.buildAnnotation {
34+
annotationTypeRef = buildResolvedTypeRef {
35+
type = Names.COLUMN_NAME_ANNOTATION.defaultType(emptyList())
36+
}
37+
argumentMapping = buildAnnotationArgumentMapping {
38+
mapping[Names.COLUMN_NAME_ARGUMENT] = buildLiteralExpression(
39+
source = null,
40+
kind = ConstantValueKind.String,
41+
value = name,
42+
setType = true
43+
)
44+
}
45+
}
46+
}
47+
48+
fun of(identifier: Name, columnNameAnnotation: FirAnnotation?): PropertyName {
49+
return PropertyName(identifier, columnNameAnnotation)
50+
}
51+
}
52+
}

‎plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/extensions/impl/SchemaProperty.kt

+12
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,12 @@
1+
package org.jetbrains.kotlinx.dataframe.plugin.extensions.impl
2+
3+
import org.jetbrains.kotlin.fir.types.ConeKotlinType
4+
import org.jetbrains.kotlin.fir.types.ConeTypeProjection
5+
6+
data class SchemaProperty(
7+
val marker: ConeTypeProjection,
8+
val propertyName: PropertyName,
9+
val dataRowReturnType: ConeKotlinType,
10+
val columnContainerReturnType: ConeKotlinType,
11+
val override: Boolean = false
12+
)

‎plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/utils/firFactories.kt

+8-6
Original file line numberDiff line numberDiff line change
@@ -21,23 +21,25 @@ import org.jetbrains.kotlin.fir.toFirResolvedTypeRef
2121
import org.jetbrains.kotlin.fir.types.FirResolvedTypeRef
2222
import org.jetbrains.kotlin.fir.types.impl.ConeClassLikeTypeImpl
2323
import org.jetbrains.kotlin.name.CallableId
24-
import org.jetbrains.kotlin.name.Name
2524
import org.jetbrains.kotlinx.dataframe.plugin.extensions.DataFramePlugin
25+
import org.jetbrains.kotlinx.dataframe.plugin.extensions.impl.PropertyName
2626

2727
internal fun FirDeclarationGenerationExtension.generateExtensionProperty(
2828
callableId: CallableId,
2929
receiverType: ConeClassLikeTypeImpl,
30-
propertyName: Name,
30+
propertyName: PropertyName,
3131
returnTypeRef: FirResolvedTypeRef,
32-
columnName: String? = null,
3332
symbol: FirClassSymbol<*>? = null,
3433
effectiveVisibility: EffectiveVisibility = EffectiveVisibility.Public
3534
): FirProperty {
3635
val firPropertySymbol = FirPropertySymbol(callableId)
3736
return buildProperty {
37+
propertyName.columnNameAnnotation?.let {
38+
annotations += it
39+
}
3840
moduleData = session.moduleData
3941
resolvePhase = FirResolvePhase.BODY_RESOLVE
40-
origin = FirDeclarationOrigin.Plugin(DataFramePlugin(columnName))
42+
origin = FirDeclarationOrigin.Plugin(DataFramePlugin)
4143
status = FirResolvedDeclarationStatusImpl(
4244
Visibilities.Public,
4345
Modality.FINAL,
@@ -67,7 +69,7 @@ internal fun FirDeclarationGenerationExtension.generateExtensionProperty(
6769
getter = buildPropertyAccessor {
6870
moduleData = session.moduleData
6971
resolvePhase = FirResolvePhase.BODY_RESOLVE
70-
origin = FirDeclarationOrigin.Plugin(DataFramePlugin(columnName))
72+
origin = FirDeclarationOrigin.Plugin(DataFramePlugin)
7173
this.returnTypeRef = returnTypeRef
7274
dispatchReceiverType = receiverType
7375
this.symbol = firPropertyAccessorSymbol
@@ -79,7 +81,7 @@ internal fun FirDeclarationGenerationExtension.generateExtensionProperty(
7981
effectiveVisibility
8082
)
8183
}
82-
name = propertyName
84+
name = propertyName.identifier
8385
this.symbol = firPropertySymbol
8486
isVar = false
8587
isLocal = false

‎plugins/kotlin-dataframe/testData/box/columnName_invalidSymbol.kt

+10
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,10 @@
1+
import org.jetbrains.kotlinx.dataframe.*
2+
import org.jetbrains.kotlinx.dataframe.annotations.*
3+
import org.jetbrains.kotlinx.dataframe.api.*
4+
import org.jetbrains.kotlinx.dataframe.io.*
5+
6+
fun box(): String {
7+
val df = dataFrameOf("a.b")(1)
8+
df.`a b`
9+
return "OK"
10+
}

‎plugins/kotlin-dataframe/testData/box/playground.kt

+9
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,9 @@
1+
import org.jetbrains.kotlinx.dataframe.*
2+
import org.jetbrains.kotlinx.dataframe.annotations.*
3+
import org.jetbrains.kotlinx.dataframe.api.*
4+
import org.jetbrains.kotlinx.dataframe.io.*
5+
6+
fun box(): String {
7+
8+
return "OK"
9+
}

‎plugins/kotlin-dataframe/tests-gen/org/jetbrains/kotlin/fir/dataframe/DataFrameBlackBoxCodegenTestGenerated.java

+12
Original file line numberDiff line numberDiff line change
@@ -52,6 +52,12 @@ public void testColumnName() {
5252
runTest("testData/box/columnName.kt");
5353
}
5454

55+
@Test
56+
@TestMetadata("columnName_invalidSymbol.kt")
57+
public void testColumnName_invalidSymbol() {
58+
runTest("testData/box/columnName_invalidSymbol.kt");
59+
}
60+
5561
@Test
5662
@TestMetadata("columnWithStarProjection.kt")
5763
public void testColumnWithStarProjection() {
@@ -304,6 +310,12 @@ public void testPlatformType() {
304310
runTest("testData/box/platformType.kt");
305311
}
306312

313+
@Test
314+
@TestMetadata("playground.kt")
315+
public void testPlayground() {
316+
runTest("testData/box/playground.kt");
317+
}
318+
307319
@Test
308320
@TestMetadata("propertiesOrder.kt")
309321
public void testPropertiesOrder() {

0 commit comments

Comments
 (0)