init research

This commit is contained in:
2026-02-08 11:20:43 -10:00
commit bdf064f54d
3041 changed files with 1592200 additions and 0 deletions
+32
View File
@@ -0,0 +1,32 @@
## :plugins
This folder holds all our Gradle- and Kotlin Compiler plugins:
### [~~:plugins:dataframe-gradle-plugin~~](./dataframe-gradle-plugin)
The Gradle plugin for DataFrame that can generate data schemas from a data sample using the simple
task `dataframes {}`. It uses [:plugins:symbol-processor](./symbol-processor) to generate column accessors.
NOTE: This plugin is disabled as KSP1 is no longer compatible with Kotlin 2.3+.
### [~~:plugins:symbol-processor~~](./symbol-processor)
The KSP plugin that can generate data schemas from a data sample using the `@file:ImportDataSchema` annotation.
It is also used to generate column accessors for in the form of extension properties for
(both manually written- or generated) `@DataSchema` annotated classes/interfaces.
NOTE: This plugin is disabled as KSP1 is no longer compatible with Kotlin 2.3+.
### [~~:plugins:kotlin-dataframe~~](./kotlin-dataframe)
The Kotlin 2.x Compiler plugin of DataFrame.
A plugin for your Kotlin project that can generate on-the-fly column accessors for the compiler and IDE even without
having to provide data schemas!
NOTE: Development of this module was moved to the Kotlin repository:
https://github.com/JetBrains/kotlin/tree/master/plugins/kotlin-dataframe
### [:plugins:expressions-converter](./expressions-converter)
A small Kotlin Compiler plugin that provides intermediate expressions of DataFrame
operation chains, used internally by [:core](../core) to generate "explainer dataframes" on the documentation website.
### [:plugins:keywords-generator](./keywords-generator)
A small Gradle plugin that is used internally to generate enums with restricted Kotlin keywords for the
[:core](../core) module.
+18
View File
@@ -0,0 +1,18 @@
## ~~:plugins:dataframe-gradle-plugin~~
This module holds the Gradle plugin for DataFrame, published as "org.jetbrains.kotlinx.dataframe" on the
[Gradle Plugin Portal](https://plugins.gradle.org/plugin/org.jetbrains.kotlin.plugin.dataframe).
This plugin can let the user generate data schemas from a data sample using the simple
Gradle task `dataframes {}`. It also provides an automatic dependency on [:plugins:symbol-processor](./symbol-processor)
to generate column accessors and support the `@file:ImportDataSchema` notation.
Read more about how to use the Gradle plugin at
[Data Schemas in Gradle projects](https://kotlin.github.io/dataframe/schemasgradle.html).
### DISABLED!
This plugin is disabled as KSP1 is no longer compatible with Kotlin 2.3+.
See https://kotlin.github.io/dataframe/gradle-plugin.html.
The recommended alternative is to use the [Compiler Plugin](https://kotlin.github.io/dataframe/compiler-plugin.html)
and [generating schemas manually from dataframes in runtime](https://kotlin.github.io/dataframe/dataschemagenerationmethods.html).
@@ -0,0 +1,137 @@
plugins {
id("org.gradle.kotlin.kotlin-dsl")
`maven-publish`
with(convention.plugins) {
alias(kotlinJvm11)
alias(buildConfig)
}
with(libs.plugins) {
alias(plugin.publish)
}
}
repositories {
mavenCentral()
mavenLocal()
maven(url = "https://jitpack.io")
google()
}
group = "org.jetbrains.kotlinx.dataframe"
buildscript {
dependencies {
classpath(embeddedKotlin("gradle-plugin"))
}
}
dependencies {
api(libs.kotlin.reflect)
implementation(projects.dataframe)
// experimental
implementation(projects.dataframeOpenapiGenerator)
compileOnly(embeddedKotlin("gradle-plugin"))
implementation(libs.kotlin.gradle.plugin.api)
implementation(libs.serialization.core)
implementation(libs.serialization.json)
implementation(libs.ksp.gradle)
implementation(libs.ksp.api)
testImplementation(gradleTestKit())
testImplementation(embeddedKotlin("test"))
testImplementation(embeddedKotlin("test-junit"))
testImplementation(libs.kotestAssertions)
testImplementation(libs.android.gradle.api)
testImplementation(libs.android.gradle)
testImplementation(embeddedKotlin("gradle-plugin"))
testImplementation(libs.ktor.server.netty)
testImplementation(libs.h2db)
}
tasks.withType<ProcessResources> {
filesMatching("**/plugin.properties") {
filter {
it.replace("%PREPROCESSOR_VERSION%", "$version")
}
}
}
tasks.withType<ProcessResources> {
filesMatching("**/df.properties") {
filter {
it.replace(
"%DATAFRAME_JAR%",
listOf(":core", ":dataframe-csv", ":dataframe-json").joinToString("\", \"") {
project(it).configurations
.getByName("instrumentedJars")
.artifacts.single()
.file.absolutePath
.replace(File.separatorChar, '/')
},
)
}
}
}
gradlePlugin {
// These settings are set for the whole plugin bundle
website = "https://github.com/Kotlin/dataframe"
vcsUrl = "https://github.com/Kotlin/dataframe"
plugins {
create("schemaGeneratorPlugin") {
id = "org.jetbrains.kotlinx.dataframe"
implementationClass = "org.jetbrains.dataframe.gradle.ConvenienceSchemaGeneratorPlugin"
displayName = "Kotlin DataFrame gradle plugin"
description = "Gradle plugin providing task for inferring data schemas from your CSV or JSON data"
tags = listOf("dataframe", "kotlin")
}
}
}
sourceSets {
val main by getting
val test by getting
val testRuntimeClasspath by configurations
create("integrationTest") {
kotlin.srcDir("src/integrationTest/kotlin")
compileClasspath += main.output + test.output + testRuntimeClasspath
runtimeClasspath += output + compileClasspath + test.runtimeClasspath
}
}
val integrationTestConfiguration by configurations.creating {
extendsFrom(configurations.testImplementation.get())
}
tasks.pluginUnderTestMetadata {
pluginClasspath.from(integrationTestConfiguration)
}
val integrationTestTask = tasks.register<Test>("integrationTest") {
dependsOn(":plugins:symbol-processor:publishToMavenLocal")
dependsOn(":dataframe-arrow:publishToMavenLocal")
dependsOn(":dataframe-excel:publishToMavenLocal")
dependsOn(":dataframe-csv:publishToMavenLocal")
dependsOn(":dataframe-jdbc:publishToMavenLocal")
dependsOn(":dataframe-json:publishToMavenLocal")
dependsOn(":dataframe-openapi-generator:publishToMavenLocal")
dependsOn(":dataframe-openapi:publishToMavenLocal")
dependsOn(":publishApiPublicationToMavenLocal")
dependsOn(":dataframe-arrow:publishDataframeArrowPublicationToMavenLocal")
dependsOn(":dataframe-excel:publishDataframeExcelPublicationToMavenLocal")
dependsOn(":dataframe-csv:publishDataframeCsvPublicationToMavenLocal")
dependsOn(":dataframe-jdbc:publishDataframeJDBCPublicationToMavenLocal")
dependsOn(":dataframe-openapi-generator:publishDataframeOpenApiPublicationToMavenLocal")
dependsOn(":plugins:symbol-processor:publishMavenPublicationToMavenLocal")
dependsOn(":core:publishCorePublicationToMavenLocal")
description = "Runs integration tests."
group = "verification"
testClassesDirs = sourceSets["integrationTest"].output.classesDirs
classpath = sourceSets["integrationTest"].runtimeClasspath
shouldRunAfter("test")
}
tasks.check { dependsOn(integrationTestTask) }
@@ -0,0 +1,17 @@
package org.jetbrains.dataframe.gradle
import org.junit.Before
import java.util.Properties
abstract class AbstractDataFramePluginIntegrationTest {
protected val kotlinVersion = TestData.kotlinVersion
protected lateinit var dataframeJarPath: String
@Before
fun before() {
val properties = Properties().also {
it.load(javaClass.getResourceAsStream("df.properties"))
}
dataframeJarPath = properties.getProperty("DATAFRAME_JAR")
}
}
@@ -0,0 +1,86 @@
package org.jetbrains.dataframe.gradle
import io.kotest.matchers.shouldBe
import org.gradle.testkit.runner.TaskOutcome
import org.jetbrains.kotlinx.dataframe.DataFrame
import org.junit.Test
import java.io.File
class ApiChangesDetectionTest : AbstractDataFramePluginIntegrationTest() {
// GenerateDataSchemaTask::class,
// DefaultReadCsvMethod::class,
// DefaultReadJsonMethod::class
@Test
fun `cast api`() {
compiles {
"""
import ${DataFrame::class.qualifiedName!!}
import org.jetbrains.kotlinx.dataframe.api.cast
interface Marker
fun DataFrame<*>.resolveApi() {
cast<Marker>()
cast<Marker>(true)
}
""".trimIndent()
}
}
// GenerateDataSchemaTask::class,
// DefaultReadJsonMethod::class
@Test
fun `read json api`() {
compiles {
"""
import ${DataFrame::class.qualifiedName!!}
import org.jetbrains.kotlinx.dataframe.io.readJson
fun DataFrame<*>.resolveApi(s: String) {
DataFrame.readJson(s)
}
""".trimIndent()
}
}
// GenerateDataSchemaTask::class,
// DefaultReadCsvMethod::class,
@Test
fun `read csv api`() {
compiles {
"""
import ${DataFrame::class.qualifiedName!!}
import org.jetbrains.kotlinx.dataframe.io.readCSV
fun DataFrame<*>.resolveApi(s: String, ch: Char) {
DataFrame.readCSV(s, ch)
}
""".trimIndent()
}
}
private fun compiles(code: () -> String) {
val (_, result) = runGradleBuild(":build") { buildDir ->
val kotlin = File(buildDir, "src/main/kotlin").also { it.mkdirs() }
val main = File(kotlin, "Main.kt")
main.writeText(code())
"""
plugins {
kotlin("jvm") version "$kotlinVersion"
id("org.jetbrains.kotlinx.dataframe")
}
dependencies {
implementation(files("$dataframeJarPath"))
}
repositories {
mavenCentral()
mavenLocal()
}
""".trimIndent()
}
result.task(":build")?.outcome shouldBe TaskOutcome.SUCCESS
}
}
@@ -0,0 +1,685 @@
package org.jetbrains.dataframe.gradle
import io.kotest.assertions.asClue
import io.kotest.matchers.shouldBe
import org.gradle.testkit.runner.TaskOutcome
import org.junit.Ignore
import org.junit.Test
import java.io.File
import java.sql.Connection
import java.sql.DriverManager
class SchemaGeneratorPluginIntegrationTest : AbstractDataFramePluginIntegrationTest() {
private companion object {
private const val FIRST_NAME = "first.csv"
private const val SECOND_NAME = "second.csv"
}
@Test
fun `compileKotlin depends on generateAll task`() {
val (_, result) = runGradleBuild(":compileKotlin") { buildDir ->
File(buildDir, FIRST_NAME).also { it.writeText(TestData.csvSample) }
File(buildDir, SECOND_NAME).also { it.writeText(TestData.csvSample) }
"""
import org.jetbrains.dataframe.gradle.SchemaGeneratorExtension
plugins {
kotlin("jvm") version "$kotlinVersion"
id("org.jetbrains.kotlinx.dataframe")
}
repositories {
mavenCentral()
mavenLocal()
}
dependencies {
implementation(files("$dataframeJarPath"))
}
dataframes {
schema {
data = file("$FIRST_NAME")
name = "Test"
packageName = "org.test"
}
schema {
data = file("$SECOND_NAME")
name = "Schema"
packageName = "org.test"
}
}
""".trimIndent()
}
result.task(":generateDataFrameTest")?.outcome shouldBe TaskOutcome.SUCCESS
result.task(":generateDataFrameSchema")?.outcome shouldBe TaskOutcome.SUCCESS
}
@Test
fun `packageName convention is 'dataframe'`() {
val (dir, result) = runGradleBuild(":build") { buildDir ->
val dataFile = File(buildDir, TestData.csvName)
dataFile.writeText(TestData.csvSample)
"""
import org.jetbrains.dataframe.gradle.SchemaGeneratorExtension
plugins {
kotlin("jvm") version "$kotlinVersion"
id("org.jetbrains.kotlinx.dataframe")
}
repositories {
mavenCentral()
mavenLocal()
}
dependencies {
implementation(files("$dataframeJarPath"))
}
dataframes {
schema {
data = file("${TestData.csvName}")
name = "Data"
}
}
""".trimIndent()
}
result.task(":generateDataFrameData")?.outcome shouldBe TaskOutcome.SUCCESS
File(dir, "build/generated/dataframe").walkBottomUp().toList().asClue {
File(dir, "build/generated/dataframe/main/kotlin/dataframe/Data.Generated.kt").exists() shouldBe true
}
}
@Test
fun `fallback all properties to conventions`() {
val (_, result) = runGradleBuild(":build") { buildDir ->
val dataFile = File(buildDir, TestData.csvName)
dataFile.writeText(TestData.csvSample)
"""
import org.jetbrains.dataframe.gradle.SchemaGeneratorExtension
plugins {
kotlin("jvm") version "$kotlinVersion"
id("org.jetbrains.kotlinx.dataframe")
}
repositories {
mavenCentral()
mavenLocal()
}
dependencies {
implementation(files("$dataframeJarPath"))
}
dataframes {
schema {
data = file("${TestData.csvName}")
}
}
""".trimIndent()
}
result.task(":generateDataFrameData")?.outcome shouldBe TaskOutcome.SUCCESS
}
@Test
fun `generated schemas resolved`() {
val (_, result) = runGradleBuild(":build") { buildDir ->
val dataFile = File(buildDir, "data.csv")
dataFile.writeText(TestData.csvSample)
val kotlin = File(buildDir, "src/main/kotlin").also { it.mkdirs() }
val main = File(kotlin, "Main.kt")
main.writeText(
"""
import org.example.Schema
""".trimIndent(),
)
"""
import org.jetbrains.dataframe.gradle.SchemaGeneratorExtension
plugins {
kotlin("jvm") version "$kotlinVersion"
id("org.jetbrains.kotlinx.dataframe")
}
repositories {
mavenLocal()
mavenCentral()
maven(url="https://jitpack.io")
}
dependencies {
implementation(files("$dataframeJarPath"))
}
kotlin.sourceSets.getByName("main").kotlin.srcDir("build/generated/ksp/main/kotlin/")
dataframes {
schema {
data = file("${TestData.csvName}")
name = "org.example.Schema"
}
}
""".trimIndent()
}
result.task(":build")?.outcome shouldBe TaskOutcome.SUCCESS
}
@Test
fun `generated schemas resolved in jvmMain source set for multiplatform project`() {
val (_, result) = runGradleBuild(":build") { buildDir ->
val dataFile = File(buildDir, "data.csv")
dataFile.writeText(TestData.csvSample)
val kotlin = File(buildDir, "src/jvmMain/kotlin").also { it.mkdirs() }
val main = File(kotlin, "Main.kt")
main.writeText(
"""
import org.example.Schema
""".trimIndent(),
)
"""
import org.jetbrains.dataframe.gradle.SchemaGeneratorExtension
plugins {
kotlin("multiplatform") version "$kotlinVersion"
id("org.jetbrains.kotlinx.dataframe")
}
repositories {
mavenCentral()
mavenLocal()
}
kotlin {
jvm()
sourceSets {
val jvmMain by getting {
dependencies {
implementation(files("$dataframeJarPath"))
}
}
}
}
dataframes {
schema {
data = file("${TestData.csvName}")
name = "org.example.Schema"
}
}
""".trimIndent()
}
result.task(":build")?.outcome shouldBe TaskOutcome.SUCCESS
}
@Ignore
@Test
fun `kotlin identifiers generated from csv names`() {
fun escapeDoubleQuotes(it: Char) = if (it == '"') "\"\"" else it.toString()
val (_, result) = runGradleBuild(":build") { buildDir ->
val filename = "data.csv"
val dataFile = File(buildDir, filename)
val notSupportedChars = setOf('\n', '\r')
(Char.MIN_VALUE..Char.MAX_VALUE).asSequence()
.filterNot { it in notSupportedChars }
.chunked(100) {
it.joinToString(
separator = "",
prefix = "\"",
postfix = "\"",
transform = ::escapeDoubleQuotes,
)
}.let {
dataFile.writeText(it.joinToString(",") + "\n" + (0 until it.count()).joinToString(","))
}
val kotlin = File(buildDir, "src/main/kotlin").also { it.mkdirs() }
val main = File(kotlin, "Main.kt")
main.writeText(
"""
import org.jetbrains.kotlinx.dataframe.DataFrame
import org.jetbrains.kotlinx.dataframe.annotations.DataSchema
import org.jetbrains.kotlinx.dataframe.io.read
import org.jetbrains.kotlinx.dataframe.api.cast
import org.jetbrains.kotlinx.dataframe.api.filter
fun main() {
val df = DataFrame.read("${TestData.csvName}").cast<Schema>()
}
""".trimIndent(),
)
"""
import org.jetbrains.dataframe.gradle.SchemaGeneratorExtension
plugins {
kotlin("jvm") version "$kotlinVersion"
id("org.jetbrains.kotlinx.dataframe")
}
repositories {
mavenLocal()
mavenCentral()
}
dependencies {
implementation(files("$dataframeJarPath"))
}
kotlin.sourceSets.getByName("main").kotlin.srcDir("build/generated/ksp/main/kotlin/")
dataframes {
schema {
data = "${TestData.csvName}"
name = "Schema"
packageName = ""
}
}
""".trimIndent()
}
result.task(":build")?.outcome shouldBe TaskOutcome.SUCCESS
}
@Test
fun `preprocessor generates extensions for DataSchema`() {
val (_, result) = runGradleBuild(":build") { buildDir ->
val dataFile = File(buildDir, "data.csv")
dataFile.writeText(TestData.csvSample)
val kotlin = File(buildDir, "src/main/kotlin").also { it.mkdirs() }
val main = File(kotlin, "Main.kt")
main.writeText(
"""
import org.jetbrains.kotlinx.dataframe.DataFrame
import org.jetbrains.kotlinx.dataframe.annotations.DataSchema
import org.jetbrains.kotlinx.dataframe.io.read
import org.jetbrains.kotlinx.dataframe.api.cast
import org.jetbrains.kotlinx.dataframe.api.filter
@DataSchema
interface MySchema {
val age: Int
}
fun main() {
val df = DataFrame.read("${TestData.csvName}").cast<MySchema>()
val df1 = df.filter { age != null }
}
""".trimIndent(),
)
@Suppress("DuplicatedCode")
"""
import org.jetbrains.dataframe.gradle.SchemaGeneratorExtension
plugins {
kotlin("jvm") version "$kotlinVersion"
id("org.jetbrains.kotlinx.dataframe")
}
repositories {
mavenLocal()
mavenCentral()
}
dependencies {
implementation(files("$dataframeJarPath"))
}
kotlin.sourceSets.getByName("main").kotlin.srcDir("build/generated/ksp/main/kotlin/")
""".trimIndent()
}
result.task(":build")?.outcome shouldBe TaskOutcome.SUCCESS
}
@Test
fun `preprocessor imports schema from local file`() {
val (_, result) = runGradleBuild(":build") { buildDir ->
val dataFile = File(buildDir, "data.csv")
dataFile.writeText(TestData.csvSample)
val kotlin = File(buildDir, "src/main/kotlin").also { it.mkdirs() }
val main = File(kotlin, "Main.kt")
main.writeText(
"""
@file:ImportDataSchema(name = "MySchema", path = "${TestData.csvName}")
package test
import org.jetbrains.kotlinx.dataframe.annotations.ImportDataSchema
import org.jetbrains.kotlinx.dataframe.api.filter
fun main() {
val df = MySchema.readCsv()
val df1 = df.filter { age != null }
}
""".trimIndent(),
)
"""
import org.jetbrains.dataframe.gradle.SchemaGeneratorExtension
plugins {
kotlin("jvm") version "$kotlinVersion"
id("org.jetbrains.kotlinx.dataframe")
}
repositories {
mavenLocal()
mavenCentral()
}
dependencies {
implementation(files("$dataframeJarPath"))
}
kotlin.sourceSets.getByName("main").kotlin.srcDir("build/generated/ksp/main/kotlin/")
""".trimIndent()
}
result.task(":build")?.outcome shouldBe TaskOutcome.SUCCESS
}
/* TODO: test is broken
e: file://test3901867314473689900/src/main/kotlin/Main.kt:12:43 Unresolved reference: readSqlTable
e: file://test3901867314473689900/src/main/kotlin/Main.kt:13:43 Unresolved reference: DbConnectionConfig
e: file://test3901867314473689900/src/main/kotlin/Main.kt:19:28 Unresolved reference: readSqlTable
e: file://test3901867314473689900/src/main/kotlin/Main.kt:20:21 Unresolved reference: age
e: file://test3901867314473689900/src/main/kotlin/Main.kt:22:29 Unresolved reference: readSqlTable
e: file://test3901867314473689900/src/main/kotlin/Main.kt:23:22 Unresolved reference: age
e: file://test3901867314473689900/src/main/kotlin/Main.kt:25:24 Unresolved reference: DbConnectionConfig
e: file://test3901867314473689900/src/main/kotlin/Main.kt:26:29 Unresolved reference: readSqlTable
e: file://test3901867314473689900/src/main/kotlin/Main.kt:27:22 Unresolved reference: age
e: file://test3901867314473689900/src/main/kotlin/Main.kt:29:29 Unresolved reference: readSqlTable
e: file://test3901867314473689900/src/main/kotlin/Main.kt:30:22 Unresolved reference: age
*/
@Test
@Ignore
fun `preprocessor imports schema from database`() {
val connectionUrl = "jdbc:h2:mem:test;DB_CLOSE_DELAY=-1;MODE=MySQL;DATABASE_TO_UPPER=false"
DriverManager.getConnection(connectionUrl).use {
val (_, result) = runGradleBuild(":build") { buildDir ->
createTestDatabase(it)
val kotlin = File(buildDir, "src/main/kotlin").also { it.mkdirs() }
val main = File(kotlin, "Main.kt")
// this is a copy of the code snippet in the
// DataFrameJdbcSymbolProcessorTest.`schema extracted via readFromDB method is resolved`
main.writeText(
"""
@file:ImportDataSchema(name = "Customer", path = "$connectionUrl")
package test
import org.jetbrains.kotlinx.dataframe.annotations.ImportDataSchema
import org.jetbrains.kotlinx.dataframe.api.filter
import org.jetbrains.kotlinx.dataframe.DataFrame
import org.jetbrains.kotlinx.dataframe.api.cast
import java.sql.Connection
import java.sql.DriverManager
import java.sql.SQLException
import org.jetbrains.kotlinx.dataframe.io.readSqlTable
import org.jetbrains.kotlinx.dataframe.io.DbConnectionConfig
fun main() {
Class.forName("org.h2.Driver")
val tableName = "Customer"
DriverManager.getConnection("$connectionUrl").use { connection ->
val df = DataFrame.readSqlTable(connection, tableName).cast<Customer>()
df.filter { age != null && age > 30 }
val df1 = DataFrame.readSqlTable(connection, tableName, 1).cast<Customer>()
df1.filter { age != null && age > 30 }
val dbConfig = DbConnectionConfig(url = "$connectionUrl")
val df2 = DataFrame.readSqlTable(dbConfig, tableName).cast<Customer>()
df2.filter { age != null && age > 30 }
val df3 = DataFrame.readSqlTable(dbConfig, tableName, 1).cast<Customer>()
df3.filter { age != null && age > 30 }
}
}
""".trimIndent(),
)
"""
import org.jetbrains.dataframe.gradle.SchemaGeneratorExtension
plugins {
kotlin("jvm") version "$kotlinVersion"
id("org.jetbrains.kotlinx.dataframe")
}
repositories {
mavenLocal()
mavenCentral()
}
dependencies {
implementation(files("$dataframeJarPath"))
}
kotlin.sourceSets.getByName("main").kotlin.srcDir("build/generated/ksp/main/kotlin/")
""".trimIndent()
}
result.task(":build")?.outcome shouldBe TaskOutcome.SUCCESS
}
}
private fun createTestDatabase(connection: Connection) {
// Crate table Customer
connection.createStatement().execute(
"""
CREATE TABLE Customer (
id INT PRIMARY KEY,
name VARCHAR(50),
age INT
)
""".trimIndent(),
)
// Create table Sale
connection.createStatement().execute(
"""
CREATE TABLE Sale (
id INT PRIMARY KEY,
customerId INT,
amount DECIMAL(10, 2)
)
""".trimIndent(),
)
// add data to the Customer table
connection.createStatement().execute("INSERT INTO Customer (id, name, age) VALUES (1, 'John', 40)")
connection.createStatement().execute("INSERT INTO Customer (id, name, age) VALUES (2, 'Alice', 25)")
connection.createStatement().execute("INSERT INTO Customer (id, name, age) VALUES (3, 'Bob', 47)")
// add data to the Sale table
connection.createStatement().execute("INSERT INTO Sale (id, customerId, amount) VALUES (1, 1, 100.50)")
connection.createStatement().execute("INSERT INTO Sale (id, customerId, amount) VALUES (2, 2, 50.00)")
connection.createStatement().execute("INSERT INTO Sale (id, customerId, amount) VALUES (3, 1, 75.25)")
connection.createStatement().execute("INSERT INTO Sale (id, customerId, amount) VALUES (4, 3, 35.15)")
}
@Test
fun `generated code compiles in explicit api mode`() {
val (_, result) = runGradleBuild(":build") { buildDir ->
val dataFile = File(buildDir, TestData.csvName)
dataFile.writeText(TestData.csvSample)
"""
import org.jetbrains.dataframe.gradle.SchemaGeneratorExtension
plugins {
kotlin("jvm") version "$kotlinVersion"
id("org.jetbrains.kotlinx.dataframe")
}
repositories {
mavenCentral()
mavenLocal()
}
dependencies {
implementation(files("$dataframeJarPath"))
}
kotlin {
explicitApi()
}
dataframes {
schema {
data = file("${TestData.csvName}")
}
}
""".trimIndent()
}
result.task(":generateDataFrameData")?.outcome shouldBe TaskOutcome.SUCCESS
}
@Ignore
@Test
fun `plugin doesn't break multiplatform build without JVM`() {
val (_, result) = runGradleBuild(":build") { buildDir ->
val dataFile = File(buildDir, TestData.csvName)
val kotlin = File(buildDir, "src/jsMain/kotlin").also { it.mkdirs() }
val main = File(kotlin, "Main.kt")
main.writeText(
"""
fun main() {
console.log("Hello, Kotlin/JS!")
}
""".trimIndent(),
)
dataFile.writeText(TestData.csvSample)
"""
import org.jetbrains.dataframe.gradle.SchemaGeneratorExtension
plugins {
kotlin("multiplatform") version "$kotlinVersion"
id("org.jetbrains.kotlinx.dataframe")
}
repositories {
mavenCentral()
mavenLocal()
}
kotlin {
sourceSets {
js(IR) {
browser()
}
}
}
dataframes {
schema {
data = file("${TestData.csvName}")
name = "Schema"
packageName = ""
src = buildDir
}
}
""".trimIndent()
}
result.task(":build")?.outcome shouldBe TaskOutcome.SUCCESS
}
@Test
fun `plugin doesn't break multiplatform build with JVM`() {
val (_, result) = runGradleBuild(":build") { buildDir ->
val dataFile = File(buildDir, TestData.csvName)
val kotlin = File(buildDir, "src/jvmMain/kotlin").also { it.mkdirs() }
val main = File(kotlin, "Main.kt")
main.writeText(
"""
fun main() {
println("Hello, Kotlin/JVM!")
}
""".trimIndent(),
)
dataFile.writeText(TestData.csvSample)
"""
import org.jetbrains.dataframe.gradle.SchemaGeneratorExtension
plugins {
kotlin("multiplatform") version "$kotlinVersion"
id("org.jetbrains.kotlinx.dataframe")
}
repositories {
mavenCentral()
mavenLocal()
}
kotlin {
sourceSets {
jvm()
}
}
dataframes {
schema {
data = file("${TestData.csvName}")
name = "Schema"
packageName = ""
src = buildDir
}
}
""".trimIndent()
}
result.task(":build")?.outcome shouldBe TaskOutcome.SUCCESS
}
@Test
fun `companion object for csv compiles`() {
testCompanionObject(TestData.csvName, TestData.csvSample)
}
@Test
fun `companion object for json compiles`() {
testCompanionObject(TestData.jsonName, TestData.jsonSample)
}
private fun testCompanionObject(dataName: String, dataSample: String) {
val (_, result) = runGradleBuild(":build") { buildDir ->
val dataFile = File(buildDir, dataName)
dataFile.writeText(dataSample)
"""
import org.jetbrains.dataframe.gradle.SchemaGeneratorExtension
plugins {
kotlin("jvm") version "$kotlinVersion"
id("org.jetbrains.kotlinx.dataframe")
}
repositories {
mavenCentral()
mavenLocal()
}
dependencies {
implementation(files("$dataframeJarPath"))
}
dataframes {
schema {
data = file("$dataName")
name = "Schema"
}
}
""".trimIndent()
}
result.task(":build")?.outcome shouldBe TaskOutcome.SUCCESS
}
}
@@ -0,0 +1 @@
DATAFRAME_JAR=%DATAFRAME_JAR%
@@ -0,0 +1,131 @@
package org.jetbrains.dataframe.gradle
import com.google.devtools.ksp.gradle.KspExtension
import org.gradle.api.Plugin
import org.gradle.api.Project
import org.gradle.kotlin.dsl.getByType
import java.util.Properties
@Suppress("unused")
public class ConvenienceSchemaGeneratorPlugin : Plugin<Project> {
public companion object {
/**
* (boolean, default `true`) whether to add KSP plugin
*/
public const val PROP_ADD_KSP: String = "kotlin.dataframe.add.ksp"
/**
* (string, default `null`) comma-delimited list of configurations to add KSP processing to.
* Defaults to guessing configurations based on which kotlin plugin is applied (jvm or multiplatform)
*/
public const val PROP_KSP_CONFIGS: String = "kotlin.dataframe.ksp.configs"
}
override fun apply(target: Project) {
val property = target.findProperty(PROP_ADD_KSP)?.toString()
var addKsp = true
if (property != null) {
if (property.equals("true", ignoreCase = true) || property.equals("false", ignoreCase = true)) {
addKsp = property.toBoolean()
} else {
target.logger.warn(
"Invalid value '$property' for '$PROP_ADD_KSP' property. Defaulting to '$addKsp'. Please use 'true' or 'false'.",
)
}
}
val properties = Properties()
properties.load(javaClass.getResourceAsStream("plugin.properties"))
val preprocessorVersion = properties.getProperty("PREPROCESSOR_VERSION")
// regardless whether we add KSP or the user adds it, when it's added,
// configure it to depend on symbol-processor-all
target.plugins.whenPluginAdded {
if (this::class.qualifiedName?.contains("com.google.devtools.ksp") == true) {
val isMultiplatform by lazy {
when {
target.plugins.hasPlugin("org.jetbrains.kotlin.jvm") -> false
target.plugins.hasPlugin("org.jetbrains.kotlin.multiplatform") -> true
else -> {
target.logger.warn(
"Kotlin plugin must be applied first so we know whether to use multiplatform configurations or not",
)
false
}
}
}
val overriddenConfigs = target.findProperty(PROP_KSP_CONFIGS)
?.let { (it as String) }
?.split(",")
?.map { it.trim() }
val configs = when {
overriddenConfigs != null -> overriddenConfigs
isMultiplatform -> listOf("kspJvm", "kspJvmTest")
else -> listOf("ksp", "kspTest")
}
val cfgsToAdd = configs.toMutableSet()
configs.forEach { cfg ->
target.configurations.findByName(cfg)?.apply {
cfgsToAdd.remove(cfg)
dependencies.add(
target.dependencies.create(
"org.jetbrains.kotlinx.dataframe:symbol-processor-all:$preprocessorVersion",
),
)
}
}
target.gradle.projectsEvaluated {
cfgsToAdd.forEach { cfg ->
target.logger.warn(
"Configuration '$cfg' was never found. Please make sure the KSP plugin is applied.",
)
}
}
target.logger.info("Added DataFrame dependency to the KSP plugin.")
target.extensions.getByType<KspExtension>().arg(
"dataframe.resolutionDir",
target.projectDir.absolutePath,
)
}
}
if (addKsp) {
target.plugins.apply(KspPluginApplier::class.java)
} else {
target.logger.warn(
"Plugin 'org.jetbrains.kotlinx.dataframe' comes bundled with its own version of KSP which is " +
"currently disabled as 'kotlin.dataframe.add.ksp' is set to 'false' in a 'properties' file. " +
"Either set 'kotlin.dataframe.add.ksp' to 'true' or add the plugin 'com.google.devtools.ksp' " +
"manually.",
)
}
target.plugins.apply(SchemaGeneratorPlugin::class.java)
}
}
@Suppress("unused")
public class DeprecatingSchemaGeneratorPlugin : Plugin<Project> {
override fun apply(target: Project) {
target.logger.warn(
"DEPRECATION: Replace plugin id(\"org.jetbrains.kotlin.plugin.dataframe\") and kotlin(\"plugin.dataframe\") with id(\"org.jetbrains.kotlinx.dataframe\").",
)
target.plugins.apply(ConvenienceSchemaGeneratorPlugin::class.java)
}
}
/**
* Applies the KSP plugin in the target project.
*/
internal class KspPluginApplier : Plugin<Project> {
override fun apply(target: Project) {
val properties = Properties()
properties.load(javaClass.getResourceAsStream("plugin.properties"))
target.plugins.apply("com.google.devtools.ksp")
}
}
@@ -0,0 +1,7 @@
package org.jetbrains.dataframe.gradle
public enum class DataSchemaVisibility {
INTERNAL,
IMPLICIT_PUBLIC,
EXPLICIT_PUBLIC,
}
@@ -0,0 +1,21 @@
package org.jetbrains.dataframe.gradle
import java.io.File
import java.net.MalformedURLException
import java.net.URL
internal fun extractFileName(url: URL): String? =
url.path.takeIf { it.isNotEmpty() }
?.substringAfterLast("/")
?.substringBeforeLast(".")
internal fun extractFileName(file: File): String = file.nameWithoutExtension
internal fun extractFileName(path: String): String? =
try {
val url = URL(path)
extractFileName(url)
} catch (e: MalformedURLException) {
val file = File(path)
extractFileName(file)
}
@@ -0,0 +1,10 @@
package org.jetbrains.dataframe.gradle
import java.io.File
internal fun File.isMiddlePackage(): Boolean = isDirectory && (listFiles()?.singleOrNull()?.isDirectory ?: false)
internal fun File.findDeepestCommonSubdirectory(): File {
if (!exists()) return this
return walkTopDown().filterNot { it.isMiddlePackage() }.first()
}
@@ -0,0 +1,278 @@
package org.jetbrains.dataframe.gradle
import org.gradle.api.DefaultTask
import org.gradle.api.provider.Property
import org.gradle.api.provider.Provider
import org.gradle.api.provider.SetProperty
import org.gradle.api.tasks.Input
import org.gradle.api.tasks.OutputFile
import org.gradle.api.tasks.TaskAction
import org.jetbrains.kotlinx.dataframe.codeGen.CodeGenerator
import org.jetbrains.kotlinx.dataframe.codeGen.MarkerVisibility
import org.jetbrains.kotlinx.dataframe.codeGen.NameNormalizer
import org.jetbrains.kotlinx.dataframe.impl.codeGen.CodeGenerationReadResult
import org.jetbrains.kotlinx.dataframe.impl.codeGen.DfReadResult
import org.jetbrains.kotlinx.dataframe.impl.codeGen.from
import org.jetbrains.kotlinx.dataframe.impl.codeGen.toStandaloneSnippet
import org.jetbrains.kotlinx.dataframe.impl.codeGen.urlCodeGenReader
import org.jetbrains.kotlinx.dataframe.impl.codeGen.urlDfReader
import org.jetbrains.kotlinx.dataframe.io.ArrowFeather
import org.jetbrains.kotlinx.dataframe.io.CsvDeephaven
import org.jetbrains.kotlinx.dataframe.io.Excel
import org.jetbrains.kotlinx.dataframe.io.JSON
import org.jetbrains.kotlinx.dataframe.io.OpenApi
import org.jetbrains.kotlinx.dataframe.io.TsvDeephaven
import org.jetbrains.kotlinx.dataframe.io.isUrl
import org.jetbrains.kotlinx.dataframe.io.readSqlQuery
import org.jetbrains.kotlinx.dataframe.io.readSqlTable
import org.jetbrains.kotlinx.dataframe.schema.DataFrameSchema
import java.io.File
import java.net.URL
import java.nio.file.Paths
import java.sql.Connection
import java.sql.DriverManager
public abstract class GenerateDataSchemaTask : DefaultTask() {
@get:Input
public abstract val data: Property<Any>
@get:Input
public abstract val csvOptions: Property<CsvOptionsDsl>
@get:Input
public abstract val jsonOptions: Property<JsonOptionsDsl>
@get:Input
public abstract val jdbcOptions: Property<JdbcOptionsDsl>
@get:Input
public abstract val src: Property<File>
@get:Input
public abstract val interfaceName: Property<String>
@get:Input
public abstract val packageName: Property<String>
@get:Input
public abstract val schemaVisibility: Property<DataSchemaVisibility>
@get:Input
public abstract val defaultPath: Property<Boolean>
@get:Input
public abstract val delimiters: SetProperty<Char>
@get:Input
public abstract val enableExperimentalOpenApi: Property<Boolean>
@Suppress("LeakingThis")
@get:OutputFile
public val dataSchema: Provider<File> = packageName.zip(interfaceName) { packageName, interfaceName ->
val packagePath = packageName.replace('.', File.separatorChar)
Paths.get(src.get().absolutePath, packagePath, "$interfaceName.Generated.kt").toFile()
}
@TaskAction
public fun generate() {
val csvOptions = csvOptions.get()
val jsonOptions = jsonOptions.get()
val jdbcOptions = jdbcOptions.get()
val schemaFile = dataSchema.get()
val escapedPackageName = escapePackageName(packageName.get())
val rawUrl = data.get().toString()
// revisit architecture for an addition of the new data source https://github.com/Kotlin/dataframe/issues/450
if (rawUrl.startsWith("jdbc")) {
val connection = DriverManager.getConnection(rawUrl, jdbcOptions.user, jdbcOptions.password)
connection.use {
val schema = generateSchemaByJdbcOptions(jdbcOptions, connection)
val codeGenerator = CodeGenerator.create(useFqNames = false)
val additionalImports: List<String> = listOf()
val delimiters = delimiters.get()
val codeGenResult = codeGenerator.generate(
schema = schema,
name = interfaceName.get(),
fields = true,
extensionProperties = false,
isOpen = true,
visibility = when (schemaVisibility.get()) {
DataSchemaVisibility.INTERNAL -> MarkerVisibility.INTERNAL
DataSchemaVisibility.IMPLICIT_PUBLIC -> MarkerVisibility.IMPLICIT_PUBLIC
DataSchemaVisibility.EXPLICIT_PUBLIC -> MarkerVisibility.EXPLICIT_PUBLIC
},
readDfMethod = null,
fieldNameNormalizer = NameNormalizer.from(delimiters),
)
schemaFile.writeText(codeGenResult.toStandaloneSnippet(escapedPackageName, additionalImports))
return
}
} else {
val url = urlOf(data.get())
val formats = listOfNotNull(
CsvDeephaven(delimiter = csvOptions.delimiter),
JSON(
typeClashTactic = jsonOptions.typeClashTactic,
keyValuePaths = jsonOptions.keyValuePaths,
unifyNumbers = jsonOptions.unifyNumbers,
),
Excel(),
TsvDeephaven(),
ArrowFeather(),
if (enableExperimentalOpenApi.get()) OpenApi() else null,
)
// first try without creating dataframe
when (val codeGenResult = CodeGenerator.urlCodeGenReader(url, interfaceName.get(), formats, false)) {
is CodeGenerationReadResult.Success -> {
val readDfMethod = codeGenResult.getReadDfMethod(stringOf(data.get()))
val code = codeGenResult
.code
.toStandaloneSnippet(escapedPackageName, readDfMethod.additionalImports)
schemaFile.bufferedWriter().use {
it.write(code)
}
return
}
is CodeGenerationReadResult.Error ->
logger.warn("Error while reading types-only from data at $url: ${codeGenResult.reason}")
}
// on error, try with reading dataframe first
val parsedDf = when (val readResult = CodeGenerator.urlDfReader(url, formats)) {
is DfReadResult.Error -> throw Exception(
"Error while reading dataframe from data at $url",
readResult.reason,
)
is DfReadResult.Success -> readResult
}
val codeGenerator = CodeGenerator.create(useFqNames = false)
val delimiters = delimiters.get()
val readDfMethod = parsedDf.getReadDfMethod(stringOf(data.get()))
val codeGenResult = codeGenerator.generate(
schema = parsedDf.schema,
name = interfaceName.get(),
fields = true,
extensionProperties = false,
isOpen = true,
visibility = when (schemaVisibility.get()) {
DataSchemaVisibility.INTERNAL -> MarkerVisibility.INTERNAL
DataSchemaVisibility.IMPLICIT_PUBLIC -> MarkerVisibility.IMPLICIT_PUBLIC
DataSchemaVisibility.EXPLICIT_PUBLIC -> MarkerVisibility.EXPLICIT_PUBLIC
},
readDfMethod = readDfMethod,
fieldNameNormalizer = NameNormalizer.from(delimiters),
)
schemaFile.writeText(codeGenResult.toStandaloneSnippet(escapedPackageName, readDfMethod.additionalImports))
}
}
// TODO: copy pasted from symbol-processor: DataSchemaGenerator, should be refactored somehow
private fun generateSchemaByJdbcOptions(jdbcOptions: JdbcOptionsDsl, connection: Connection): DataFrameSchema {
logger.debug("Table name: ${jdbcOptions.tableName}")
logger.debug("SQL query: ${jdbcOptions.sqlQuery}")
val tableName = jdbcOptions.tableName
val sqlQuery = jdbcOptions.sqlQuery
return when {
isTableNameNotBlankAndQueryBlank(tableName, sqlQuery) -> generateSchemaForTable(connection, tableName)
isQueryNotBlankAndTableBlank(tableName, sqlQuery) -> generateSchemaForQuery(connection, sqlQuery)
areBothNotBlank(tableName, sqlQuery) -> throwBothFieldsFilledException(tableName, sqlQuery)
else -> throwBothFieldsEmptyException(tableName, sqlQuery)
}
}
private fun isTableNameNotBlankAndQueryBlank(tableName: String, sqlQuery: String) =
tableName.isNotBlank() && sqlQuery.isBlank()
private fun isQueryNotBlankAndTableBlank(tableName: String, sqlQuery: String) =
sqlQuery.isNotBlank() && tableName.isBlank()
private fun areBothNotBlank(tableName: String, sqlQuery: String) = sqlQuery.isNotBlank() && tableName.isNotBlank()
private fun generateSchemaForTable(connection: Connection, tableName: String) =
DataFrameSchema.readSqlTable(connection, tableName)
private fun generateSchemaForQuery(connection: Connection, sqlQuery: String) =
DataFrameSchema.readSqlQuery(connection, sqlQuery)
private fun throwBothFieldsFilledException(tableName: String, sqlQuery: String): Nothing =
throw RuntimeException(
"Table name '$tableName' and SQL query '$sqlQuery' both are filled! " +
"Clear 'tableName' or 'sqlQuery' properties in jdbcOptions with value to generate schema for SQL table or result of SQL query!",
)
private fun throwBothFieldsEmptyException(tableName: String, sqlQuery: String): Nothing =
throw RuntimeException(
"Table name '$tableName' and SQL query '$sqlQuery' both are empty! " +
"Populate 'tableName' or 'sqlQuery' properties in jdbcOptions with value to generate schema for SQL table or result of SQL query!",
)
private fun stringOf(data: Any): String =
when (data) {
is File -> data.absolutePath
is URL -> data.toExternalForm()
is String ->
when {
isUrl(data) -> stringOf(URL(data))
else -> {
val relativeFile = project.file(data)
val absoluteFile = File(data)
stringOf(if (relativeFile.exists()) relativeFile else absoluteFile)
}
}
else -> unsupportedType()
}
// See RegexExpectationsTest
private fun escapePackageName(packageName: String): String =
if (packageName.isNotEmpty()) {
packageName.split(NameChecker.PACKAGE_IDENTIFIER_DELIMITER)
.joinToString(".") { part -> "`$part`" }
} else {
packageName
}
private fun urlOf(data: Any): URL =
when (data) {
is File -> data.toURI()
is URL -> data.toURI()
is String -> when {
isUrl(data) -> URL(data).toURI()
else -> {
val relativeFile = project.file(data)
val absoluteFile = File(data)
if (relativeFile.exists()) {
relativeFile
} else {
absoluteFile
}.toURI()
}
}
else -> unsupportedType()
}.toURL()
private fun unsupportedType(): Nothing =
throw IllegalArgumentException("data for schema \"${interfaceName.get()}\" must be File, URL or String")
}
@@ -0,0 +1,28 @@
package org.jetbrains.dataframe.gradle
internal object NameChecker {
fun checkValidIdentifier(identifiers: String) {
check(identifiers.none { char -> char in NAME_RESTRICTED_CHARS }) {
val illegalChars = NAME_RESTRICTED_CHARS.intersect(identifiers.toSet()).joinToString(",")
"$identifiers contains illegal characters: $illegalChars"
}
}
fun checkValidPackageName(name: String) {
val identifiers = name
.split(PACKAGE_IDENTIFIER_DELIMITER)
.joinToString("")
checkValidIdentifier(identifiers)
}
// https://github.com/JetBrains/kotlin/blob/1.5.30/compiler/frontend.java/src/org/jetbrains/kotlin/resolve/jvm/checkers/JvmSimpleNameBacktickChecker.kt
private val INVALID_CHARS = setOf('.', ';', '[', ']', '/', '<', '>', ':', '\\')
private val DANGEROUS_CHARS = setOf('?', '*', '"', '|', '%')
// QuotedSymbol https://kotlinlang.org/spec/syntax-and-grammar.html#identifiers
private val RESTRICTED_CHARS = setOf('`', '\r', '\n')
private val NAME_RESTRICTED_CHARS = INVALID_CHARS + DANGEROUS_CHARS + RESTRICTED_CHARS
// https://kotlinlang.org/spec/syntax-and-grammar.html#grammar-rule-identifier
val PACKAGE_IDENTIFIER_DELIMITER = "(\\n|\\r\\n)*\\.".toRegex()
}
@@ -0,0 +1,155 @@
@file:Suppress("unused")
package org.jetbrains.dataframe.gradle
import groovy.lang.Closure
import org.gradle.api.Project
import org.jetbrains.kotlinx.dataframe.api.JsonPath
import org.jetbrains.kotlinx.dataframe.io.JSON
import java.io.File
import java.io.Serializable
import java.net.URL
public open class SchemaGeneratorExtension {
public lateinit var project: Project
internal val schemas: MutableList<Schema> = mutableListOf()
public var packageName: String? = null
public var sourceSet: String? = null
public var visibility: DataSchemaVisibility? = null
internal var defaultPath: Boolean? = null
internal var withNormalizationBy: Set<Char>? = null
/** Can be set to `true` to enable experimental OpenAPI 3.0.0 types support */
public var enableExperimentalOpenApi: Boolean = false
public fun schema(config: Schema.() -> Unit) {
val schema = Schema(project).apply(config)
schemas.add(schema)
}
public fun schema(config: Closure<*>) {
val schema = Schema(project)
project.configure(schema, config)
schemas.add(schema)
}
public fun withoutDefaultPath() {
defaultPath = false
}
public fun withNormalizationBy(vararg delimiter: Char) {
withNormalizationBy = delimiter.toSet()
}
// Overload for Groovy.
// It's impossible to call a method with char argument without type cast in groovy, because it only has string literals
public fun withNormalizationBy(delimiter: String) {
withNormalizationBy = delimiter.toSet()
}
public fun withoutNormalization() {
withNormalizationBy = emptySet()
}
public fun enableExperimentalOpenApi(enable: Boolean) {
enableExperimentalOpenApi = enable
}
}
public class Schema(
private val project: Project,
public var data: Any? = null,
public var src: File? = null,
public var name: String? = null,
public var packageName: String? = null,
public var sourceSet: String? = null,
public var visibility: DataSchemaVisibility? = null,
internal var defaultPath: Boolean? = null,
internal var withNormalizationBy: Set<Char>? = null,
public val csvOptions: CsvOptionsDsl = CsvOptionsDsl(),
public val jsonOptions: JsonOptionsDsl = JsonOptionsDsl(),
public val jdbcOptions: JdbcOptionsDsl = JdbcOptionsDsl(),
) {
public fun setData(file: File) {
data = file
}
public fun setData(path: String) {
data = path
}
public fun setData(url: URL) {
data = url
}
public fun csvOptions(config: CsvOptionsDsl.() -> Unit) {
csvOptions.apply(config)
}
public fun csvOptions(config: Closure<*>) {
project.configure(csvOptions, config)
}
public fun jsonOptions(config: JsonOptionsDsl.() -> Unit) {
jsonOptions.apply(config)
}
public fun jsonOptions(config: Closure<*>) {
project.configure(jsonOptions, config)
}
public fun jdbcOptions(config: JdbcOptionsDsl.() -> Unit) {
jdbcOptions.apply(config)
}
public fun jdbcOptions(config: Closure<*>) {
project.configure(jdbcOptions, config)
}
public fun withoutDefaultPath() {
defaultPath = false
}
public fun withDefaultPath() {
defaultPath = true
}
public fun withNormalizationBy(vararg delimiter: Char) {
withNormalizationBy = delimiter.toSet()
}
// Overload for Groovy.
// It's impossible to call a method with char argument without type cast in groovy, because it only has string literals
public fun withNormalizationBy(delimiter: String) {
withNormalizationBy = delimiter.toSet()
}
public fun withoutNormalization() {
withNormalizationBy = emptySet()
}
}
// Without Serializable GradleRunner tests fail
// TODO add more options
public data class CsvOptionsDsl(var delimiter: Char = ',') : Serializable
public data class JsonOptionsDsl(
var typeClashTactic: JSON.TypeClashTactic = JSON.TypeClashTactic.ARRAY_AND_VALUE_COLUMNS,
var keyValuePaths: List<JsonPath> = emptyList(),
var unifyNumbers: Boolean = true,
) : Serializable
/**
* Represents the configuration options for JDBC data source.
*
* @property [user] The username used to authenticate with the database. Default is an empty string.
* @property [password] The password used to authenticate with the database. Default is an empty string.
* @property [tableName] The name of the table to generate schema for. Default is an empty string.
* @property [sqlQuery] The SQL query used to generate schema. Default is an empty string.
*/
public data class JdbcOptionsDsl(
var user: String = "",
var password: String = "",
var tableName: String = "",
var sqlQuery: String = "",
) : Serializable
@@ -0,0 +1,240 @@
package org.jetbrains.dataframe.gradle
import com.google.devtools.ksp.gradle.KspAATask
import com.google.devtools.ksp.gradle.KspTask
import com.google.devtools.ksp.gradle.KspTaskJvm
import org.gradle.api.Plugin
import org.gradle.api.Project
import org.gradle.api.file.FileCollection
import org.gradle.api.logging.LogLevel
import org.gradle.api.tasks.TaskProvider
import org.gradle.internal.logging.services.DefaultLoggingManager
import org.gradle.kotlin.dsl.create
import org.jetbrains.kotlin.gradle.dsl.ExplicitApiMode
import org.jetbrains.kotlin.gradle.dsl.KotlinAndroidProjectExtension
import org.jetbrains.kotlin.gradle.dsl.KotlinJvmProjectExtension
import org.jetbrains.kotlin.gradle.dsl.KotlinMultiplatformExtension
import org.jetbrains.kotlin.gradle.dsl.KotlinProjectExtension
import org.jetbrains.kotlin.gradle.tasks.BaseKotlinCompile
import java.io.File
import java.net.URL
import java.nio.file.Path
import java.nio.file.Paths
import java.util.Locale
public class SchemaGeneratorPlugin : Plugin<Project> {
override fun apply(target: Project) {
val extension = target.extensions.create<SchemaGeneratorExtension>("dataframes")
extension.project = target
target.afterEvaluate {
val appliedPlugin =
KOTLIN_EXTENSIONS.firstNotNullOfOrNull {
target.extensions.findByType(it.extensionClass)?.let { ext -> AppliedPlugin(ext, it) }
}
if (appliedPlugin == null) {
target.logger.warn("Schema generator plugin applied, but no Kotlin plugin was found")
}
val generationTasks = extension.schemas.map {
registerTask(target, extension, appliedPlugin, it)
}
val generateAll = target.tasks.register("generateDataFrames") {
group = GROUP
dependsOn(*generationTasks.toTypedArray())
}
tasks.withType(KspTask::class.java).configureEach {
dependsOn(generateAll)
dependsOn(*generationTasks.toTypedArray())
}
tasks.withType(KspAATask::class.java).configureEach {
error(
"Detected KSP2. This is not supported by the DataFrame Gradle/Ksp plugin. Add 'ksp.useKSP2=false' to 'gradle.properties'.",
)
}
tasks.withType(BaseKotlinCompile::class.java).configureEach {
dependsOn(generateAll)
}
}
}
private fun registerTask(
target: Project,
extension: SchemaGeneratorExtension,
appliedPlugin: AppliedPlugin?,
schema: Schema,
): TaskProvider<GenerateDataSchemaTask> {
val interfaceName = getInterfaceName(schema)
fun propertyError(property: String): Nothing {
error(
"No supported Kotlin plugin was found. Please apply one or specify property $property for schema $interfaceName explicitly",
)
}
val sourceSetName by lazy {
schema.sourceSet
?: extension.sourceSet
?: (appliedPlugin ?: propertyError("sourceSet")).sourceSetConfiguration.defaultSourceSet
}
val src: File = schema.src
?: run {
appliedPlugin ?: propertyError("src")
val sourceSet = appliedPlugin.kotlinExtension.sourceSets.getByName(sourceSetName)
val src = target.file(Paths.get("build/generated/dataframe/", sourceSetName, "kotlin").toFile())
// Add the new sources to the source set
sourceSet.kotlin.srcDir(src)
// Configure the right ksp task to be aware of these new sources
val kspTaskName = "ksp${sourceSetName.replaceFirstChar { it.uppercase() }}Kotlin"
target.tasks.withType(KspTaskJvm::class.java).configureEach {
if ((sourceSetName == "main" && name == "kspKotlin") || name == kspTaskName) {
source(src)
}
}
src
}
val packageName = schema.name?.let { name -> extractPackageName(name) }
?: schema.packageName
?: extension.packageName
?: run {
(appliedPlugin ?: propertyError("packageName"))
val sourceSet = appliedPlugin.kotlinExtension.sourceSets.getByName(sourceSetName)
val path = appliedPlugin.sourceSetConfiguration.getKotlinRoot(
sourceSet.kotlin.sourceDirectories,
sourceSetName,
)
val src = target.file(path)
inferPackageName(src)
}
val visibility = schema.visibility
?: extension.visibility
?: run {
if (appliedPlugin != null) {
when (appliedPlugin.kotlinExtension.explicitApi) {
null -> DataSchemaVisibility.IMPLICIT_PUBLIC
ExplicitApiMode.Strict -> DataSchemaVisibility.EXPLICIT_PUBLIC
ExplicitApiMode.Warning -> DataSchemaVisibility.EXPLICIT_PUBLIC
ExplicitApiMode.Disabled -> DataSchemaVisibility.IMPLICIT_PUBLIC
}
} else {
DataSchemaVisibility.IMPLICIT_PUBLIC
}
}
val defaultPath = schema.defaultPath ?: extension.defaultPath ?: true
val delimiters = schema.withNormalizationBy ?: extension.withNormalizationBy ?: setOf('\t', ' ', '_')
return target.tasks.register("generateDataFrame$interfaceName", GenerateDataSchemaTask::class.java) {
(logging as? DefaultLoggingManager)?.setLevelInternal(LogLevel.QUIET)
group = GROUP
data.set(schema.data)
this.interfaceName.set(interfaceName)
this.packageName.set(packageName)
this.src.set(src)
this.schemaVisibility.set(visibility)
this.csvOptions.set(schema.csvOptions)
this.jsonOptions.set(schema.jsonOptions)
this.jdbcOptions.set(schema.jdbcOptions)
this.defaultPath.set(defaultPath)
this.delimiters.set(delimiters)
this.enableExperimentalOpenApi.set(extension.enableExperimentalOpenApi)
}
}
private fun getInterfaceName(schema: Schema): String? {
val rawName = schema.name?.substringAfterLast('.')
?: fileName(schema.data)
?.toCamelCaseByDelimiters(delimiters)
?.replaceFirstChar { if (it.isLowerCase()) it.titlecase(Locale.getDefault()) else it.toString() }
?.removeSurrounding("`")
?: return null
NameChecker.checkValidIdentifier(rawName)
return rawName
}
private val delimiters = "[\\s_]".toRegex()
private class AppliedPlugin(
val kotlinExtension: KotlinProjectExtension,
val sourceSetConfiguration: SourceSetConfiguration<*>,
)
private class SourceSetConfiguration<T : KotlinProjectExtension>(
val extensionClass: Class<T>,
val defaultSourceSet: String,
) {
fun getKotlinRoot(sourceDirectories: FileCollection, sourceSetName: String): File {
fun sourceSet(lang: String) = Paths.get("src", sourceSetName, lang)
val ktSet = sourceSet("kotlin")
val javaSet = sourceSet("java")
val isKotlinRoot: (Path) -> Boolean = { f -> f.endsWith(ktSet) }
val genericRoot = sourceDirectories.find { isKotlinRoot(it.toPath()) }
if (genericRoot != null) return genericRoot
val androidSpecificRoot = if (extensionClass == KotlinAndroidProjectExtension::class.java) {
val isAndroidKotlinRoot: (Path) -> Boolean = { f -> f.endsWith(javaSet) }
sourceDirectories.find { isAndroidKotlinRoot(it.toPath()) }
} else {
error("Directory '$ktSet' was not found in $sourceSetName. Please, specify 'src' explicitly")
}
return androidSpecificRoot
?: error(
"Directory '$ktSet' or '$javaSet' was not found in $sourceSetName. Please, specify 'src' explicitly",
)
}
}
private fun fileName(data: Any?): String? =
when (data) {
is String -> extractFileName(data)
is URL -> extractFileName(data)
is File -> extractFileName(data)
else -> throw IllegalArgumentException(
"data for schema must be File, URL or String, but was ${data?.javaClass ?: ""}($data)",
)
}
private fun extractPackageName(fqName: String): String? {
val packageName = fqName
.substringBeforeLast('.')
.takeIf { it != fqName }
if (packageName != null) {
NameChecker.checkValidPackageName(packageName)
}
return packageName
}
private fun inferPackageName(root: File): String {
val node = root.findDeepestCommonSubdirectory()
val parentPath = root.absolutePath
return node.absolutePath
.removePrefix(parentPath)
.removePrefix(File.separator)
.replace(File.separatorChar, '.')
.let {
when {
it.isEmpty() -> "dataframe"
it.endsWith(".dataframe") -> it
else -> "$it.dataframe"
}
}
}
private companion object {
private val KOTLIN_EXTENSIONS = sequenceOf(
SourceSetConfiguration(KotlinJvmProjectExtension::class.java, "main"),
SourceSetConfiguration(KotlinMultiplatformExtension::class.java, "jvmMain"),
SourceSetConfiguration(KotlinAndroidProjectExtension::class.java, "main"),
)
private const val GROUP = "dataframe"
}
}
@@ -0,0 +1,15 @@
package org.jetbrains.dataframe.gradle
import java.util.Locale
public fun String.toCamelCaseByDelimiters(delimiters: Regex): String =
split(delimiters)
.joinToCamelCaseString()
.replaceFirstChar { it.lowercase(Locale.getDefault()) }
public fun List<String>.joinToCamelCaseString(): String =
joinToString(separator = "") { s ->
s.replaceFirstChar {
if (it.isLowerCase()) it.titlecase(Locale.getDefault()) else it.toString()
}
}
@@ -0,0 +1 @@
PREPROCESSOR_VERSION=%PREPROCESSOR_VERSION%
@@ -0,0 +1,143 @@
package org.jetbrains.dataframe.gradle
import io.kotest.assertions.asClue
import io.kotest.assertions.throwables.shouldNotThrowAny
import io.kotest.assertions.throwables.shouldThrow
import io.kotest.assertions.throwables.shouldThrowAny
import io.kotest.matchers.shouldBe
import kotlinx.serialization.SerializationException
import org.jetbrains.kotlinx.dataframe.DataFrame
import org.jetbrains.kotlinx.dataframe.api.isEmpty
import org.jetbrains.kotlinx.dataframe.io.read
import org.jetbrains.kotlinx.dataframe.io.readCsv
import org.jetbrains.kotlinx.dataframe.io.readSqlTable
import org.junit.Test
import java.io.File
import java.io.FileNotFoundException
import java.io.IOException
import java.net.URL
import java.nio.file.Files
import java.nio.file.Paths
import java.sql.DriverManager
import kotlin.io.path.absolutePathString
class DataFrameReadTest {
@Test
fun `file that does not exists`() {
val temp = Files.createTempDirectory("").toFile()
val definitelyDoesNotExists = File(temp, "absolutelyRandomName")
shouldThrow<FileNotFoundException> {
DataFrame.read(definitelyDoesNotExists)
}
}
@Test
fun `file with invalid json`() {
val temp = Files.createTempDirectory("").toFile()
val invalidJson = File(temp, "test.json").also { it.writeText(".") }
shouldThrow<IllegalStateException> {
DataFrame.read(invalidJson)
}
}
@Test
fun `file with invalid csv`() {
val temp = Files.createTempDirectory("").toFile()
val invalidCsv = File(temp, "test.csv").also { it.writeText("") }
DataFrame.read(invalidCsv).isEmpty() shouldBe true
}
@Test
fun `invalid url`() {
val exception = shouldThrowAny {
DataFrame.read("http:://example.com")
}
exception.asClue {
(exception is IllegalArgumentException || exception is IOException) shouldBe true
}
}
@Test
fun `valid url`() {
useHostedJson("{}") {
DataFrame.read(URL(it))
}
}
@Test
fun `path that is valid url`() {
useHostedJson("{}") {
DataFrame.read(it)
}
}
@Test
fun `URL with invalid JSON`() {
useHostedJson("<invalid json>") { url ->
shouldThrow<SerializationException> {
DataFrame.read(url).also { println(it) }
}
}
}
@Test
fun `data accessible and readable`() {
shouldNotThrowAny {
DataFrame.readCsv(Paths.get("../../data/jetbrains repositories.csv").absolutePathString(), skipLines = 1)
}
}
@Test
fun `csvSample is valid csv`() {
val temp = Files.createTempFile("f", "csv").toFile()
temp.writeText(TestData.csvSample)
val df = DataFrame.read(temp)
df.columnNames() shouldBe listOf("name", "age")
}
@Test
fun `jdbcSample is valid jdbc`() {
DriverManager.getConnection("jdbc:h2:mem:test;DB_CLOSE_DELAY=-1;MODE=MySQL;DATABASE_TO_UPPER=false")
.use { connection ->
// Create table Customer
connection.createStatement().execute(
"""
CREATE TABLE Customer (
id INT PRIMARY KEY,
name VARCHAR(50),
age INT
)
""".trimIndent(),
)
// Create table Sale
connection.createStatement().execute(
"""
CREATE TABLE Sale (
id INT PRIMARY KEY,
customerId INT,
amount DECIMAL(10, 2)
)
""".trimIndent(),
)
// add data to the Customer table
connection.createStatement().execute("INSERT INTO Customer (id, name, age) VALUES (1, 'John', 40)")
connection.createStatement().execute("INSERT INTO Customer (id, name, age) VALUES (2, 'Alice', 25)")
connection.createStatement().execute("INSERT INTO Customer (id, name, age) VALUES (3, 'Bob', 47)")
// add data to the Sale table
connection.createStatement().execute("INSERT INTO Sale (id, customerId, amount) VALUES (1, 1, 100.50)")
connection.createStatement().execute("INSERT INTO Sale (id, customerId, amount) VALUES (2, 2, 50.00)")
connection.createStatement().execute("INSERT INTO Sale (id, customerId, amount) VALUES (3, 1, 75.25)")
connection.createStatement().execute("INSERT INTO Sale (id, customerId, amount) VALUES (4, 3, 35.15)")
val dfCustomer = DataFrame.readSqlTable(connection, "Customer")
dfCustomer.columnNames() shouldBe listOf("id", "name", "age")
val dfSale = DataFrame.readSqlTable(connection, "Sale")
dfSale.columnNames() shouldBe listOf("id", "customerId", "amount")
}
}
}
@@ -0,0 +1,23 @@
package org.jetbrains.dataframe.gradle
import io.ktor.http.ContentType
import io.ktor.server.application.call
import io.ktor.server.engine.embeddedServer
import io.ktor.server.netty.Netty
import io.ktor.server.response.respondText
import io.ktor.server.routing.get
import io.ktor.server.routing.routing
fun useHostedJson(json: String, f: (url: String) -> Unit) {
// duplicated in ksp/EmbeddedServerRunners.kt
val port = 14771
val server = embeddedServer(Netty, port = port) {
routing {
get("/test.json") {
call.respondText(json, ContentType.Application.Json)
}
}
}.start()
f("http://0.0.0.0:$port/test.json")
server.stop(500, 1000)
}
@@ -0,0 +1,45 @@
package org.jetbrains.dataframe.gradle
import org.junit.Test
import java.io.File
import java.net.URL
class ExtractFileNameTest {
@Test
fun `1`() {
val name =
extractFileName(
"https://raw.githubusercontent.com/Kotlin/dataframe/8ea139c35aaf2247614bb227756d6fdba7359f6a/data/playlistItems.json",
)
assert(name == "playlistItems")
}
@Test
fun `2`() {
val name =
extractFileName(
URL(
"https://raw.githubusercontent.com/Kotlin/dataframe/8ea139c35aaf2247614bb227756d6fdba7359f6a/data/playlistItems.json",
),
)
assert(name == "playlistItems")
}
@Test
fun `3`() {
val name = extractFileName("/etc/example/file.json")
assert(name == "file")
}
@Test
fun `4`() {
val name = extractFileName(File("/etc/example/file.json"))
assert(name == "file")
}
@Test
fun `5`() {
val name = extractFileName("abc")
assert(name == "abc")
}
}
@@ -0,0 +1,8 @@
package org.jetbrains.dataframe.gradle
import io.kotest.matchers.should
import java.io.File
import java.nio.file.Paths
fun File.shouldEndWith(first: String, vararg path: String) =
this should { it.endsWith(Paths.get(first, *path).toFile()) }
@@ -0,0 +1,59 @@
package org.jetbrains.dataframe.gradle
import io.kotest.matchers.shouldBe
import org.junit.Before
import org.junit.Test
import java.io.File
import java.nio.file.Files
internal class FileTraversalTest {
private lateinit var temp: File
@Before
fun init() {
temp = Files.createTempDirectory("temp").toFile()
}
@Test
fun directoryWithFileIsNotMiddlePackage() {
val leaf = File(temp, "a/b/c").also { it.mkdirs() }
val file = File(leaf, "test.txt").also { it.createNewFile() }
leaf.isMiddlePackage() shouldBe false
file.isMiddlePackage() shouldBe false
listOf(leaf, file).filterNot { it.isMiddlePackage() }.first() shouldBe leaf
}
@Test
fun emptySubdirectory() {
val leaf = File(temp, "a/b/c").also { it.mkdirs() }
temp.findDeepestCommonSubdirectory() shouldBe leaf
}
@Test
fun subdirectoryWithFile() {
val leaf = File(temp, "a/b/c").also { it.mkdirs() }
File(leaf, "test.txt").also { it.createNewFile() }
temp.findDeepestCommonSubdirectory() shouldBe leaf
}
@Test
fun forkAtDepth0() {
File(temp, "a/b/c").also { it.mkdirs() }
File(temp, "b/c/d").also { it.mkdirs() }
temp.findDeepestCommonSubdirectory() shouldBe temp
}
@Test
fun forkAtDepth1() {
val a = File(temp, "a").also { it.mkdirs() }
File(a, "b/c").also { it.mkdirs() }
File(a, "c/d").also { it.mkdirs() }
temp.findDeepestCommonSubdirectory() shouldBe a
}
@Test
fun noSubdirectories() {
temp.findDeepestCommonSubdirectory() shouldBe temp
}
}
@@ -0,0 +1,34 @@
package org.jetbrains.dataframe.gradle
import org.gradle.testkit.runner.BuildResult
import org.gradle.testkit.runner.GradleRunner
import org.intellij.lang.annotations.Language
import java.io.File
import java.nio.file.Files
fun runGradleBuild(
task: String,
@Language("kts") settingsGradle: (File) -> String = { "" },
@Language("kts") build: (File) -> String,
): Build {
val buildDir = Files.createTempDirectory("test").toFile()
val buildFile = File(buildDir, "build.gradle.kts")
buildFile.writeText(build(buildDir))
val settingsFile = File(buildDir, "settings.gradle.kts")
settingsFile.writeText(settingsGradle(buildDir))
val propertiesFile = File(buildDir, "gradle.properties")
propertiesFile.writeText("ksp.useKSP2=false")
return Build(buildDir, gradleRunner(buildDir, task).build())
}
fun gradleRunner(buildDir: File, task: String): GradleRunner =
GradleRunner.create()
.withProjectDir(buildDir)
// if we use api from the newest Gradle release, a user project will fail with NoSuchMethod
// testing compatibility with an older Gradle version ensures our plugin can run on as many versions as possible
.withGradleVersion("8.5")
.withPluginClasspath()
.withArguments(task, "--stacktrace", "--info")
.withDebug(true)
data class Build(val buildDir: File, val buildResult: BuildResult)
@@ -0,0 +1,38 @@
package org.jetbrains.dataframe.gradle
import io.kotest.assertions.throwables.shouldThrow
import org.gradle.api.internal.plugins.PluginApplicationException
import org.junit.Test
class KotlinPluginExpectationsTest {
@Test
fun `project can only apply 1 Kotlin plugin 1`() {
val project = makeProject()
project.plugins.apply("org.jetbrains.kotlin.jvm")
shouldThrow<PluginApplicationException> {
project.plugins.apply("org.jetbrains.kotlin.android")
project.evaluate()
}
}
@Test
fun `project can only apply 1 Kotlin plugin 2`() {
val project = makeProject()
project.plugins.apply("org.jetbrains.kotlin.multiplatform")
shouldThrow<PluginApplicationException> {
project.plugins.apply("org.jetbrains.kotlin.android")
project.evaluate()
}
}
@Test
fun `project can only apply 1 Kotlin plugin 3`() {
val project = makeProject()
project.plugins.apply("org.jetbrains.kotlin.multiplatform")
shouldThrow<PluginApplicationException> {
project.plugins.apply("org.jetbrains.kotlin.jvm")
project.evaluate()
}
}
}
@@ -0,0 +1,17 @@
package org.jetbrains.dataframe.gradle
import io.kotest.matchers.shouldBe
import org.junit.Test
class RegexExpectationsTest {
// for package name validation
@Test
fun `regex split returns non empty list for empty string`() {
"".split("123".toRegex()) shouldBe listOf("")
}
@Test
fun `regex split ignore delimiter`() {
"1.2.3".split("\\.".toRegex()) shouldBe listOf("1", "2", "3")
}
}
@@ -0,0 +1,413 @@
package org.jetbrains.dataframe.gradle
import io.kotest.assertions.asClue
import io.kotest.inspectors.forOne
import io.kotest.matchers.shouldBe
import io.kotest.matchers.string.shouldContain
import org.gradle.testkit.runner.TaskOutcome
import org.jetbrains.kotlin.gradle.dsl.KotlinJvmProjectExtension
import org.junit.Test
import java.io.File
import java.nio.file.Files
internal class SchemaGeneratorPluginTest {
private companion object {
private val KOTLIN_VERSION = TestData.kotlinVersion
}
@Test
fun `plugin configured via configure`() {
val (_, result) = runGradleBuild(":generateDataFrameTest") {
// language=kts
"""
import java.net.URI
import org.jetbrains.dataframe.gradle.SchemaGeneratorExtension
plugins {
kotlin("jvm") version "$KOTLIN_VERSION"
id("org.jetbrains.kotlinx.dataframe")
}
repositories {
mavenCentral()
}
configure<SchemaGeneratorExtension> {
schema {
data = URI("https://raw.githubusercontent.com/Kotlin/dataframe/8ea139c35aaf2247614bb227756d6fdba7359f6a/data/playlistItems.json").toURL()
name = "Test"
packageName = "org.test"
}
}
""".trimIndent()
}
result.task(":generateDataFrameTest")?.outcome shouldBe TaskOutcome.SUCCESS
}
@Test
fun `plugin configured via extension DSL`() {
val (_, result) = runGradleBuild(":generateDataFrameTest") {
// language=kts
"""
import java.net.URI
import org.jetbrains.dataframe.gradle.SchemaGeneratorExtension
plugins {
kotlin("jvm") version "$KOTLIN_VERSION"
id("org.jetbrains.kotlinx.dataframe")
}
repositories {
mavenCentral()
}
dataframes {
schema {
data = URI("https://raw.githubusercontent.com/Kotlin/dataframe/8ea139c35aaf2247614bb227756d6fdba7359f6a/data/playlistItems.json").toURL()
name = "Test"
packageName = "org.test"
}
}
""".trimIndent()
}
result.task(":generateDataFrameTest")?.outcome shouldBe TaskOutcome.SUCCESS
}
@Test
fun `plugin configured via extension DSL with Groovy`() {
val buildDir = Files.createTempDirectory("test").toFile()
val buildFile = File(buildDir, "build.gradle")
buildFile.writeText(
// language=groovy
"""
import java.net.URI
import org.jetbrains.dataframe.gradle.SchemaGeneratorExtension
plugins {
id "org.jetbrains.kotlin.jvm" version "$KOTLIN_VERSION"
id "org.jetbrains.kotlinx.dataframe"
}
repositories {
mavenCentral()
}
dataframes {
schema {
data = new URI("https://raw.githubusercontent.com/Kotlin/dataframe/8ea139c35aaf2247614bb227756d6fdba7359f6a/data/playlistItems.json").toURL()
name = "Test"
packageName = "org.test"
}
}
""".trimIndent(),
)
val result = gradleRunner(buildDir, ":generateDataFrameTest").build()
result.task(":generateDataFrameTest")?.outcome shouldBe TaskOutcome.SUCCESS
}
@Test
fun `delimiters configured with Groovy`() {
val buildDir = Files.createTempDirectory("test").toFile()
val buildFile = File(buildDir, "build.gradle")
buildFile.writeText(
// language=groovy
"""
import java.net.URL
import org.jetbrains.dataframe.gradle.SchemaGeneratorExtension
plugins {
id "org.jetbrains.kotlin.jvm" version "$KOTLIN_VERSION"
id "org.jetbrains.kotlinx.dataframe"
}
repositories {
mavenCentral()
}
dataframes {
schema {
data = new URI("https://raw.githubusercontent.com/Kotlin/dataframe/8ea139c35aaf2247614bb227756d6fdba7359f6a/data/playlistItems.json").toURL()
name = "Test"
packageName = "org.test"
withNormalizationBy('-_\t ')
}
}
""".trimIndent(),
)
val result = gradleRunner(buildDir, ":generateDataFrameTest").build()
result.task(":generateDataFrameTest")?.outcome shouldBe TaskOutcome.SUCCESS
}
@Test
fun `plugin configure multiple schemas from URLs via extension`() {
val (_, result) = runGradleBuild(":generateDataFrames") {
// language=kts
"""
import java.net.URI
import org.jetbrains.dataframe.gradle.SchemaGeneratorExtension
plugins {
kotlin("jvm") version "$KOTLIN_VERSION"
id("org.jetbrains.kotlinx.dataframe")
}
repositories {
mavenCentral()
}
dataframes {
schema {
data = URI("https://raw.githubusercontent.com/Kotlin/dataframe/8ea139c35aaf2247614bb227756d6fdba7359f6a/data/playlistItems.json").toURL()
name = "Test"
packageName = "org.test"
}
schema {
data = URI("https://raw.githubusercontent.com/Kotlin/dataframe/master/data/jetbrains_repositories.csv").toURL()
name = "Schema"
packageName = "org.test"
}
}
""".trimIndent()
}
result.task(":generateDataFrameTest")?.outcome shouldBe TaskOutcome.SUCCESS
result.task(":generateDataFrameSchema")?.outcome shouldBe TaskOutcome.SUCCESS
}
@Test
fun `plugin configure multiple schemas from strings via extension`() {
val (_, result) = runGradleBuild(":generateDataFrames") { buildDir ->
File(buildDir, "data").also {
it.mkdirs()
File(it, TestData.csvName).writeText(TestData.csvSample)
File(it, TestData.jsonName).writeText(TestData.jsonSample)
}
// language=kts
"""
import org.jetbrains.dataframe.gradle.SchemaGeneratorExtension
plugins {
kotlin("jvm") version "$KOTLIN_VERSION"
id("org.jetbrains.kotlinx.dataframe")
}
repositories {
mavenCentral()
}
dataframes {
schema {
data = "data/${TestData.csvName}"
name = "Test"
packageName = "org.test"
}
schema {
data = "data/${TestData.jsonName}"
name = "Schema"
packageName = "org.test"
}
}
""".trimIndent()
}
result.task(":generateDataFrameTest")?.outcome shouldBe TaskOutcome.SUCCESS
result.task(":generateDataFrameSchema")?.outcome shouldBe TaskOutcome.SUCCESS
}
@Test
fun `data is string and relative path`() {
val (_, result) = runGradleBuild(":generateDataFrameTest") { buildDir ->
val dataDir = File(buildDir, "data").also { it.mkdirs() }
File(dataDir, TestData.jsonName).writeText(TestData.jsonSample)
// language=kts
"""
import org.jetbrains.dataframe.gradle.SchemaGeneratorExtension
plugins {
kotlin("jvm") version "$KOTLIN_VERSION"
id("org.jetbrains.kotlinx.dataframe")
}
repositories {
mavenCentral()
}
dataframes {
schema {
data = "data/${TestData.jsonName}"
name = "Test"
packageName = "org.test"
}
}
""".trimIndent()
}
result.task(":generateDataFrameTest")?.outcome shouldBe TaskOutcome.SUCCESS
}
@Test
fun `data is string and absolute path`() {
val (_, result) = runGradleBuild(":generateDataFrameTest") { buildDir ->
val dataDir = File(buildDir, "data").also { it.mkdirs() }
val file = File(dataDir, TestData.jsonName).also { it.writeText(TestData.jsonSample) }
val absolutePath = file.absolutePath.replace(File.separatorChar, '/')
// language=kts
"""
import org.jetbrains.dataframe.gradle.SchemaGeneratorExtension
plugins {
kotlin("jvm") version "$KOTLIN_VERSION"
id("org.jetbrains.kotlinx.dataframe")
}
repositories {
mavenCentral()
}
dataframes {
schema {
data = "$absolutePath"
name = "Test"
packageName = "org.test"
}
}
""".trimIndent()
}
result.task(":generateDataFrameTest")?.outcome shouldBe TaskOutcome.SUCCESS
}
@Test
fun `data is string and url`() {
val (_, result) = runGradleBuild(":generateDataFrameTest") { buildDir ->
println("Build dir: $buildDir")
// language=kts
"""
import org.jetbrains.dataframe.gradle.SchemaGeneratorExtension
plugins {
kotlin("jvm") version "$KOTLIN_VERSION"
id("org.jetbrains.kotlinx.dataframe")
}
repositories {
mavenCentral()
}
dataframes {
schema {
data = "https://raw.githubusercontent.com/Kotlin/dataframe/8ea139c35aaf2247614bb227756d6fdba7359f6a/data/ghost.json"
name = "Test"
packageName = "org.test"
}
}
""".trimIndent()
}
result.task(":generateDataFrameTest")?.outcome shouldBe TaskOutcome.SUCCESS
}
@Test
fun `data is OpenApi string and url`() {
val (_, result) = runGradleBuild(":generateDataFrameTest") { buildDir ->
println("Build dir: $buildDir")
// language=kts
"""
import org.jetbrains.dataframe.gradle.SchemaGeneratorExtension
plugins {
kotlin("jvm") version "$KOTLIN_VERSION"
id("org.jetbrains.kotlinx.dataframe")
}
repositories {
mavenCentral()
}
dataframes {
schema {
data = "https://raw.githubusercontent.com/Kotlin/dataframe/2ad1c2f5d27267fa9b2abde1bf611fa50c5abce9/data/petstore.json"
name = "Test"
packageName = "org.test"
}
}
""".trimIndent()
}
result.task(":generateDataFrameTest")?.outcome shouldBe TaskOutcome.SUCCESS
}
@Test
fun `custom csv delimiter`() {
val (buildDir, result) = runGradleBuild(":generateDataFrameTest") { buildDir ->
val csv = "semicolons.csv"
val data = File(buildDir, csv)
data.writeText(
"""
a;b;c
1;2;3
""".trimIndent(),
)
// language=kts
"""
import org.jetbrains.dataframe.gradle.SchemaGeneratorExtension
plugins {
kotlin("jvm") version "$KOTLIN_VERSION"
id("org.jetbrains.kotlinx.dataframe")
}
repositories {
mavenCentral()
}
dataframes {
schema {
data = "$csv"
name = "Test"
packageName = "org.test"
csvOptions {
delimiter = ';'
}
}
}
""".trimIndent()
}
result.task(":generateDataFrameTest")?.outcome shouldBe TaskOutcome.SUCCESS
File(buildDir, "build/generated/dataframe/main/kotlin/org/test/Test.Generated.kt").asClue {
it.readLines().let {
it.forOne {
it.shouldContain("val a")
}
it.forOne {
it.shouldContain("val b")
}
it.forOne {
it.shouldContain("val c")
}
}
}
}
@Test
fun `most specific sourceSet is used in the packageName inference`() {
val project = makeProject()
project.plugins.apply(SchemaGeneratorPlugin::class.java)
project.plugins.apply("org.jetbrains.kotlin.jvm")
project.extensions.getByType(KotlinJvmProjectExtension::class.java).apply {
sourceSets.create("main1")
}
project.extensions.getByType(SchemaGeneratorExtension::class.java).apply {
sourceSet = "main"
schema {
sourceSet = "main1"
data = "123"
name = "321"
}
}
project.file("src/main1/kotlin/org/example/test").also { it.mkdirs() }
project.evaluate()
(project.tasks.getByName("generateDataFrame321") as GenerateDataSchemaTask)
.dataSchema
.get()
.shouldBe(
project.file("build/generated/dataframe/main1/kotlin/org/example/test/dataframe/321.Generated.kt"),
)
}
}
@@ -0,0 +1,30 @@
package org.jetbrains.dataframe.gradle
import io.kotest.inspectors.forAny
import io.kotest.matchers.shouldBe
import org.jetbrains.kotlin.gradle.dsl.KotlinJvmProjectExtension
import org.jetbrains.kotlin.gradle.dsl.KotlinMultiplatformExtension
import org.junit.Test
class SourceSetsExpectationsTest {
@Test
fun `there is main in default JVM project`() {
val project = makeProject()
project.plugins.apply("org.jetbrains.kotlin.jvm")
project.extensions.getByType(KotlinJvmProjectExtension::class.java).let { extension ->
val main = extension.sourceSets.getByName("main")
main.kotlin.sourceDirectories.toList().forAny {
it.shouldEndWith("src", "main", "kotlin")
}
}
}
@Test
fun `there is no jvmMain in default Multiplatform project`() {
val project = makeProject()
project.plugins.apply("org.jetbrains.kotlin.multiplatform")
project.extensions.getByType(KotlinMultiplatformExtension::class.java).let {
it.sourceSets.findByName("jvmMain") shouldBe null
}
}
}
@@ -0,0 +1,21 @@
package org.jetbrains.dataframe.gradle
import org.jetbrains.kotlinx.dataframe.BuildConfig
object TestData {
val csvSample =
"""
name, age
Alice, 15
Bob,
""".trimIndent()
val csvName = "data.csv"
val jsonSample = """{"name": "Test"}"""
val jsonName = "test.json"
val kotlinVersion = BuildConfig.KOTLIN_VERSION
}
@@ -0,0 +1,49 @@
@file:Suppress("UnstableApiUsage")
package org.jetbrains.dataframe.gradle
import org.gradle.api.Project
import org.gradle.api.internal.project.DefaultProject
import org.gradle.api.internal.project.ProjectInternal
import org.gradle.api.provider.Provider
import org.gradle.build.event.BuildEventsListenerRegistry
import org.gradle.internal.service.DefaultServiceRegistry
import org.gradle.internal.service.scopes.Scope
import org.gradle.internal.service.scopes.ServiceScope
import org.gradle.testfixtures.ProjectBuilder
import org.gradle.tooling.events.OperationCompletionListener
import java.lang.reflect.Field
import java.util.concurrent.atomic.AtomicReference
internal fun makeProject(): ProjectInternal {
val project = ProjectBuilder.builder().build() as ProjectInternal
addBuildEventsListenerRegistryMock(project)
return project
}
/**
* In Gradle 6.7-rc-1 BuildEventsListenerRegistry service is not created in we need it in order
* to instantiate AGP. This creates a fake one and injects it - http://b/168630734.
*/
internal fun addBuildEventsListenerRegistryMock(project: Project) {
try {
val projectScopeServices = (project as DefaultProject).services as DefaultServiceRegistry
val state: Field = DefaultServiceRegistry::class.java.getDeclaredField("state")
state.isAccessible = true
@Suppress("UNCHECKED_CAST")
val stateValue: AtomicReference<Any> = state.get(projectScopeServices) as AtomicReference<Any>
val enumClass = Class.forName(DefaultServiceRegistry::class.java.name + "\$State")
stateValue.set(enumClass.enumConstants[0])
// add service and set state so that future mutations are not allowed
projectScopeServices.add(BuildEventsListenerRegistryMock::class.java, BuildEventsListenerRegistryMock)
stateValue.set(enumClass.enumConstants[1])
} catch (e: Throwable) {
throw RuntimeException(e)
}
}
@ServiceScope(Scope.Project::class)
object BuildEventsListenerRegistryMock : BuildEventsListenerRegistry {
override fun onTaskCompletion(listener: Provider<out OperationCompletionListener>?) = Unit
}
@@ -0,0 +1,31 @@
package org.jetbrains.dataframe.gradle.taskProperties
import io.kotest.matchers.shouldBe
import org.jetbrains.dataframe.gradle.GenerateDataSchemaTask
import org.jetbrains.dataframe.gradle.SchemaGeneratorExtension
import org.jetbrains.dataframe.gradle.SchemaGeneratorPlugin
import org.jetbrains.dataframe.gradle.makeProject
import org.junit.Test
class TaskCsvOptionsPropertyTest {
@Test
fun `configure delimiter`() {
val project = makeProject()
project.plugins.apply(SchemaGeneratorPlugin::class.java)
val tab = '\t'
project.extensions.getByType(SchemaGeneratorExtension::class.java).apply {
schema {
data = "/test/data.csv"
name = "org.example.Data"
src = project.projectDir
csvOptions {
delimiter = tab
}
}
}
project.evaluate()
(project.tasks.getByName("generateDataFrameData") as GenerateDataSchemaTask)
.csvOptions.get()
.delimiter shouldBe tab
}
}
@@ -0,0 +1,36 @@
package org.jetbrains.dataframe.gradle.taskProperties
import io.kotest.assertions.throwables.shouldNotThrow
import io.kotest.matchers.shouldBe
import org.gradle.api.ProjectConfigurationException
import org.jetbrains.dataframe.gradle.GenerateDataSchemaTask
import org.jetbrains.dataframe.gradle.SchemaGeneratorExtension
import org.jetbrains.dataframe.gradle.SchemaGeneratorPlugin
import org.jetbrains.dataframe.gradle.makeProject
import org.jetbrains.kotlin.gradle.dsl.KotlinJvmProjectExtension
import org.junit.Test
class TaskDataSchemaPropertyTest {
@Test
fun `extension sourceSet present in project`() {
val project = makeProject()
project.plugins.apply(SchemaGeneratorPlugin::class.java)
project.plugins.apply("org.jetbrains.kotlin.jvm")
project.extensions.getByType(KotlinJvmProjectExtension::class.java).apply {
sourceSets.create("main1")
}
project.extensions.getByType(SchemaGeneratorExtension::class.java).apply {
sourceSet = "main1"
schema {
data = "123"
name = "org.example.my.321"
}
}
shouldNotThrow<ProjectConfigurationException> {
project.evaluate()
}
(project.tasks.getByName("generateDataFrame321") as GenerateDataSchemaTask)
.dataSchema.get()
.shouldBe(project.file("build/generated/dataframe/main1/kotlin/org/example/my/321.Generated.kt"))
}
}
@@ -0,0 +1,93 @@
package org.jetbrains.dataframe.gradle.taskProperties
import io.kotest.assertions.throwables.shouldThrow
import io.kotest.matchers.shouldNotBe
import io.kotest.matchers.string.shouldContain
import org.gradle.api.ProjectConfigurationException
import org.jetbrains.dataframe.gradle.SchemaGeneratorExtension
import org.jetbrains.dataframe.gradle.SchemaGeneratorPlugin
import org.jetbrains.dataframe.gradle.makeProject
import org.junit.Test
class TaskNamePropertyTest {
@Test
fun `task name is last part of FQ name`() {
val project = makeProject()
project.plugins.apply(SchemaGeneratorPlugin::class.java)
project.extensions.getByType(SchemaGeneratorExtension::class.java).apply {
schema {
data = "123"
name = "org.example.my.321"
src = project.projectDir
}
}
project.evaluate()
(project.tasks.findByName("generateDataFrame321") shouldNotBe null)
}
@Test
fun `name from task property have higher priority then inferred from data`() {
val project = makeProject()
project.plugins.apply(SchemaGeneratorPlugin::class.java)
project.extensions.getByType(SchemaGeneratorExtension::class.java).apply {
schema {
data = "/test/data.json"
name = "org.example.my.321"
src = project.projectDir
}
}
project.evaluate()
(project.tasks.findByName("generateDataFrame321") shouldNotBe null)
}
@Test
fun `task name contains invalid characters`() {
val project = makeProject()
project.plugins.apply(SchemaGeneratorPlugin::class.java)
project.extensions.getByType(SchemaGeneratorExtension::class.java).apply {
schema {
data = "123"
name = "org.test.example.[321]"
src = project.projectDir
}
}
val exception = shouldThrow<ProjectConfigurationException> {
project.evaluate()
}
exception.causes.single().message shouldContain "[321] contains illegal characters: [,]"
}
@Test
fun `data name should not override invalid name`() {
val project = makeProject()
project.plugins.apply(SchemaGeneratorPlugin::class.java)
project.extensions.getByType(SchemaGeneratorExtension::class.java).apply {
schema {
data =
"https://datalore-samples.s3-eu-west-1.amazonaws.com/datalore_gallery_of_samples/city_population.csv"
name = "org.test.example.[321]"
src = project.projectDir
}
}
val exception = shouldThrow<ProjectConfigurationException> {
project.evaluate()
}
exception.causes.single().message shouldContain "[321] contains illegal characters: [,]"
}
@Test
fun `name convention is data file name`() {
val project = makeProject()
project.plugins.apply(SchemaGeneratorPlugin::class.java)
project.extensions.getByType(SchemaGeneratorExtension::class.java).apply {
schema {
data =
"https://datalore-samples.s3-eu-west-1.amazonaws.com/datalore_gallery_of_samples/city_population.csv"
packageName = ""
src = project.projectDir
}
}
project.evaluate()
project.tasks.getByName("generateDataFrameCityPopulation") shouldNotBe null
}
}
@@ -0,0 +1,131 @@
package org.jetbrains.dataframe.gradle.taskProperties
import io.kotest.assertions.throwables.shouldThrow
import io.kotest.matchers.shouldBe
import org.gradle.api.ProjectConfigurationException
import org.jetbrains.dataframe.gradle.GenerateDataSchemaTask
import org.jetbrains.dataframe.gradle.SchemaGeneratorExtension
import org.jetbrains.dataframe.gradle.SchemaGeneratorPlugin
import org.jetbrains.dataframe.gradle.makeProject
import org.junit.Test
import java.io.File
class TaskPackageNamePropertyTest {
@Test
fun `task inherit default packageName from extension`() {
val project = makeProject()
project.plugins.apply(SchemaGeneratorPlugin::class.java)
project.extensions.getByType(SchemaGeneratorExtension::class.java).apply {
packageName = "org.example.test"
schema {
data = "123"
name = "321"
src = project.projectDir
}
}
project.evaluate()
(project.tasks.getByName("generateDataFrame321") as GenerateDataSchemaTask)
.packageName.get() shouldBe "org.example.test"
}
@Test
fun `task packageName overrides packageName from extension`() {
val project = makeProject()
project.plugins.apply(SchemaGeneratorPlugin::class.java)
project.extensions.getByType(SchemaGeneratorExtension::class.java).apply {
packageName = "org.example.test"
schema {
data = "123"
packageName = "org.example.my"
name = "321"
src = project.projectDir
}
}
project.evaluate()
(project.tasks.getByName("generateDataFrame321") as GenerateDataSchemaTask)
.packageName.get() shouldBe "org.example.my"
}
@Test
fun `task packageName convention is package part of FQ name`() {
val project = makeProject()
project.plugins.apply(SchemaGeneratorPlugin::class.java)
project.extensions.getByType(SchemaGeneratorExtension::class.java).apply {
packageName = "org.example.test"
schema {
data = "123"
name = "org.example.my.321"
src = project.projectDir
}
}
project.evaluate()
(project.tasks.findByName("generateDataFrame321") as GenerateDataSchemaTask)
.packageName.get() shouldBe "org.example.my"
}
@Test
fun `name package part overrides packageName`() {
val project = makeProject()
project.plugins.apply(SchemaGeneratorPlugin::class.java)
project.extensions.getByType(SchemaGeneratorExtension::class.java).apply {
schema {
data = "123"
packageName = "org.example.test"
name = "org.example.my.321"
src = project.projectDir
}
}
project.evaluate()
(project.tasks.findByName("generateDataFrame321") as GenerateDataSchemaTask)
.packageName.get() shouldBe "org.example.my"
}
@Test
fun `illegal characters in package part of name cause exception`() {
val project = makeProject()
project.plugins.apply(SchemaGeneratorPlugin::class.java)
project.extensions.getByType(SchemaGeneratorExtension::class.java).apply {
schema {
data = "123"
name = "`[org]`.321"
}
}
shouldThrow<ProjectConfigurationException> {
project.evaluate()
}
}
@Test
fun `task infers packageName from directory structure`() {
val project = makeProject()
project.plugins.apply(SchemaGeneratorPlugin::class.java)
project.plugins.apply("org.jetbrains.kotlin.jvm")
File(project.projectDir, "/src/main/kotlin/org/test/").also { it.mkdirs() }
project.extensions.getByType(SchemaGeneratorExtension::class.java).apply {
schema {
data = "123"
name = "321"
}
}
project.evaluate()
(project.tasks.getByName("generateDataFrame321") as GenerateDataSchemaTask)
.packageName.get() shouldBe "org.test.dataframe"
}
@Test
fun `task will not add _dataframe_ if inferred package ends with _dataframe_`() {
val project = makeProject()
project.plugins.apply(SchemaGeneratorPlugin::class.java)
project.plugins.apply("org.jetbrains.kotlin.jvm")
File(project.projectDir, "/src/main/kotlin/org/dataframe/").also { it.mkdirs() }
project.extensions.getByType(SchemaGeneratorExtension::class.java).apply {
schema {
data = "123"
name = "321"
}
}
project.evaluate()
(project.tasks.getByName("generateDataFrame321") as GenerateDataSchemaTask)
.packageName.get() shouldBe "org.dataframe"
}
}
@@ -0,0 +1,164 @@
package org.jetbrains.dataframe.gradle.taskProperties
import io.kotest.assertions.throwables.shouldNotThrow
import io.kotest.assertions.throwables.shouldThrow
import io.kotest.inspectors.forOne
import io.kotest.matchers.collections.shouldHaveSize
import io.kotest.matchers.shouldBe
import io.kotest.matchers.string.shouldContain
import org.gradle.api.ProjectConfigurationException
import org.jetbrains.dataframe.gradle.GenerateDataSchemaTask
import org.jetbrains.dataframe.gradle.SchemaGeneratorExtension
import org.jetbrains.dataframe.gradle.SchemaGeneratorPlugin
import org.jetbrains.dataframe.gradle.makeProject
import org.jetbrains.kotlin.gradle.dsl.KotlinJvmProjectExtension
import org.junit.Test
class TaskSourceSetPropertyTest {
@Test
fun `extension sourceSet present in project`() {
val project = makeProject()
project.plugins.apply(SchemaGeneratorPlugin::class.java)
project.plugins.apply("org.jetbrains.kotlin.jvm")
project.extensions.getByType(KotlinJvmProjectExtension::class.java).apply {
sourceSets.create("main1")
}
project.extensions.getByType(SchemaGeneratorExtension::class.java).apply {
sourceSet = "main1"
schema {
data = "123"
name = "org.example.my.321"
}
}
shouldNotThrow<ProjectConfigurationException> {
project.evaluate()
}
(project.tasks.getByName("generateDataFrame321") as GenerateDataSchemaTask)
.src.get()
.shouldBe(project.file("build/generated/dataframe/main1/kotlin/"))
}
@Test
fun `extension sourceSet not present in project`() {
val project = makeProject()
project.plugins.apply(SchemaGeneratorPlugin::class.java)
project.plugins.apply("org.jetbrains.kotlin.jvm")
project.extensions.getByType(SchemaGeneratorExtension::class.java).apply {
sourceSet = "main1"
schema {
data = "123"
name = "org.example.my.321"
}
}
val exception = shouldThrow<ProjectConfigurationException> {
project.evaluate()
}
exception.causes.shouldHaveSize(1)
exception.causes.forOne { it.message shouldContain "KotlinSourceSet with name 'main1' not found" }
}
@Test
fun `extension sourceSet not specified`() {
val project = makeProject()
project.plugins.apply(SchemaGeneratorPlugin::class.java)
project.plugins.apply("org.jetbrains.kotlin.jvm")
project.extensions.getByType(SchemaGeneratorExtension::class.java).apply {
schema {
data = "123"
name = "org.example.my.321"
}
}
project.evaluate()
(project.tasks.getByName("generateDataFrame321") as GenerateDataSchemaTask)
.src.get()
.shouldBe(project.file("build/generated/dataframe/main/kotlin/"))
}
@Test
fun `choose most specific sourceSet`() {
val project = makeProject()
project.plugins.apply(SchemaGeneratorPlugin::class.java)
project.plugins.apply("org.jetbrains.kotlin.jvm")
project.extensions.getByType(KotlinJvmProjectExtension::class.java).apply {
sourceSets.create("main1")
}
project.extensions.getByType(SchemaGeneratorExtension::class.java).apply {
sourceSet = "main"
schema {
sourceSet = "main1"
data = "123"
name = "321"
}
}
project.evaluate()
(project.tasks.getByName("generateDataFrame321") as GenerateDataSchemaTask)
.src.get()
.shouldBe(project.file("build/generated/dataframe/main1/kotlin/"))
}
@Test
fun `extension sourceSet specified but no kotlin plugin found`() {
val project = makeProject()
project.plugins.apply(SchemaGeneratorPlugin::class.java)
project.extensions.getByType(SchemaGeneratorExtension::class.java).apply {
packageName = "org.example.test"
sourceSet = "myMain"
schema {
data = "123"
name = "org.example.my.321"
}
}
val exception = shouldThrow<ProjectConfigurationException> {
project.evaluate()
}
exception.causes.shouldHaveSize(1)
exception.causes.forOne {
it.message shouldContain
"No supported Kotlin plugin was found. Please apply one or specify property src for schema 321 explicitly"
}
}
@Test
fun `task with explicit src don't evaluates sourceSet`() {
val project = makeProject()
project.plugins.apply(SchemaGeneratorPlugin::class.java)
project.extensions.getByType(SchemaGeneratorExtension::class.java).apply {
packageName = "org.example.test"
sourceSet = "myMain"
schema {
data = "123"
src = project.file("src/main/kotlin")
name = "org.example.my.321"
}
}
shouldNotThrow<ProjectConfigurationException> {
project.evaluate()
}
(project.tasks.getByName("generateDataFrame321") as GenerateDataSchemaTask)
.dataSchema.get()
.shouldBe(project.file("src/main/kotlin/org/example/my/321.Generated.kt"))
}
@Test
fun `task and extension sourceSet not present in project`() {
val project = makeProject()
project.plugins.apply(SchemaGeneratorPlugin::class.java)
project.extensions.getByType(SchemaGeneratorExtension::class.java).apply {
packageName = "org.example.test"
sourceSet = "myMain"
schema {
data = "123"
sourceSet = "myMain1"
name = "org.example.my.321"
}
}
val exception = shouldThrow<ProjectConfigurationException> {
project.evaluate()
}
exception.causes.shouldHaveSize(1)
exception.causes.forOne {
it.message shouldContain
"No supported Kotlin plugin was found. Please apply one or specify property src for schema 321 explicitly"
}
}
}
@@ -0,0 +1,81 @@
package org.jetbrains.dataframe.gradle.taskProperties
import io.kotest.assertions.throwables.shouldNotThrow
import io.kotest.matchers.shouldBe
import org.gradle.api.ProjectConfigurationException
import org.jetbrains.dataframe.gradle.DataSchemaVisibility
import org.jetbrains.dataframe.gradle.GenerateDataSchemaTask
import org.jetbrains.dataframe.gradle.SchemaGeneratorExtension
import org.jetbrains.dataframe.gradle.SchemaGeneratorPlugin
import org.jetbrains.dataframe.gradle.makeProject
import org.jetbrains.kotlin.gradle.dsl.KotlinJvmProjectExtension
import org.junit.Test
class TaskVisibilityPropertyTest {
@Test
fun `extension sourceSet present in project 1 `() {
val project = makeProject()
project.plugins.apply(SchemaGeneratorPlugin::class.java)
project.extensions.getByType(SchemaGeneratorExtension::class.java).apply {
sourceSet = "main1"
visibility = DataSchemaVisibility.INTERNAL
schema {
src = project.file("src")
data = "123"
name = "org.example.my.321"
}
}
shouldNotThrow<ProjectConfigurationException> {
project.evaluate()
}
(project.tasks.getByName("generateDataFrame321") as GenerateDataSchemaTask)
.schemaVisibility.get()
.shouldBe(DataSchemaVisibility.INTERNAL)
}
@Test
fun `extension sourceSet present in project 2`() {
val project = makeProject()
project.plugins.apply(SchemaGeneratorPlugin::class.java)
project.extensions.getByType(SchemaGeneratorExtension::class.java).apply {
sourceSet = "main1"
visibility = DataSchemaVisibility.INTERNAL
schema {
src = project.file("src")
visibility = DataSchemaVisibility.EXPLICIT_PUBLIC
data = "123"
name = "org.example.my.321"
}
}
shouldNotThrow<ProjectConfigurationException> {
project.evaluate()
}
(project.tasks.getByName("generateDataFrame321") as GenerateDataSchemaTask)
.schemaVisibility.get()
.shouldBe(DataSchemaVisibility.EXPLICIT_PUBLIC)
}
@Test
fun `extension sourceSet present in project`() {
val project = makeProject()
project.plugins.apply(SchemaGeneratorPlugin::class.java)
project.plugins.apply("org.jetbrains.kotlin.jvm")
project.extensions.getByType(KotlinJvmProjectExtension::class.java).apply {
sourceSets.create("main1")
explicitApi()
}
project.extensions.getByType(SchemaGeneratorExtension::class.java).apply {
sourceSet = "main1"
schema {
data = "123"
name = "org.example.my.321"
}
}
shouldNotThrow<ProjectConfigurationException> {
project.evaluate()
}
(project.tasks.getByName("generateDataFrame321") as GenerateDataSchemaTask)
.schemaVisibility.get()
.shouldBe(DataSchemaVisibility.EXPLICIT_PUBLIC)
}
}
+12
View File
@@ -0,0 +1,12 @@
## :plugins:expressions-converter
This Kotlin Compiler plugin, used by [:core](../../core), can extract intermediate
DataFrame expressions from `@TransformDataFrameExpressions` annotated functions.
It is used to generate sample "explainer dataframe" HTML files that can be used as iFrames on the documentation website.
Annotated functions in [core/.../test/.../samples/api](../../core/src/test/kotlin/org/jetbrains/kotlinx/dataframe/samples/api)
are tested, generated, and copied over to [docs/StardustDocs/resources/snippets](../../docs/StardustDocs/resources/snippets) by
our "explainer" [plugin callback proxy](../../core/src/test/kotlin/org/jetbrains/kotlinx/dataframe/explainer),
which hooks into [the TestBase class](../../core/src/test/kotlin/org/jetbrains/kotlinx/dataframe/samples/api/TestBase.kt) and
retrieves the intermediate DataFrame expressions thanks to this module.
@@ -0,0 +1,70 @@
plugins {
with(convention.plugins) {
alias(kotlinJvm8)
}
with(libs.plugins) {
alias(shadow)
alias(publisher)
}
}
group = "org.jetbrains.kotlinx.dataframe"
dependencies {
compileOnly(libs.kotlin.compiler)
testImplementation(libs.kotlin.compiler)
testImplementation(libs.kotlin.compiler.internal.test.framework)
testRuntimeOnly(projects.core)
testRuntimeOnly(libs.kotlin.test)
testRuntimeOnly(libs.kotlin.script.runtime)
testRuntimeOnly(libs.kotlin.annotations.jvm)
testImplementation(platform(libs.junit.bom))
testImplementation(libs.junit.jupiter)
testImplementation(libs.junit.platform.commons)
testImplementation(libs.junit.platform.launcher)
testImplementation(libs.junit.platform.runner)
testImplementation(libs.junit.platform.suite.api)
testRuntimeOnly(libs.junit.jupiter.engine)
}
tasks.test {
useJUnitPlatform()
doFirst {
setLibraryProperty("org.jetbrains.kotlin.test.kotlin-stdlib", "kotlin-stdlib")
setLibraryProperty("org.jetbrains.kotlin.test.kotlin-reflect", "kotlin-reflect")
setLibraryProperty("org.jetbrains.kotlin.test.kotlin-test", "kotlin-test")
setLibraryProperty("org.jetbrains.kotlin.test.kotlin-script-runtime", "kotlin-script-runtime")
setLibraryProperty("org.jetbrains.kotlin.test.kotlin-annotations-jvm", "kotlin-annotations-jvm")
}
}
sourceSets {
main {
java.setSrcDirs(listOf("src"))
resources.setSrcDirs(listOf("resources"))
}
test {
java.setSrcDirs(listOf("tests", "tests-gen"))
resources.setSrcDirs(listOf("testResources"))
}
}
tasks.register<JavaExec>("generateTests") {
classpath = sourceSets.test.get().runtimeClasspath
mainClass = "org.jetbrains.kotlinx.dataframe.GenerateTestsKt"
}
fun Test.setLibraryProperty(propName: String, jarName: String) {
val path = project.configurations
.testRuntimeClasspath.get()
.files
.find { """$jarName-\d.*jar""".toRegex().matches(it.name) }
?.absolutePath
?: return
systemProperty(propName, path)
}
@@ -0,0 +1,6 @@
#
# Copyright 2010-2020 JetBrains s.r.o. and Kotlin Programming Language contributors.
# Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
#
org.jetbrains.kotlinx.dataframe.ExplainerComponentRegistrar
@@ -0,0 +1,18 @@
package org.jetbrains.kotlinx.dataframe
import org.jetbrains.kotlin.backend.common.extensions.IrGenerationExtension
import org.jetbrains.kotlin.compiler.plugin.CompilerPluginRegistrar
import org.jetbrains.kotlin.compiler.plugin.ExperimentalCompilerApi
import org.jetbrains.kotlin.config.CompilerConfiguration
@OptIn(ExperimentalCompilerApi::class)
public class ExplainerComponentRegistrar : CompilerPluginRegistrar() {
override val supportsK2: Boolean
get() = true
override val pluginId: String = "org.jetbrains.kotlinx.dataframe.ExplainerComponentRegistrar"
override fun ExtensionStorage.registerExtensions(configuration: CompilerConfiguration) {
IrGenerationExtension.registerExtension(ExplainerIrGenerationExtension())
}
}
@@ -0,0 +1,14 @@
package org.jetbrains.kotlinx.dataframe
import org.jetbrains.kotlin.backend.common.extensions.IrGenerationExtension
import org.jetbrains.kotlin.backend.common.extensions.IrPluginContext
import org.jetbrains.kotlin.ir.declarations.IrModuleFragment
public class ExplainerIrGenerationExtension : IrGenerationExtension {
override fun generate(moduleFragment: IrModuleFragment, pluginContext: IrPluginContext) {
val lowering = ExplainerIrTransformer(pluginContext)
for (file in moduleFragment.files) {
lowering.lower(file)
}
}
}
@@ -0,0 +1,330 @@
@file:Suppress("ktlint:standard:no-unused-imports")
package org.jetbrains.kotlinx.dataframe
import org.jetbrains.kotlin.backend.common.FileLoweringPass
import org.jetbrains.kotlin.backend.common.extensions.IrPluginContext
import org.jetbrains.kotlin.descriptors.DescriptorVisibilities
import org.jetbrains.kotlin.descriptors.Modality
import org.jetbrains.kotlin.ir.IrElement
import org.jetbrains.kotlin.ir.IrElementBase
import org.jetbrains.kotlin.ir.IrStatement
import org.jetbrains.kotlin.ir.declarations.IrClass
import org.jetbrains.kotlin.ir.declarations.IrDeclarationOrigin
import org.jetbrains.kotlin.ir.declarations.IrField
import org.jetbrains.kotlin.ir.declarations.IrFile
import org.jetbrains.kotlin.ir.declarations.IrFunction
import org.jetbrains.kotlin.ir.declarations.IrParameterKind
import org.jetbrains.kotlin.ir.declarations.path
import org.jetbrains.kotlin.ir.expressions.IrBlockBody
import org.jetbrains.kotlin.ir.expressions.IrBody
import org.jetbrains.kotlin.ir.expressions.IrCall
import org.jetbrains.kotlin.ir.expressions.IrDeclarationReference
import org.jetbrains.kotlin.ir.expressions.IrExpression
import org.jetbrains.kotlin.ir.expressions.IrExpressionBody
import org.jetbrains.kotlin.ir.expressions.IrGetValue
import org.jetbrains.kotlin.ir.expressions.IrStatementOrigin
import org.jetbrains.kotlin.ir.expressions.impl.IrCallImplWithShape
import org.jetbrains.kotlin.ir.expressions.impl.IrConstImpl
import org.jetbrains.kotlin.ir.expressions.impl.IrFunctionExpressionImpl
import org.jetbrains.kotlin.ir.expressions.impl.IrGetObjectValueImpl
import org.jetbrains.kotlin.ir.expressions.impl.IrGetValueImpl
import org.jetbrains.kotlin.ir.symbols.UnsafeDuringIrConstructionAPI
import org.jetbrains.kotlin.ir.symbols.impl.IrSimpleFunctionSymbolImpl
import org.jetbrains.kotlin.ir.symbols.impl.IrValueParameterSymbolImpl
import org.jetbrains.kotlin.ir.types.classFqName
import org.jetbrains.kotlin.ir.types.makeNullable
import org.jetbrains.kotlin.ir.types.typeWith
import org.jetbrains.kotlin.ir.util.SetDeclarationsParentVisitor
import org.jetbrains.kotlin.ir.util.defaultType
import org.jetbrains.kotlin.ir.util.fqNameWhenAvailable
import org.jetbrains.kotlin.ir.util.isLocal
import org.jetbrains.kotlin.ir.visitors.IrElementTransformerVoid
import org.jetbrains.kotlin.ir.visitors.IrTransformer
import org.jetbrains.kotlin.ir.visitors.transformChildrenVoid
import org.jetbrains.kotlin.name.CallableId
import org.jetbrains.kotlin.name.ClassId
import org.jetbrains.kotlin.name.FqName
import org.jetbrains.kotlin.name.Name
import java.io.File
public data class ContainingDeclarations(val clazz: IrClass?, val function: IrFunction?, val statementIndex: Int = 0)
@OptIn(UnsafeDuringIrConstructionAPI::class)
public class ExplainerIrTransformer(public val pluginContext: IrPluginContext) :
IrTransformer<ContainingDeclarations>(),
FileLoweringPass {
public lateinit var file: IrFile
public lateinit var source: String
override fun lower(irFile: IrFile) {
var file: File
file = File("testData/box/${irFile.path}")
if (!file.exists()) {
file = File(irFile.path)
}
this.file = irFile
source = file.readText()
irFile.transformChildren(this, ContainingDeclarations(null, null))
irFile.transformChildrenVoid(object : IrElementTransformerVoid() {
override fun visitFunction(declaration: IrFunction): IrStatement {
declaration.acceptChildren(SetDeclarationsParentVisitor, declaration)
return super.visitFunction(declaration)
}
})
}
override fun visitBlockBody(body: IrBlockBody, data: ContainingDeclarations): IrBody {
for (i in 0 until body.statements.size) {
@Suppress("UNCHECKED_CAST")
body.statements.set(
index = i,
element = (body.statements[i] as IrElementBase)
.transform(
transformer = this,
data = data.copy(statementIndex = i),
) as IrStatement,
)
}
return body
}
override fun visitClass(declaration: IrClass, data: ContainingDeclarations): IrStatement =
super.visitClass(declaration, data.copy(clazz = declaration))
override fun visitFunction(declaration: IrFunction, data: ContainingDeclarations): IrStatement {
val annotated = declaration.annotations.any {
it.type.classFqName
?.shortName()
?.identifierOrNullIfSpecial
?.equals("TransformDataFrameExpressions") == true
}
return if (annotated) {
super.visitFunction(declaration, data.copy(function = declaration))
} else {
declaration
}
}
override fun visitElement(element: IrElement, data: ContainingDeclarations): IrElement {
element.transformChildren(this, data)
return element
}
override fun visitField(declaration: IrField, data: ContainingDeclarations): IrStatement {
if (declaration.isLocal) {
declaration.transformChildren(this, data)
}
return declaration
}
override fun visitExpressionBody(body: IrExpressionBody, data: ContainingDeclarations): IrBody = body
public val dataFrameLike: Set<FqName> = setOf(
FqName("org.jetbrains.kotlinx.dataframe.api.Pivot"),
FqName("org.jetbrains.kotlinx.dataframe.api.ReducedPivot"),
FqName("org.jetbrains.kotlinx.dataframe.api.PivotGroupBy"),
FqName("org.jetbrains.kotlinx.dataframe.api.ReducedPivotGroupBy"),
FqName("org.jetbrains.kotlinx.dataframe.api.SplitWithTransform"),
FqName("org.jetbrains.kotlinx.dataframe.api.Merge"),
FqName("org.jetbrains.kotlinx.dataframe.api.Split"),
FqName("org.jetbrains.kotlinx.dataframe.api.Gather"),
FqName("org.jetbrains.kotlinx.dataframe.api.Update"),
FqName("org.jetbrains.kotlinx.dataframe.api.Convert"),
FqName("org.jetbrains.kotlinx.dataframe.api.FormattedFrame"),
FqName("org.jetbrains.kotlinx.dataframe.api.GroupBy"),
FqName("org.jetbrains.kotlinx.dataframe.DataFrame"),
FqName("org.jetbrains.kotlinx.dataframe.DataRow"),
)
public val explainerPackage: FqName = FqName("org.jetbrains.kotlinx.dataframe.explainer")
override fun visitGetValue(expression: IrGetValue, data: ContainingDeclarations): IrExpression {
if (expression.startOffset < 0) return expression
if (expression.type.classFqName in dataFrameLike) {
return transformDataFrameExpression(expression, expression.symbol.owner.name, receiver = null, data)
}
return super.visitExpression(expression, data)
}
// also, what if expression type is not DataFrame, but Unit? and receiver expression is DataFrame at some point
override fun visitCall(expression: IrCall, data: ContainingDeclarations): IrExpression {
if (expression.startOffset < 0) return expression
if (expression.type.classFqName in dataFrameLike) {
if (expression.symbol.owner.name == Name.identifier("component1")) return expression
val extensionReceiverIndex =
expression.symbol.owner.parameters.indexOfFirst { it.kind == IrParameterKind.ExtensionReceiver }
var receiver: IrExpression?
// expression.arguments[extensionReceiverIndex] = extension callables,
// expression.dispatchReceiver = member callables such as "GroupBy.aggregate"
if (extensionReceiverIndex >= 0) {
receiver = expression.arguments[extensionReceiverIndex]!!
val transformedExtensionReceiver = receiver.transform(this, data)
expression.arguments[extensionReceiverIndex] = transformedExtensionReceiver
} else {
receiver = expression.dispatchReceiver
val transformedExtensionReceiver = expression.dispatchReceiver?.transform(this, data)
expression.dispatchReceiver = transformedExtensionReceiver
}
return transformDataFrameExpression(expression, expression.symbol.owner.name, receiver = receiver, data)
}
return super.visitExpression(expression, data)
}
private fun transformDataFrameExpression(
expression: IrDeclarationReference,
ownerName: Name,
receiver: IrExpression?,
data: ContainingDeclarations,
): IrCall {
val alsoReference = pluginContext
.referenceFunctions(
CallableId(FqName("kotlin"), Name.identifier("also")),
).single()
val result = IrCallImplWithShape(
startOffset = -1,
endOffset = -1,
type = expression.type,
symbol = alsoReference,
typeArgumentsCount = 1,
valueArgumentsCount = 1,
contextParameterCount = 0,
hasDispatchReceiver = true,
hasExtensionReceiver = true,
).apply {
val extensionReceiverIndex =
this.symbol.owner.parameters.indexOfFirst { it.kind == IrParameterKind.ExtensionReceiver }
if (extensionReceiverIndex >= 0) {
this.arguments[extensionReceiverIndex] = expression
} else {
this.insertExtensionReceiver(expression)
}
typeArguments[0] = expression.type
val symbol = IrSimpleFunctionSymbolImpl()
val alsoLambda = pluginContext.irFactory
.createSimpleFunction(
startOffset = -1,
endOffset = -1,
origin = IrDeclarationOrigin.LOCAL_FUNCTION_FOR_LAMBDA,
symbol = symbol,
name = Name.special("<anonymous>"),
visibility = DescriptorVisibilities.LOCAL,
modality = Modality.FINAL,
returnType = pluginContext.irBuiltIns.unitType,
isInline = false,
isExternal = false,
isTailrec = false,
isSuspend = false,
isOperator = false,
isInfix = false,
isExpect = false,
).apply {
// replace all regular value parameters with a single one `it`
parameters = parameters.filterNot { it.kind == IrParameterKind.Regular } +
pluginContext.irFactory.createValueParameter(
startOffset = -1,
endOffset = -1,
origin = IrDeclarationOrigin.DEFINED,
kind = IrParameterKind.Regular,
name = Name.identifier("it"),
type = expression.type,
isAssignable = false,
symbol = IrValueParameterSymbolImpl(),
varargElementType = null,
isCrossinline = false,
isNoinline = false,
isHidden = false,
)
val itSymbol = parameters.first { it.kind == IrParameterKind.Regular }.symbol
val source = try {
source.substring(expression.startOffset, expression.endOffset)
} catch (e: Exception) {
throw Exception("$expression ${ownerName.asString()} $source", e)
}
val expressionId = expressionId(expression)
val receiverId = receiver?.let { expressionId(it) }
val valueArguments = buildList<IrExpression?> {
add(source.irConstImpl()) // source: String
add(ownerName.asStringStripSpecialMarkers().irConstImpl()) // name: String
add(IrGetValueImpl(-1, -1, itSymbol)) // df: Any
add(expressionId.irConstImpl()) // id: String
add(receiverId.irConstImpl()) // receiverId: String?
add(data.clazz?.fqNameWhenAvailable?.asString().irConstImpl()) // containingClassFqName: String?
add(data.function?.name?.asString().irConstImpl()) // containingFunName: String?
add(
IrConstImpl.int(
-1,
-1,
pluginContext.irBuiltIns.intType,
data.statementIndex,
),
) // statementIndex: Int
}
body = pluginContext.irFactory.createBlockBody(-1, -1).apply {
val callableId = CallableId(
explainerPackage,
FqName("PluginCallbackProxy"),
Name.identifier("doAction"),
)
val doAction = pluginContext.referenceFunctions(callableId).single()
statements += IrCallImplWithShape(
startOffset = -1,
endOffset = -1,
type = doAction.owner.returnType,
symbol = doAction,
typeArgumentsCount = 0,
valueArgumentsCount = valueArguments.size,
contextParameterCount = 0,
hasDispatchReceiver = true,
hasExtensionReceiver = false,
).apply {
val clazz = ClassId(explainerPackage, Name.identifier("PluginCallbackProxy"))
val plugin = pluginContext.referenceClass(clazz)!!
val pluginType = plugin.owner.defaultType
dispatchReceiver = IrGetObjectValueImpl(-1, -1, pluginType, plugin)
val firstValueArgumentIndex = 1 // skipping dispatch receiver
valueArguments.forEachIndexed { i, argument ->
this.arguments[firstValueArgumentIndex + i] = argument
}
}
}
}
val alsoLambdaExpression = IrFunctionExpressionImpl(
startOffset = -1,
endOffset = -1,
type = pluginContext.irBuiltIns
.functionN(2)
.typeWith(listOf(expression.type, pluginContext.irBuiltIns.unitType)),
function = alsoLambda,
origin = IrStatementOrigin.LAMBDA,
)
val firstValueArgumentIndex = this.symbol.owner.parameters
.indexOfFirst { it.kind == IrParameterKind.Regular }
.takeUnless { it < 0 } ?: this.symbol.owner.parameters.size
this.arguments[firstValueArgumentIndex] = alsoLambdaExpression
}
return result
}
private fun String?.irConstImpl(): IrConstImpl {
val nullableString = pluginContext.irBuiltIns.stringType.makeNullable()
val argument = if (this == null) {
IrConstImpl.constNull(-1, -1, nullableString)
} else {
IrConstImpl.string(-1, -1, nullableString, this)
}
return argument
}
private fun expressionId(expression: IrExpression): String {
val line = file.fileEntry.getLineNumber(expression.startOffset)
val column = file.fileEntry.getColumnNumber(expression.startOffset)
return "${file.path}:${line + 1}:${column + 1}"
}
}
@@ -0,0 +1,16 @@
package org.jetbrains.kotlinx.dataframe.explainer
import org.jetbrains.kotlinx.dataframe.*
import org.jetbrains.kotlinx.dataframe.api.*
object PluginCallback {
fun doAction(any: Any) {}
}
fun box(): String {
val a: GroupBy<*, *> = dataFrameOf("a")(1).groupBy("a")
a.also {
PluginCallback.doAction(it)
}
return "OK"
}
@@ -0,0 +1,25 @@
package
package org {
package org.jetbrains {
package org.jetbrains.kotlinx {
package org.jetbrains.kotlinx.dataframe {
package org.jetbrains.kotlinx.dataframe.explainer {
public fun box(): kotlin.String
public object PluginCallback {
private constructor PluginCallback()
public final fun doAction(/*0*/ any: kotlin.Any): kotlin.Unit
public open override /*1*/ /*fake_override*/ fun equals(/*0*/ other: kotlin.Any?): kotlin.Boolean
public open override /*1*/ /*fake_override*/ fun hashCode(): kotlin.Int
public open override /*1*/ /*fake_override*/ fun toString(): kotlin.String
}
}
}
}
}
}
@@ -0,0 +1,156 @@
package org.jetbrains.kotlinx.dataframe.explainer
import org.jetbrains.kotlinx.dataframe.*
import org.jetbrains.kotlinx.dataframe.api.*
@TransformDataFrameExpressions
fun callChain(df: DataFrame<*>) {
val df1 = df
.filter { "age"<Int>() > 20 }
.groupBy("value")
.sum()
}
interface Person
class Wrapper {
val df = dataFrameOf("name", "age", "city", "weight")(
"Alice", 15, "London", 54,
"Bob", 45, "Dubai", 87,
"Charlie", 20, "Moscow", null,
"Charlie", 40, "Milan", null,
"Bob", 30, "Tokyo", 68,
"Alice", 20, null, 55,
"Charlie", 30, "Moscow", 90
)
val typed: DataFrame<Person> = df.cast()
@TransformDataFrameExpressions
fun ff() {
val name by column<String>()
val aggregated = typed.groupBy { name() }.aggregate {
(if (name().first().startsWith("A")) first() else null) into "agg"
}["agg"]
}
}
@TransformDataFrameExpressions
fun aggregateDf() {
val df = dataFrameOf("firstName", "lastName", "age", "city", "weight", "isHappy")(
"Alice", "Cooper", 15, "London", 54, true,
"Bob", "Dylan", 45, "Dubai", 87, true,
"Charlie", "Daniels", 20, "Moscow", null, false,
"Charlie", "Chaplin", 40, "Milan", null, true,
"Bob", "Marley", 30, "Tokyo", 68, true,
"Alice", "Wolf", 20, null, 55, false,
"Charlie", "Byrd", 30, "Moscow", 90, true
).group("firstName", "lastName").into("name")
df.groupBy("city").aggregate {
count() into "total"
count { "age"<Int>() > 18 } into "adults"
median("age") into "median age"
min("age") into "min age"
maxBy("age")["name"] into "oldest"
}
}
@TransformDataFrameExpressions
fun move() {
val df = dataFrameOf("firstName", "lastName", "age", "city", "weight", "isHappy")(
"Alice", "Cooper", 15, "London", 54, true,
"Bob", "Dylan", 45, "Dubai", 87, true,
"Charlie", "Daniels", 20, "Moscow", null, false,
"Charlie", "Chaplin", 40, "Milan", null, true,
"Bob", "Marley", 30, "Tokyo", 68, true,
"Alice", "Wolf", 20, null, 55, false,
"Charlie", "Byrd", 30, "Moscow", 90, true
).group("firstName", "lastName").into("name")
df.move("age", "weight").into { pathOf("info", it.name()) }
}
fun interface PluginCallback {
fun doAction(
source: String,
name: String,
df: Any,
id: String,
receiverId: String?,
containingClassFqName: String?,
containingFunName: String?,
statementIndex: Int
)
}
object PluginCallbackProxy : PluginCallback {
var action: PluginCallback = PluginCallback { _, _, _, _, _, _, _, _ -> Unit }
override fun doAction(source: String, name: String, df: Any, id: String, receiverId: String?, containingClassFqName: String?, containingFunName: String?, statemenIndex: Int) {
action.doAction(source, name, df, id, receiverId, containingClassFqName, containingFunName, statemenIndex)
}
}
//fun callChainTransformed(df: DataFrame<*>) {
// val df1 = df
// .filter { "age"<Int>() > 20 }
// .also { PluginCallbackProxy.doAction(""".filter { "age"<Int>() > 20 }""", "filter", it) }
//}
//fun callChainTransformed(df: DataFrame<*>) {
// val df1 = df
// .filter { "age"<Int>() > 20 }
// .also { PluginCallbackProxy.action(""".filter { "age"<Int>() > 20 }""", it) }
// .groupBy("something")
// .sum()
// .also { PluginCallbackProxy.action(""".groupBy("something").sum()""", it) }
//}
annotation class TransformDataFrameExpressions
fun box(): String {
val age by columnOf(10, 21, 30, 1)
val value by columnOf("a", "b", "c", "c")
val df = dataFrameOf(age, value)
val expressions = mutableListOf<String>()
PluginCallbackProxy.action = PluginCallback { source, _, df, id, receiverId, containingClassFqName, containingFunName, statementIndex ->
println("== Call ==")
expressions += source
if (df is AnyFrame) {
println(source)
df.print()
println(id)
println(receiverId)
println(containingClassFqName)
println(containingFunName)
println("statementIndex = ${statementIndex}")
} else {
println(df::class)
}
println("== End ==")
}
println("CallChain")
callChain(df)
println("expressions = ${expressions}")
expressions.clear()
println("ff")
Wrapper().ff()
println("expressions = ${expressions}")
expressions.clear()
// callChainTransformed(df)
println("aggregateDf")
aggregateDf()
println("expressions = ${expressions}")
expressions.clear()
println("move")
move()
println("expressions = ${expressions}")
expressions.clear()
return "OK"
}
@@ -0,0 +1,59 @@
package
package org {
package org.jetbrains {
package org.jetbrains.kotlinx {
package org.jetbrains.kotlinx.dataframe {
package org.jetbrains.kotlinx.dataframe.explainer {
@org.jetbrains.kotlinx.dataframe.explainer.TransformDataFrameExpressions public fun aggregateDf(): kotlin.Unit
public fun box(): kotlin.String
@org.jetbrains.kotlinx.dataframe.explainer.TransformDataFrameExpressions public fun callChain(/*0*/ df: org.jetbrains.kotlinx.dataframe.DataFrame<*>): kotlin.Unit
@org.jetbrains.kotlinx.dataframe.explainer.TransformDataFrameExpressions public fun move(): kotlin.Unit
public interface Person {
public open override /*1*/ /*fake_override*/ fun equals(/*0*/ other: kotlin.Any?): kotlin.Boolean
public open override /*1*/ /*fake_override*/ fun hashCode(): kotlin.Int
public open override /*1*/ /*fake_override*/ fun toString(): kotlin.String
}
public fun interface PluginCallback {
public abstract fun doAction(/*0*/ source: kotlin.String, /*1*/ name: kotlin.String, /*2*/ df: kotlin.Any, /*3*/ id: kotlin.String, /*4*/ receiverId: kotlin.String?, /*5*/ containingClassFqName: kotlin.String?, /*6*/ containingFunName: kotlin.String?, /*7*/ statementIndex: kotlin.Int): kotlin.Unit
public open override /*1*/ /*fake_override*/ fun equals(/*0*/ other: kotlin.Any?): kotlin.Boolean
public open override /*1*/ /*fake_override*/ fun hashCode(): kotlin.Int
public open override /*1*/ /*fake_override*/ fun toString(): kotlin.String
}
public object PluginCallbackProxy : org.jetbrains.kotlinx.dataframe.explainer.PluginCallback {
private constructor PluginCallbackProxy()
public final var action: org.jetbrains.kotlinx.dataframe.explainer.PluginCallback
public open override /*1*/ fun doAction(/*0*/ source: kotlin.String, /*1*/ name: kotlin.String, /*2*/ df: kotlin.Any, /*3*/ id: kotlin.String, /*4*/ receiverId: kotlin.String?, /*5*/ containingClassFqName: kotlin.String?, /*6*/ containingFunName: kotlin.String?, /*7*/ statemenIndex: kotlin.Int): kotlin.Unit
public open override /*1*/ /*fake_override*/ fun equals(/*0*/ other: kotlin.Any?): kotlin.Boolean
public open override /*1*/ /*fake_override*/ fun hashCode(): kotlin.Int
public open override /*1*/ /*fake_override*/ fun toString(): kotlin.String
}
public final annotation class TransformDataFrameExpressions : kotlin.Annotation {
public constructor TransformDataFrameExpressions()
public open override /*1*/ /*fake_override*/ fun equals(/*0*/ other: kotlin.Any?): kotlin.Boolean
public open override /*1*/ /*fake_override*/ fun hashCode(): kotlin.Int
public open override /*1*/ /*fake_override*/ fun toString(): kotlin.String
}
public final class Wrapper {
public constructor Wrapper()
public final val df: org.jetbrains.kotlinx.dataframe.DataFrame<*>
public final val typed: org.jetbrains.kotlinx.dataframe.DataFrame<org.jetbrains.kotlinx.dataframe.explainer.Person>
public open override /*1*/ /*fake_override*/ fun equals(/*0*/ other: kotlin.Any?): kotlin.Boolean
@org.jetbrains.kotlinx.dataframe.explainer.TransformDataFrameExpressions public final fun ff(): kotlin.Unit
public open override /*1*/ /*fake_override*/ fun hashCode(): kotlin.Int
public open override /*1*/ /*fake_override*/ fun toString(): kotlin.String
}
}
}
}
}
}
@@ -0,0 +1,21 @@
package org.jetbrains.kotlinx.dataframe.explainer
fun print(any: Any) {}
fun printM(pair: Pair<String, Any>) {
val (text, obj) = pair
println("# ${text}: $obj")
}
fun ir() {
val df = Any()
printM("Any()" to df)
}
fun box(): String {
val df = Any()
val df1 = 1 + 2
print(df)
print(df1)
return "OK"
}
@@ -0,0 +1,20 @@
package
package org {
package org.jetbrains {
package org.jetbrains.kotlinx {
package org.jetbrains.kotlinx.dataframe {
package org.jetbrains.kotlinx.dataframe.explainer {
public fun box(): kotlin.String
public fun ir(): kotlin.Unit
public fun print(/*0*/ any: kotlin.Any): kotlin.Unit
public fun printM(/*0*/ pair: kotlin.Pair<kotlin.String, kotlin.Any>): kotlin.Unit
}
}
}
}
}
@@ -0,0 +1,40 @@
package org.jetbrains.kotlinx.dataframe;
import com.intellij.testFramework.TestDataPath;
import org.jetbrains.kotlin.test.util.KtTestUtil;
import org.jetbrains.kotlin.test.TestMetadata;
import org.junit.jupiter.api.Test;
import java.io.File;
import java.util.regex.Pattern;
/** This class is generated by {@link org.jetbrains.kotlinx.dataframe.GenerateTestsKt}. DO NOT MODIFY MANUALLY */
@SuppressWarnings("all")
@TestMetadata("testData/box")
@TestDataPath("$PROJECT_ROOT")
public class ExplainerBlackBoxCodegenTestGenerated extends AbstractExplainerBlackBoxCodegenTest {
@Test
public void testAllFilesPresentInBox() {
KtTestUtil.assertAllTestsPresentByMetadataWithExcluded(this.getClass(), new File("testData/box"), Pattern.compile("^(.+)\\.kt$"), null, true);
}
@Test
@TestMetadata("any.kt")
public void testAny() {
runTest("testData/box/any.kt");
}
@Test
@TestMetadata("df.kt")
public void testDf() {
runTest("testData/box/df.kt");
}
@Test
@TestMetadata("test.kt")
public void testTest() {
runTest("testData/box/test.kt");
}
}
@@ -0,0 +1,78 @@
package org.jetbrains.kotlinx.dataframe
import org.jetbrains.kotlin.config.JvmTarget
import org.jetbrains.kotlin.platform.jvm.JvmPlatforms
import org.jetbrains.kotlin.test.TargetBackend
import org.jetbrains.kotlin.test.TestJdkKind
import org.jetbrains.kotlin.test.backend.BlackBoxCodegenSuppressor
import org.jetbrains.kotlin.test.backend.handlers.IrPrettyKotlinDumpHandler
import org.jetbrains.kotlin.test.backend.handlers.IrTextDumpHandler
import org.jetbrains.kotlin.test.backend.handlers.IrTreeVerifierHandler
import org.jetbrains.kotlin.test.backend.handlers.JvmBoxRunner
import org.jetbrains.kotlin.test.backend.ir.JvmIrBackendFacade
import org.jetbrains.kotlin.test.builders.TestConfigurationBuilder
import org.jetbrains.kotlin.test.builders.classicFrontendHandlersStep
import org.jetbrains.kotlin.test.builders.irHandlersStep
import org.jetbrains.kotlin.test.builders.jvmArtifactsHandlersStep
import org.jetbrains.kotlin.test.builders.psi2IrStep
import org.jetbrains.kotlin.test.directives.JvmEnvironmentConfigurationDirectives
import org.jetbrains.kotlin.test.frontend.classic.ClassicFrontendFacade
import org.jetbrains.kotlin.test.frontend.classic.handlers.ClassicDiagnosticsHandler
import org.jetbrains.kotlin.test.frontend.classic.handlers.DeclarationsDumpHandler
import org.jetbrains.kotlin.test.model.DependencyKind
import org.jetbrains.kotlin.test.model.FrontendKinds
import org.jetbrains.kotlin.test.model.TestModule
import org.jetbrains.kotlin.test.services.RuntimeClasspathProvider
import org.jetbrains.kotlin.test.services.TemporaryDirectoryManager
import org.jetbrains.kotlin.test.services.TestServices
import org.jetbrains.kotlin.test.services.configuration.CommonEnvironmentConfigurator
import org.jetbrains.kotlin.test.services.configuration.JvmEnvironmentConfigurator
import org.jetbrains.kotlinx.dataframe.services.TemporaryDirectoryManagerImplFixed
import java.io.File
open class AbstractExplainerBlackBoxCodegenTest : BaseTestRunner() {
override fun configure(builder: TestConfigurationBuilder): Unit =
with(builder) {
globalDefaults {
frontend = FrontendKinds.ClassicAndFIR
targetPlatform = JvmPlatforms.jvm8
dependencyKind = DependencyKind.Binary
targetBackend = TargetBackend.JVM_IR
}
defaultDirectives {
JvmEnvironmentConfigurationDirectives.JDK_KIND with TestJdkKind.FULL_JDK
JvmEnvironmentConfigurationDirectives.JVM_TARGET with JvmTarget.JVM_1_8
+JvmEnvironmentConfigurationDirectives.WITH_REFLECT
}
facadeStep(::ClassicFrontendFacade)
commonFirWithPluginFrontendConfiguration()
classicFrontendHandlersStep {
useHandlers(
::ClassicDiagnosticsHandler,
::DeclarationsDumpHandler,
)
}
psi2IrStep()
irHandlersStep {
useHandlers(
::IrPrettyKotlinDumpHandler,
::IrTextDumpHandler,
::IrTreeVerifierHandler,
)
}
facadeStep(::JvmIrBackendFacade)
jvmArtifactsHandlersStep {
useHandlers(::JvmBoxRunner)
}
useConfigurators(::JvmEnvironmentConfigurator, ::CommonEnvironmentConfigurator, ::PluginAnnotationsProvider)
useCustomRuntimeClasspathProviders(::MyClasspathProvider)
useAfterAnalysisCheckers(::BlackBoxCodegenSuppressor)
useAdditionalService<TemporaryDirectoryManager>(::TemporaryDirectoryManagerImplFixed)
}
class MyClasspathProvider(testServices: TestServices) : RuntimeClasspathProvider(testServices) {
override fun runtimeClassPaths(module: TestModule): List<File> =
(classpathFromClassloader(javaClass.classLoader) ?: error("no classpath"))
}
}
@@ -0,0 +1,20 @@
package org.jetbrains.kotlinx.dataframe
import org.jetbrains.kotlin.test.initIdeaConfiguration
import org.jetbrains.kotlin.test.runners.AbstractKotlinCompilerTest
import org.jetbrains.kotlin.test.services.EnvironmentBasedStandardLibrariesPathProvider
import org.jetbrains.kotlin.test.services.KotlinStandardLibrariesPathProvider
import org.junit.jupiter.api.BeforeAll
abstract class BaseTestRunner : AbstractKotlinCompilerTest() {
companion object {
@BeforeAll
@JvmStatic
fun setUp() {
initIdeaConfiguration()
}
}
override fun createKotlinStandardLibrariesPathProvider(): KotlinStandardLibrariesPathProvider =
EnvironmentBasedStandardLibrariesPathProvider
}
@@ -0,0 +1,19 @@
package org.jetbrains.kotlinx.dataframe
import org.jetbrains.kotlin.backend.common.extensions.IrGenerationExtension
import org.jetbrains.kotlin.compiler.plugin.CompilerPluginRegistrar
import org.jetbrains.kotlin.compiler.plugin.ExperimentalCompilerApi
import org.jetbrains.kotlin.config.CompilerConfiguration
import org.jetbrains.kotlin.test.model.TestModule
import org.jetbrains.kotlin.test.services.EnvironmentConfigurator
import org.jetbrains.kotlin.test.services.TestServices
class ExtensionRegistrarConfigurator(testServices: TestServices) : EnvironmentConfigurator(testServices) {
@OptIn(ExperimentalCompilerApi::class)
override fun CompilerPluginRegistrar.ExtensionStorage.registerCompilerExtensions(
module: TestModule,
configuration: CompilerConfiguration,
) {
IrGenerationExtension.registerExtension(ExplainerIrGenerationExtension())
}
}
@@ -0,0 +1,13 @@
package org.jetbrains.kotlinx.dataframe
import org.jetbrains.kotlin.generators.dsl.junit5.generateTestGroupSuiteWithJUnit5
fun main() {
generateTestGroupSuiteWithJUnit5 {
testGroup(testDataRoot = "testData", testsRoot = "tests-gen") {
testClass<AbstractExplainerBlackBoxCodegenTest> {
model("box")
}
}
}
}
@@ -0,0 +1,13 @@
package org.jetbrains.kotlinx.dataframe
import org.jetbrains.kotlin.cli.jvm.config.addJvmClasspathRoots
import org.jetbrains.kotlin.config.CompilerConfiguration
import org.jetbrains.kotlin.test.model.TestModule
import org.jetbrains.kotlin.test.services.EnvironmentConfigurator
import org.jetbrains.kotlin.test.services.TestServices
class PluginAnnotationsProvider(testServices: TestServices) : EnvironmentConfigurator(testServices) {
override fun configureCompilerConfiguration(configuration: CompilerConfiguration, module: TestModule) {
configuration.addJvmClasspathRoots(classpathFromClassloader(javaClass.classLoader) ?: error("no classpath"))
}
}
@@ -0,0 +1,479 @@
/*
* Copyright 2010-2018 JetBrains s.r.o. and Kotlin Programming Language contributors.
* Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
*/
package org.jetbrains.kotlinx.dataframe
import java.io.File
import java.io.FileInputStream
import java.io.FileNotFoundException
import java.net.JarURLConnection
import java.net.URL
import java.net.URLClassLoader
import java.util.jar.JarInputStream
import kotlin.reflect.KClass
// Kotlin Compiler dependencies
internal const val KOTLIN_JAVA_STDLIB_JAR = "kotlin-stdlib.jar"
internal const val KOTLIN_JAVA_REFLECT_JAR = "kotlin-reflect.jar"
internal const val KOTLIN_JAVA_SCRIPT_RUNTIME_JAR = "kotlin-script-runtime.jar"
internal const val TROVE4J_JAR = "trove4j.jar"
internal const val KOTLIN_SCRIPTING_COMPILER_JAR = "kotlin-scripting-compiler.jar"
internal const val KOTLIN_SCRIPTING_COMPILER_EMBEDDABLE_JAR = "kotlin-scripting-compiler-embeddable.jar"
internal const val KOTLIN_SCRIPTING_COMPILER_IMPL_JAR = "kotlin-scripting-compiler-impl.jar"
internal const val KOTLIN_SCRIPTING_COMPILER_IMPL_EMBEDDABLE_JAR = "kotlin-scripting-compiler-impl-embeddable.jar"
internal const val KOTLIN_SCRIPTING_COMMON_JAR = "kotlin-scripting-common.jar"
internal const val KOTLIN_SCRIPTING_JVM_JAR = "kotlin-scripting-jvm.jar"
internal const val KOTLIN_COMPILER_NAME = "kotlin-compiler"
internal const val KOTLIN_COMPILER_JAR = "$KOTLIN_COMPILER_NAME.jar"
private val JAR_COLLECTIONS_CLASSES_PATHS = arrayOf("BOOT-INF/classes", "WEB-INF/classes")
private val JAR_COLLECTIONS_LIB_PATHS = arrayOf("BOOT-INF/lib", "WEB-INF/lib")
private val JAR_COLLECTIONS_KEY_PATHS = JAR_COLLECTIONS_CLASSES_PATHS + JAR_COLLECTIONS_LIB_PATHS
internal const val JAR_MANIFEST_RESOURCE_NAME = "META-INF/MANIFEST.MF"
internal const val KOTLIN_SCRIPT_CLASSPATH_PROPERTY = "kotlin.script.classpath"
internal const val KOTLIN_COMPILER_CLASSPATH_PROPERTY = "kotlin.compiler.classpath"
internal const val KOTLIN_COMPILER_JAR_PROPERTY = "kotlin.compiler.jar"
internal const val KOTLIN_STDLIB_JAR_PROPERTY = "kotlin.java.stdlib.jar"
internal const val KOTLIN_REFLECT_JAR_PROPERTY = "kotlin.java.reflect.jar"
// obsolete name, but maybe still used in the wild
// TODO: consider removing
internal const val KOTLIN_RUNTIME_JAR_PROPERTY = "kotlin.java.runtime.jar"
internal const val KOTLIN_SCRIPT_RUNTIME_JAR_PROPERTY = "kotlin.script.runtime.jar"
private val validClasspathFilesExtensions = setOf("jar", "zip", "java")
private val validJarCollectionFilesExtensions = setOf("jar", "war", "zip")
class ClasspathExtractionException(message: String) : Exception(message)
fun classpathFromClassloader(currentClassLoader: ClassLoader, unpackJarCollections: Boolean = false): List<File>? {
val processedJars = hashSetOf<File>()
val unpackJarCollectionsDir by lazy {
File.createTempFile("unpackedJarCollections", null).canonicalFile.apply {
delete()
mkdir()
setReadable(false, false)
setWritable(false, false)
setExecutable(false, false)
setReadable(true, true)
setWritable(true, true)
setExecutable(true, true)
Runtime.getRuntime().addShutdownHook(
Thread {
deleteRecursively()
},
)
}
}
return allRelatedClassLoaders(currentClassLoader).flatMap { classLoader ->
var classPath = emptySequence<File>()
if (unpackJarCollections &&
JAR_COLLECTIONS_KEY_PATHS.any { classLoader.getResource(it)?.file?.isNotEmpty() == true }
) {
// if cache dir is specified, find all jar collections (spring boot fat jars and WARs so far, and unpack it accordingly
val jarCollections = JAR_COLLECTIONS_KEY_PATHS
.asSequence()
.flatMap { currentClassLoader.getResources(it).asSequence() }
.mapNotNull {
it.toContainingJarOrNull()?.takeIf { file ->
// additionally mark/check processed collection jars since unpacking is expensive
file.extension in validJarCollectionFilesExtensions && processedJars.add(file)
}
}
classPath +=
jarCollections
.flatMap { it.unpackJarCollection(unpackJarCollectionsDir) }
.filter { it.isValidClasspathFile() }
}
classPath += when (classLoader) {
is URLClassLoader -> {
classLoader.urLs.asSequence().mapNotNull { url -> url.toValidClasspathFileOrNull() }
}
else -> {
classLoader.classPathFromGetUrlsMethodOrNull()
?: classLoader.classPathFromTypicalResourceUrls()
}
}
classPath
}.filter { processedJars.add(it) }
.toList()
.takeIf { it.isNotEmpty() }
}
internal fun URL.toValidClasspathFileOrNull(): File? =
(toContainingJarOrNull() ?: toFileOrNull())?.takeIf { it.isValidClasspathFile() }
internal fun File.isValidClasspathFile(): Boolean =
isDirectory || (isFile && extension in validClasspathFilesExtensions)
private fun ClassLoader.classPathFromGetUrlsMethodOrNull(): Sequence<File>? =
try {
// e.g. for IDEA platform UrlClassLoader
val getUrls = this::class.java.getMethod("getUrls")
getUrls.isAccessible = true
val result = getUrls.invoke(this) as? List<Any?>
result?.asSequence()?.filterIsInstance<URL>()?.mapNotNull { it.toValidClasspathFileOrNull() }
} catch (e: Throwable) {
null
}
internal class ClassLoaderResourceRootFIlePathCalculator(private val keyResourcePath: String) {
private var keyResourcePathDepth = -1
operator fun invoke(resourceFile: File): File {
if (keyResourcePathDepth < 0) {
keyResourcePathDepth =
if (keyResourcePath.isBlank()) {
0
} else {
keyResourcePath.trim('/').count { it == '/' } + 1
}
}
var root = resourceFile
for (i in 0 until keyResourcePathDepth) {
root = root.parentFile
}
return root
}
}
internal fun ClassLoader.rawClassPathFromKeyResourcePath(keyResourcePath: String): Sequence<File> {
val resourceRootCalc = ClassLoaderResourceRootFIlePathCalculator(keyResourcePath)
return getResources(keyResourcePath).asSequence().mapNotNull { url ->
if (url.protocol == "jar") {
(url.openConnection() as? JarURLConnection)?.jarFileURL?.toFileOrNull()
} else {
url.toFileOrNull()?.let { resourceRootCalc(it) }
}
}
}
fun ClassLoader.classPathFromTypicalResourceUrls(): Sequence<File> =
// roots without manifest cases are detected in some test scenarios
// manifests without containing directory entries are detected in some optimized jars, e.g. after proguard
// TODO: investigate whether getting resources with empty name works in all situations
(rawClassPathFromKeyResourcePath("") + rawClassPathFromKeyResourcePath(JAR_MANIFEST_RESOURCE_NAME))
.distinct()
.filter { it.isValidClasspathFile() }
private fun File.unpackJarCollection(rootTempDir: File): Sequence<File> {
val targetDir = File.createTempFile(nameWithoutExtension, null, rootTempDir).apply {
delete()
mkdir()
}
return try {
ArrayList<File>().apply {
JarInputStream(FileInputStream(this@unpackJarCollection)).use { jarInputStream ->
for (classesDir in JAR_COLLECTIONS_CLASSES_PATHS) {
add(File(targetDir, classesDir))
}
do {
val entry = jarInputStream.nextJarEntry
if (entry != null) {
try {
if (!entry.isDirectory) {
val file = File(targetDir, entry.name)
if (JAR_COLLECTIONS_LIB_PATHS.any { entry.name.startsWith("$it/") }) {
add(file)
}
file.parentFile.mkdirs()
file.outputStream().use { outputStream ->
jarInputStream.copyTo(outputStream)
outputStream.flush()
}
}
} finally {
jarInputStream.closeEntry()
}
}
} while (entry != null)
}
}.asSequence()
} catch (e: Throwable) {
targetDir.deleteRecursively()
throw e
}
}
fun classpathFromClasspathProperty(): List<File>? =
System.getProperty("java.class.path")
?.split(String.format("\\%s", File.pathSeparatorChar).toRegex())
?.dropLastWhile(String::isEmpty)
?.map(::File)
fun classpathFromClass(classLoader: ClassLoader, klass: KClass<out Any>): List<File>? =
classpathFromFQN(classLoader, klass.qualifiedName!!)
fun classpathFromClass(klass: KClass<out Any>): List<File>? = classpathFromClass(klass.java.classLoader, klass)
inline fun <reified T : Any> classpathFromClass(): List<File>? = classpathFromClass(T::class)
fun classpathFromFQN(classLoader: ClassLoader, fqn: String): List<File>? {
val clp = "${fqn.replace('.', '/')}.class"
return classLoader
.rawClassPathFromKeyResourcePath(clp)
.filter { it.isValidClasspathFile() }
.toList()
.takeIf { it.isNotEmpty() }
}
fun File.matchMaybeVersionedFile(baseName: String) =
name == baseName ||
// for classes dirs
name == baseName.removeSuffix(".jar") ||
Regex(Regex.escape(baseName.removeSuffix(".jar")) + "(-\\d.*)?\\.jar").matches(name)
fun File.hasParentNamed(baseName: String): Boolean =
nameWithoutExtension == baseName || parentFile?.hasParentNamed(baseName) ?: false
private const val KOTLIN_COMPILER_EMBEDDABLE_JAR = "$KOTLIN_COMPILER_NAME-embeddable.jar"
// Iterating over classloaders tree in a regular, parent-first order
private fun allRelatedClassLoaders(
clsLoader: ClassLoader,
visited: MutableSet<ClassLoader> = HashSet(),
): Sequence<ClassLoader> {
if (!visited.add(clsLoader)) return emptySequence()
val singleParent = clsLoader.parent
if (singleParent != null) {
return sequenceOf(singleParent).flatMap { allRelatedClassLoaders(it, visited) } + clsLoader
}
return try {
val arrayOfClassLoaders = getParentClassLoaders(clsLoader)
// TODO: PluginClassLoader uses filtering (mustBeLoadedByPlatform), consider using the same logic, if possible
// (untill proper compiling from classloader instead of classpath is implemented)
arrayOfClassLoaders.asSequence().flatMap { allRelatedClassLoaders(it, visited) } + clsLoader
} catch (e: Throwable) {
sequenceOf(clsLoader)
}
}
private fun getParentClassLoaders(clsLoader: ClassLoader): Array<ClassLoader> =
try {
getParentsForNewClassLoader(clsLoader)
} catch (exception: NoSuchMethodException) {
try {
getParentsForOldClassLoader(clsLoader)
} catch (exception: NoSuchFieldException) {
// Possibly idea sources and kotlin compiler had diverged
emptyArray()
}
}
@Throws(NoSuchFieldException::class)
private fun getParentsForOldClassLoader(clsLoader: ClassLoader): Array<ClassLoader> {
// Correct way of getting parents in com.intellij.ide.plugins.cl.PluginClassLoader from IDEA 202 and earlier
val field = clsLoader.javaClass.getDeclaredField("myParents") // com.intellij.ide.plugins.cl.PluginClassLoader
field.isAccessible = true
@Suppress("UNCHECKED_CAST")
return field.get(clsLoader) as Array<ClassLoader>
}
@Throws(NoSuchMethodException::class)
private fun getParentsForNewClassLoader(clsLoader: ClassLoader): Array<ClassLoader> {
// Correct way of getting parents in com.intellij.ide.plugins.cl.PluginClassLoader from IDEA 203+
val method = clsLoader.javaClass.getDeclaredMethod("getAllParents")
method.isAccessible = true
@Suppress("UNCHECKED_CAST")
return method.invoke(clsLoader) as Array<ClassLoader>
}
internal fun List<File>.takeIfContainsAll(vararg keyNames: String): List<File>? =
takeIf { classpath ->
keyNames.all { key -> classpath.any { it.matchMaybeVersionedFile(key) } }
}
internal fun List<File>.filterIfContainsAll(vararg keyNames: String): List<File>? {
val foundKeys = mutableSetOf<String>()
val res = arrayListOf<File>()
for (cpentry in this) {
for (prefix in keyNames) {
if (cpentry.matchMaybeVersionedFile(prefix) || (cpentry.isDirectory && cpentry.hasParentNamed(prefix))) {
foundKeys.add(prefix)
res.add(cpentry)
break
}
}
}
return res.takeIf { foundKeys.containsAll(keyNames.asList()) }
}
internal fun List<File>.takeIfContainsAny(vararg keyNames: String): List<File>? =
takeIf { classpath ->
keyNames.any { key -> classpath.any { it.matchMaybeVersionedFile(key) } }
}
fun scriptCompilationClasspathFromContextOrNull(
vararg keyNames: String,
classLoader: ClassLoader = Thread.currentThread().contextClassLoader,
wholeClasspath: Boolean = false,
unpackJarCollections: Boolean = false,
): List<File>? {
fun List<File>.takeAndFilter() =
when {
isEmpty() -> null
wholeClasspath -> takeIfContainsAll(*keyNames)
else -> filterIfContainsAll(*keyNames)
}
val fromProperty = System.getProperty(KOTLIN_SCRIPT_CLASSPATH_PROPERTY)?.split(File.pathSeparator)?.map(::File)
if (fromProperty != null) return fromProperty
return classpathFromClassloader(classLoader, unpackJarCollections)?.takeAndFilter()
?: classpathFromClasspathProperty()?.takeAndFilter()
}
fun scriptCompilationClasspathFromContext(
vararg keyNames: String,
classLoader: ClassLoader = Thread.currentThread().contextClassLoader,
wholeClasspath: Boolean = false,
unpackJarCollections: Boolean = false,
): List<File> =
scriptCompilationClasspathFromContextOrNull(
*keyNames,
classLoader = classLoader,
wholeClasspath = wholeClasspath,
unpackJarCollections = unpackJarCollections,
) ?: throw ClasspathExtractionException(
"Unable to get script compilation classpath from context, please specify explicit classpath via \"$KOTLIN_SCRIPT_CLASSPATH_PROPERTY\" property",
)
object KotlinJars {
private val explicitCompilerClasspath: List<File>? by lazy {
System.getProperty(KOTLIN_COMPILER_CLASSPATH_PROPERTY)?.split(File.pathSeparator)?.map(::File)
?: System.getProperty(KOTLIN_COMPILER_JAR_PROPERTY)
?.let(::File)
?.takeIf(File::exists)
?.let { listOf(it) }
}
val compilerClasspath: List<File> by lazy {
findCompilerClasspath(withScripting = false)
}
val compilerWithScriptingClasspath: List<File> by lazy {
findCompilerClasspath(withScripting = true)
}
private fun findCompilerClasspath(withScripting: Boolean): List<File> {
val kotlinCompilerJars = listOf(
KOTLIN_COMPILER_JAR,
KOTLIN_COMPILER_EMBEDDABLE_JAR,
)
val kotlinLibsJars = listOf(
KOTLIN_JAVA_STDLIB_JAR,
KOTLIN_JAVA_REFLECT_JAR,
KOTLIN_JAVA_SCRIPT_RUNTIME_JAR,
TROVE4J_JAR,
)
val kotlinScriptingJars = if (withScripting) {
listOf(
KOTLIN_SCRIPTING_COMPILER_JAR,
KOTLIN_SCRIPTING_COMPILER_EMBEDDABLE_JAR,
KOTLIN_SCRIPTING_COMPILER_IMPL_JAR,
KOTLIN_SCRIPTING_COMPILER_IMPL_EMBEDDABLE_JAR,
KOTLIN_SCRIPTING_COMMON_JAR,
KOTLIN_SCRIPTING_JVM_JAR,
)
} else {
emptyList()
}
val kotlinBaseJars = kotlinCompilerJars + kotlinLibsJars + kotlinScriptingJars
val classpath = explicitCompilerClasspath
// search classpath from context classloader and `java.class.path` property
?: (
classpathFromFQN(
classLoader = Thread.currentThread().contextClassLoader,
fqn = "org.jetbrains.kotlin.cli.jvm.K2JVMCompiler",
)
?: classpathFromClassloader(Thread.currentThread().contextClassLoader)?.takeIf { it.isNotEmpty() }
?: classpathFromClasspathProperty()
)?.filter { f -> kotlinBaseJars.any { f.matchMaybeVersionedFile(it) } }
?.takeIf { it.isNotEmpty() }
// if autodetected, additionally check for presence of the compiler jars
if (classpath == null ||
(
explicitCompilerClasspath == null &&
classpath.none { f ->
kotlinCompilerJars.any { f.matchMaybeVersionedFile(it) }
}
)
) {
throw FileNotFoundException(
"Cannot find kotlin compiler jar, set kotlin.compiler.classpath property to proper location",
)
}
return classpath
}
fun getLib(
propertyName: String,
jarName: String,
markerClass: KClass<*>,
classLoader: ClassLoader? = null,
): File? =
getExplicitLib(propertyName, jarName)
?: run {
val requestedClassloader = classLoader ?: Thread.currentThread().contextClassLoader
val byName =
if (requestedClassloader == markerClass.java.classLoader) {
null
} else {
tryGetResourcePathForClassByName(markerClass.java.name, requestedClassloader)
}
byName ?: tryGetResourcePathForClass(markerClass.java)
}?.takeIf(File::exists)
fun getLib(
propertyName: String,
jarName: String,
markerClassName: String,
classLoader: ClassLoader? = null,
): File? =
getExplicitLib(propertyName, jarName)
?: tryGetResourcePathForClassByName(
markerClassName,
classLoader ?: Thread.currentThread().contextClassLoader,
)?.takeIf(File::exists)
private fun getExplicitLib(propertyName: String, jarName: String) =
System.getProperty(propertyName)?.let(::File)?.takeIf(File::exists)
?: explicitCompilerClasspath?.firstOrNull { it.matchMaybeVersionedFile(jarName) }?.takeIf(File::exists)
val stdlibOrNull: File? by lazy {
System.getProperty(KOTLIN_STDLIB_JAR_PROPERTY)?.let(::File)?.takeIf(File::exists)
?: getLib(
KOTLIN_RUNTIME_JAR_PROPERTY,
KOTLIN_JAVA_STDLIB_JAR,
JvmStatic::class,
)
}
val stdlib: File by lazy {
stdlibOrNull ?: throw Exception(
"Unable to find kotlin stdlib, please specify it explicitly via \"$KOTLIN_STDLIB_JAR_PROPERTY\" property",
)
}
val reflectOrNull: File? by lazy {
getLib(
KOTLIN_REFLECT_JAR_PROPERTY,
KOTLIN_JAVA_REFLECT_JAR,
"kotlin.reflect.full.KClasses", // using a class that is a part of the kotlin-reflect.jar
)
}
}
@@ -0,0 +1,9 @@
package org.jetbrains.kotlinx.dataframe
import org.jetbrains.kotlin.test.builders.TestConfigurationBuilder
fun TestConfigurationBuilder.commonFirWithPluginFrontendConfiguration() {
useConfigurators(
::ExtensionRegistrarConfigurator,
)
}
@@ -0,0 +1,115 @@
/*
* Copyright 2010-2018 JetBrains s.r.o. and Kotlin Programming Language contributors.
* Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
*/
package org.jetbrains.kotlinx.dataframe
import java.io.File
import java.net.JarURLConnection
import java.net.URI
import java.net.URL
// Based on an implementation in com.intellij.openapi.application.PathManager.getResourceRoot
internal fun getResourceRoot(context: Class<*>, path: String): String? {
var url: URL? = context.getResource(path)
if (url == null) {
url = ClassLoader.getSystemResource(path.substring(1))
}
return if (url != null) extractRoot(url, path) else null
}
private const val JAR_PROTOCOL = "jar"
private const val FILE_PROTOCOL = "file"
private const val JAR_SEPARATOR = "!/"
private const val SCHEME_SEPARATOR = "://"
private fun extractRoot(resourceURL: URL, resourcePath: String): String? {
if (!resourcePath.startsWith('/') || resourcePath.startsWith('\\')) return null
var resultPath: String? = null
val protocol = resourceURL.protocol
if (protocol == FILE_PROTOCOL) {
val path = resourceURL.toFileOrNull()!!.path
val testPath = path.replace('\\', '/')
val testResourcePath = resourcePath.replace('\\', '/')
if (testPath.endsWith(testResourcePath, ignoreCase = true)) {
resultPath = path.substring(0, path.length - resourcePath.length)
}
} else if (protocol == JAR_PROTOCOL) {
val paths = splitJarUrl(resourceURL.file)
if (paths?.first != null) {
resultPath = File(paths.first).canonicalPath
}
}
return resultPath?.trimEnd(File.separatorChar)
}
private fun splitJarUrl(url: String): Pair<String, String>? {
val pivot = url.indexOf(JAR_SEPARATOR).takeIf { it >= 0 } ?: return null
val resourcePath = url.substring(pivot + 2)
var jarPath = url.substring(0, pivot)
if (jarPath.startsWith(JAR_PROTOCOL + ":")) {
jarPath = jarPath.substring(JAR_PROTOCOL.length + 1)
}
if (jarPath.startsWith(FILE_PROTOCOL)) {
try {
jarPath = URI(jarPath).toURL().toFileOrNull()!!.path.replace('\\', '/')
} catch (e: Exception) {
jarPath = jarPath.substring(FILE_PROTOCOL.length)
if (jarPath.startsWith(SCHEME_SEPARATOR)) {
jarPath = jarPath.substring(SCHEME_SEPARATOR.length)
} else if (jarPath.startsWith(':')) {
jarPath = jarPath.substring(1)
}
}
}
return Pair(jarPath, resourcePath)
}
fun tryGetResourcePathForClass(aClass: Class<*>): File? {
val path = "/" + aClass.name.replace('.', '/') + ".class"
return getResourceRoot(aClass, path)?.let {
File(it).absoluteFile
}
}
fun getResourcePathForClass(aClass: Class<*>): File =
tryGetResourcePathForClass(aClass)
?: throw IllegalStateException("Resource for class: ${aClass.name} not found")
fun tryGetResourcePathForClassByName(name: String, classLoader: ClassLoader): File? =
try {
classLoader.loadClass(name)?.let(::tryGetResourcePathForClass)
} catch (_: ClassNotFoundException) {
null
} catch (_: NoClassDefFoundError) {
null
}
internal fun URL.toFileOrNull() =
try {
File(toURI())
} catch (e: IllegalArgumentException) {
null
} catch (e: java.net.URISyntaxException) {
null
} ?: run {
if (protocol != "file") {
null
} else {
File(file)
}
}
internal fun URL.toContainingJarOrNull(): File? =
if (protocol == "jar") {
(openConnection() as? JarURLConnection)?.jarFileURL?.toFileOrNull()
} else {
null
}
@@ -0,0 +1,58 @@
package org.jetbrains.kotlinx.dataframe.services
import org.jetbrains.kotlin.test.services.TemporaryDirectoryManager
import org.jetbrains.kotlin.test.services.TestServices
import org.jetbrains.kotlin.test.services.testInfo
import org.jetbrains.kotlin.test.util.KtTestUtil
import java.io.File
import java.nio.file.Paths
import java.util.Locale
import kotlin.io.path.ExperimentalPathApi
import kotlin.io.path.deleteRecursively
// Copied from org.jetbrains.kotlin.test.services.impl.TemporaryDirectoryManagerImpl
// because it uses NioFiles#deleteRecursively and throws method not found as a result.
class TemporaryDirectoryManagerImplFixed(testServices: TestServices) : TemporaryDirectoryManager(testServices) {
private val cache = mutableMapOf<String, File>()
private val rootTempDir = lazy {
val testInfo = testServices.testInfo
val className = testInfo.className
val methodName = testInfo.methodName
if (!onWindows && className.length + methodName.length < 255) {
return@lazy KtTestUtil.tmpDirForTest(className, methodName)
}
// This code will simplify directory name for windows. This is needed because there can occur errors due to long name
val lastDot = className.lastIndexOf('.')
val simplifiedClassName = className.substring(lastDot + 1).getOnlyUpperCaseSymbols()
val simplifiedMethodName = methodName.getOnlyUpperCaseSymbols()
KtTestUtil.tmpDirForTest(simplifiedClassName, simplifiedMethodName)
}
override val rootDir: File
get() = rootTempDir.value
override fun getOrCreateTempDirectory(name: String): File =
cache.getOrPut(name) {
KtTestUtil.tmpDir(rootDir, name)
}
@OptIn(ExperimentalPathApi::class)
override fun cleanupTemporaryDirectories() {
cache.clear()
if (rootTempDir.isInitialized()) {
Paths.get(rootDir.path).deleteRecursively()
}
}
companion object {
private val onWindows: Boolean =
System.getProperty("os.name").lowercase(Locale.getDefault()).contains("windows")
private fun String.getOnlyUpperCaseSymbols(): String =
this.filter { it.isUpperCase() || it == '$' }
.toList()
.joinToString(separator = "")
}
}
+11
View File
@@ -0,0 +1,11 @@
## :plugins:keywords-generator
This module holds a little Gradle plugin whose sole purpose is to provide
[:core](../../core) with the `generateKeywordsSrc` task.
This task, generates three enum classes: `HardKeywords`, `ModifierKeywords`, and `SoftKeywords`.
These enums together contain all restricted Kotlin keywords to be taken into account when generating our own
code in Notebooks or any of our [plugins](..). Words like "package", "fun", "suspend", etc...
As the Kotlin language can change over time, this task ensures that any changes to the language
will be reflected in our code generation.
+28
View File
@@ -0,0 +1,28 @@
@file:OptIn(ExperimentalBuildToolsApi::class, ExperimentalKotlinGradlePluginApi::class)
import org.jetbrains.kotlin.buildtools.api.ExperimentalBuildToolsApi
import org.jetbrains.kotlin.gradle.ExperimentalKotlinGradlePluginApi
plugins {
`java-gradle-plugin`
`kotlin-dsl`
with(convention.plugins) {
alias(kotlinJvm8)
alias(buildConfig)
}
}
dependencies {
compileOnly(kotlin("compiler-embeddable", kotlin.compilerVersion.get()))
implementation(libs.kotlinpoet)
}
gradlePlugin {
plugins {
create("dependencies") {
id = "org.jetbrains.dataframe.generator"
version = "1.0"
implementationClass = "org.jetbrains.dataframe.keywords.KeywordsGeneratorPlugin"
}
}
}
@@ -0,0 +1,9 @@
pluginManagement {
includeBuild("../../build-settings-logic")
}
plugins {
id("dfsettings.catalogs")
}
includeBuild("../../build-logic")
@@ -0,0 +1,6 @@
package org.jetbrains.dataframe.keywords
public data class EnumEntry(
public val name: String,
public val strValue: String
)
@@ -0,0 +1,87 @@
package org.jetbrains.dataframe.keywords
import com.squareup.kotlinpoet.FileSpec
import com.squareup.kotlinpoet.FunSpec
import com.squareup.kotlinpoet.ParameterizedTypeName.Companion.parameterizedBy
import com.squareup.kotlinpoet.PropertySpec
import com.squareup.kotlinpoet.TypeSpec
import org.gradle.workers.WorkAction
import org.gradle.workers.WorkParameters
import org.jetbrains.kotlin.com.intellij.psi.tree.TokenSet
import org.jetbrains.kotlin.config.KotlinCompilerVersion
import org.jetbrains.kotlin.lexer.KtKeywordToken
import org.jetbrains.kotlin.lexer.KtTokens
import java.io.File
public abstract class KeywordsGeneratorAction : WorkAction<KeywordsGeneratorAction.Parameters> {
public interface Parameters : WorkParameters {
public var srcDir: File
}
private val taskPackageName = "org.jetbrains.kotlinx.dataframe.keywords"
override fun execute() {
println("Generating keywords using Kotlin compiler: ${KotlinCompilerVersion.getVersion()}")
parameters.srcDir.deleteRecursively()
generateKeywordEnums()
}
private fun generateKeywordEnums() {
listOf(
"HardKeywords" to KtTokens.KEYWORDS,
"SoftKeywords" to KtTokens.SOFT_KEYWORDS,
"ModifierKeywords" to KtTokens.MODIFIER_KEYWORDS,
).forEach { (name, set) ->
generateKeywordsEnum(name, set)
}
}
private fun generateKeywordsEnum(name: String, tokenSet: TokenSet) {
buildKwEnum(name, getKeywords(tokenSet)).writeTo(parameters.srcDir)
}
private fun getKeywords(tokenSet: TokenSet): List<EnumEntry> {
fun id(value: String) = value.uppercase().replace("!", "NOT_")
return tokenSet.types.map { t ->
t as KtKeywordToken
EnumEntry(id(t.value), t.value)
}
}
private fun buildKwEnum(name: String, values: List<EnumEntry>): FileSpec {
val fileBuilder = FileSpec.builder(taskPackageName, name)
val valList = mutableListOf<String>()
val enumBuilder = TypeSpec.enumBuilder(name).apply {
primaryConstructor(
FunSpec.constructorBuilder()
.addParameter("value", String::class)
.build()
)
values.forEach { entry ->
valList.add("\"${entry.strValue}\"")
addEnumConstant(
entry.name, TypeSpec.anonymousClassBuilder()
.addSuperclassConstructorParameter("%S", entry.strValue)
.build()
)
}
val compObj = TypeSpec.companionObjectBuilder().addProperty(
PropertySpec
.builder("VALUES", List::class.parameterizedBy(String::class))
.initializer(valList.joinToString(", ", "listOf(", ")"))
.build()
).build()
addType(compObj)
}
fileBuilder.addType(enumBuilder.build())
return fileBuilder.build()
}
}
@@ -0,0 +1,48 @@
package org.jetbrains.dataframe.keywords
import org.gradle.api.Plugin
import org.gradle.api.Project
import org.gradle.api.artifacts.DependencyScopeConfiguration
import org.gradle.api.artifacts.ResolvableConfiguration
import org.gradle.api.tasks.SourceSet
import org.gradle.api.tasks.SourceSetContainer
import org.gradle.kotlin.dsl.get
import org.gradle.kotlin.dsl.register
import org.jetbrains.kotlinx.dataframe.BuildConfig
import java.io.File
@Suppress("UnstableApiUsage")
public abstract class KeywordsGeneratorPlugin : Plugin<Project> {
override fun apply(target: Project): Unit = with(target) {
// from https://kotlinlang.org/docs/whatsnew21.html#compiler-symbols-hidden-from-the-kotlin-gradle-plugin-api
val dependencyScopeConfiguration: DependencyScopeConfiguration = configurations.dependencyScope("keywordsGeneratorDependencyScope").get()
dependencies.add(dependencyScopeConfiguration.name, "$KOTLIN_COMPILER_EMBEDDABLE:$KOTLIN_COMPILER_VERSION")
val resolvableConfiguration: ResolvableConfiguration = configurations.resolvable("keywordGeneratorResolvable") {
extendsFrom(dependencyScopeConfiguration)
}.get()
val genSrcDir = layout.buildDirectory.asFile.get().resolve("generatedSrc")
val sourceSets = project.extensions.getByName("sourceSets") as SourceSetContainer
val mainSourceSet = sourceSets.named("main").get()
mainSourceSet.addDir(genSrcDir)
val genTask = tasks.register<KeywordsGeneratorTask>(KeywordsGeneratorTask.NAME) {
kotlinCompiler.from(resolvableConfiguration)
srcDir = genSrcDir
}
tasks["compileKotlin"].dependsOn(genTask)
}
private fun SourceSet.addDir(dir: File) {
java.setSrcDirs(java.srcDirs + dir)
}
public companion object {
public const val KOTLIN_COMPILER_EMBEDDABLE: String = "org.jetbrains.kotlin:kotlin-compiler-embeddable"
public const val KOTLIN_COMPILER_VERSION: String = BuildConfig.KOTLIN_COMPILER_VERSION
}
}
@@ -0,0 +1,41 @@
package org.jetbrains.dataframe.keywords
import org.gradle.api.DefaultTask
import org.gradle.api.file.ConfigurableFileCollection
import org.gradle.api.tasks.Classpath
import org.gradle.api.tasks.Input
import org.gradle.api.tasks.OutputDirectory
import org.gradle.api.tasks.TaskAction
import org.gradle.kotlin.dsl.submit
import org.gradle.workers.WorkerExecutor
import java.io.File
import javax.inject.Inject
public abstract class KeywordsGeneratorTask: DefaultTask() {
@get:Inject
public abstract val executor: WorkerExecutor
@get:Classpath
public abstract val kotlinCompiler: ConfigurableFileCollection
@OutputDirectory
public lateinit var srcDir: File
@Input
override fun getGroup(): String = "codegen"
@TaskAction
public fun generate() {
val workQueue = executor.classLoaderIsolation {
classpath.from(kotlinCompiler)
}
workQueue.submit(KeywordsGeneratorAction::class) {
srcDir = this@KeywordsGeneratorTask.srcDir
}
}
public companion object {
public const val NAME: String = "generateKeywordsSrc"
}
}
+11
View File
@@ -0,0 +1,11 @@
## ~~:plugins:kotlin-dataframe~~
The Kotlin 2.x Compiler plugin of DataFrame.
A plugin for your Kotlin project that can generate on-the-fly column accessors for the compiler and IDE even without
having to provide data schemas!
### DISABLED!
Development of this module was moved to the Kotlin repository:
https://github.com/JetBrains/kotlin/tree/master/plugins/kotlin-dataframe.
These files are out of date.
+121
View File
@@ -0,0 +1,121 @@
import org.jetbrains.kotlin.gradle.dsl.JvmTarget
import org.jetbrains.kotlin.gradle.dsl.KotlinVersion
plugins {
id("java")
kotlin("jvm")
kotlin("plugin.serialization")
}
group = "org.jetbrains.kotlinx.dataframe"
val kotlinVersion: String by project.properties
repositories {
mavenCentral()
maven("https://maven.pkg.jetbrains.space/kotlin/p/kotlin/dev/")
}
sourceSets {
main {
java.setSrcDirs(listOf("src"))
resources.setSrcDirs(listOf("resources"))
}
test {
java.setSrcDirs(listOf("tests", "tests-gen"))
resources.setSrcDirs(listOf("testResources"))
}
}
dependencies {
"org.jetbrains.kotlin:kotlin-compiler:$kotlinVersion".let {
compileOnly(it)
testImplementation(it)
}
testRuntimeOnly("org.jetbrains.kotlin:kotlin-test:$kotlinVersion")
testRuntimeOnly("org.jetbrains.kotlin:kotlin-script-runtime:$kotlinVersion")
testRuntimeOnly("org.jetbrains.kotlin:kotlin-annotations-jvm:$kotlinVersion")
implementation(project(projects.dataframeCompilerPluginCore.path, "shadow"))
testRuntimeOnly(projects.core)
testRuntimeOnly(projects.dataframeCsv)
testImplementation("org.jetbrains.kotlin:kotlin-reflect:$kotlinVersion")
testImplementation("org.jetbrains.kotlin:kotlin-compiler-internal-test-framework:$kotlinVersion")
testImplementation(platform("org.junit:junit-bom:5.11.3"))
testImplementation("org.junit.jupiter:junit-jupiter")
testImplementation("org.junit.platform:junit-platform-commons")
testImplementation("org.junit.platform:junit-platform-launcher")
testImplementation("org.junit.platform:junit-platform-runner")
testImplementation("org.junit.platform:junit-platform-suite-api")
}
tasks.test {
useJUnitPlatform()
jvmArgs("-Xmx2G")
environment("TEST_RESOURCES", project.layout.projectDirectory)
doFirst {
setLibraryProperty("org.jetbrains.kotlin.test.kotlin-stdlib", "kotlin-stdlib")
setLibraryProperty("org.jetbrains.kotlin.test.kotlin-reflect", "kotlin-reflect")
setLibraryProperty("org.jetbrains.kotlin.test.kotlin-test", "kotlin-test")
setLibraryProperty("org.jetbrains.kotlin.test.kotlin-script-runtime", "kotlin-script-runtime")
setLibraryProperty("org.jetbrains.kotlin.test.kotlin-annotations-jvm", "kotlin-annotations-jvm")
}
}
tasks.withType<org.jetbrains.kotlin.gradle.tasks.KotlinCompile>().configureEach {
friendPaths.from(project(projects.core.path).projectDir)
compilerOptions {
freeCompilerArgs.addAll(
"-Xcontext-receivers",
)
optIn.addAll(
"org.jetbrains.kotlin.compiler.plugin.ExperimentalCompilerApi",
)
}
}
tasks.withType<JavaCompile> {
sourceCompatibility = JavaVersion.VERSION_1_8.toString()
targetCompatibility = JavaVersion.VERSION_1_8.toString()
}
tasks.compileKotlin {
compilerOptions {
languageVersion = KotlinVersion.KOTLIN_2_0
jvmTarget = JvmTarget.JVM_1_8
}
}
tasks.compileTestKotlin {
compilerOptions {
languageVersion = KotlinVersion.KOTLIN_2_0
jvmTarget = JvmTarget.JVM_1_8
}
}
tasks.register<JavaExec>("generateTests") {
classpath = sourceSets.test.get().runtimeClasspath
mainClass = "org.jetbrains.kotlin.fir.dataframe.GenerateTestsKt"
}
fun Test.setLibraryProperty(propName: String, jarName: String) {
val path = project.configurations
.testRuntimeClasspath.get()
.files
.find { """$jarName-\d.*jar""".toRegex().matches(it.name) }
?.absolutePath
?: return
systemProperty(propName, path)
}
// Disabling all tests before removing the compiler plugin here
// because we're moving to the Kotlin repo: #1290
tasks.filter {
":plugins:kotlin-dataframe" in it.path &&
"test" in it.name.lowercase()
}.forEach {
println("disabling compiler plugin test task: ${it.path}. See #1290")
it.onlyIf { false }
}
+2
View File
@@ -0,0 +1,2 @@
kotlin.code.style=official
kotlinVersion=2.0.20
@@ -0,0 +1,5 @@
distributionBase=GRADLE_USER_HOME
distributionPath=wrapper/dists
distributionUrl=https\://services.gradle.org/distributions/gradle-8.10-bin.zip
zipStoreBase=GRADLE_USER_HOME
zipStorePath=wrapper/dists
@@ -0,0 +1,6 @@
#
# Copyright 2010-2020 JetBrains s.r.o. and Kotlin Programming Language contributors.
# Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
#
org.jetbrains.kotlinx.dataframe.plugin.FirDataFrameComponentRegistrar
@@ -0,0 +1,53 @@
/*
* Copyright 2010-2021 JetBrains s.r.o. and Kotlin Programming Language contributors.
* Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
*/
package org.jetbrains.kotlinx.dataframe.plugin
import org.jetbrains.kotlin.backend.common.extensions.IrGenerationExtension
import org.jetbrains.kotlin.compiler.plugin.CompilerPluginRegistrar
import org.jetbrains.kotlin.compiler.plugin.ExperimentalCompilerApi
import org.jetbrains.kotlin.config.CompilerConfiguration
import org.jetbrains.kotlin.fir.FirSession
import org.jetbrains.kotlin.fir.extensions.FirExtensionApiInternals
import org.jetbrains.kotlin.fir.extensions.FirExtensionRegistrar
import org.jetbrains.kotlin.fir.extensions.FirExtensionRegistrarAdapter
import org.jetbrains.kotlinx.dataframe.plugin.extensions.DataRowSchemaSupertype
import org.jetbrains.kotlinx.dataframe.plugin.extensions.ExpressionAnalysisAdditionalChecker
import org.jetbrains.kotlinx.dataframe.plugin.extensions.FunctionCallTransformer
import org.jetbrains.kotlinx.dataframe.plugin.extensions.IrBodyFiller
import org.jetbrains.kotlinx.dataframe.plugin.extensions.ReturnTypeBasedReceiverInjector
import org.jetbrains.kotlinx.dataframe.plugin.extensions.TokenGenerator
import org.jetbrains.kotlinx.dataframe.plugin.extensions.TopLevelExtensionsGenerator
class FirDataFrameExtensionRegistrar(
val isTest: Boolean,
val dumpSchemas: Boolean,
) : FirExtensionRegistrar() {
@OptIn(FirExtensionApiInternals::class)
override fun ExtensionRegistrarContext.configurePlugin() {
+::TopLevelExtensionsGenerator
+::ReturnTypeBasedReceiverInjector
+{ it: FirSession ->
FunctionCallTransformer(it, isTest)
}
+::TokenGenerator
+::DataRowSchemaSupertype
+{ it: FirSession ->
ExpressionAnalysisAdditionalChecker(it, isTest, dumpSchemas)
}
}
}
@OptIn(ExperimentalCompilerApi::class)
class FirDataFrameComponentRegistrar : CompilerPluginRegistrar() {
override fun ExtensionStorage.registerExtensions(configuration: CompilerConfiguration) {
FirExtensionRegistrarAdapter.registerExtension(
FirDataFrameExtensionRegistrar(isTest = false, dumpSchemas = true)
)
IrGenerationExtension.registerExtension(IrBodyFiller())
}
override val supportsK2: Boolean = true
}
@@ -0,0 +1,63 @@
/*
* Copyright 2010-2022 JetBrains s.r.o. and Kotlin Programming Language contributors.
* Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
*/
package org.jetbrains.kotlinx.dataframe.plugin
import org.jetbrains.kotlin.fir.expressions.FirExpression
import org.jetbrains.kotlin.fir.expressions.FirFunctionCall
import org.jetbrains.kotlin.fir.types.ConeClassLikeType
import org.jetbrains.kotlin.fir.types.classId
import org.jetbrains.kotlin.fir.types.resolvedType
import org.jetbrains.kotlin.name.ClassId
import org.jetbrains.kotlin.name.Name
import org.jetbrains.kotlinx.dataframe.plugin.extensions.KotlinTypeFacade
import org.jetbrains.kotlinx.dataframe.plugin.impl.Interpreter
internal inline fun <reified T> KotlinTypeFacade.analyzeRefinedCallShape(
call: FirFunctionCall,
expectedReturnType: ClassId,
reporter: InterpretationErrorReporter
): CallResult<T>? {
val callReturnType = call.resolvedType
if (callReturnType.classId != expectedReturnType) return null
// rootMarker is expected to be a token generated by the plugin.
// it's implied by "refined call"
// thus ConeClassLikeType
val rootMarkers = callReturnType.typeArguments.filterIsInstance<ConeClassLikeType>()
if (rootMarkers.size != callReturnType.typeArguments.size) return null
val newSchema: T? = call.interpreterName(session)?.let { name ->
when (name) {
else -> name.load<Interpreter<*>>().let { processor ->
val dataFrameSchema = interpret(call, processor, reporter = reporter)
.let {
val value = it?.value
if (value !is T) {
if (!reporter.errorReported) {
reporter.reportInterpretationError(call, "${processor::class} must return ${T::class}, but was $value")
}
null
} else {
value
}
}
dataFrameSchema
}
}
}
return CallResult(rootMarkers, newSchema)
}
data class CallResult<T>(val markers: List<ConeClassLikeType>, val result: T?)
class RefinedArguments(val refinedArguments: List<RefinedArgument>) : List<RefinedArgument> by refinedArguments
data class RefinedArgument(val name: Name, val expression: FirExpression) {
override fun toString(): String {
return "RefinedArgument(name=$name, expression=${expression})"
}
}
@@ -0,0 +1,21 @@
package org.jetbrains.kotlinx.dataframe.plugin.extensions
import org.jetbrains.kotlin.KtSourceElement
import org.jetbrains.kotlinx.dataframe.plugin.extensions.impl.SchemaProperty
import org.jetbrains.kotlin.fir.declarations.FirClass
import org.jetbrains.kotlin.fir.declarations.FirDeclarationDataKey
import org.jetbrains.kotlin.fir.declarations.FirDeclarationDataRegistry
import org.jetbrains.kotlin.fir.symbols.impl.FirRegularClassSymbol
sealed interface CallShapeData {
class Schema(val columns: List<SchemaProperty>) : CallShapeData
class Scope(val columns: List<SchemaProperty>, val source: KtSourceElement?) : CallShapeData
class RefinedType(val scopes: List<FirRegularClassSymbol>) : CallShapeData
}
object CallShapeAttribute : FirDeclarationDataKey()
var FirClass.callShapeData: CallShapeData? by FirDeclarationDataRegistry.data(CallShapeAttribute)
@@ -0,0 +1,5 @@
package org.jetbrains.kotlinx.dataframe.plugin.extensions
import org.jetbrains.kotlin.GeneratedDeclarationKey
data object DataFramePlugin : GeneratedDeclarationKey()
@@ -0,0 +1,52 @@
package org.jetbrains.kotlinx.dataframe.plugin.extensions
import org.jetbrains.kotlin.descriptors.ClassKind
import org.jetbrains.kotlin.fir.FirSession
import org.jetbrains.kotlin.fir.analysis.checkers.toClassLikeSymbol
import org.jetbrains.kotlin.fir.declarations.FirClassLikeDeclaration
import org.jetbrains.kotlin.fir.declarations.FirRegularClass
import org.jetbrains.kotlin.fir.extensions.AnnotationFqn
import org.jetbrains.kotlin.fir.extensions.FirDeclarationPredicateRegistrar
import org.jetbrains.kotlin.fir.extensions.FirSupertypeGenerationExtension
import org.jetbrains.kotlin.fir.extensions.predicate.LookupPredicate
import org.jetbrains.kotlin.fir.extensions.predicateBasedProvider
import org.jetbrains.kotlin.fir.types.FirResolvedTypeRef
import org.jetbrains.kotlin.fir.types.builder.buildResolvedTypeRef
import org.jetbrains.kotlin.fir.types.constructClassLikeType
import org.jetbrains.kotlin.name.ClassId
import org.jetbrains.kotlin.name.FqName
import org.jetbrains.kotlin.name.Name
import org.jetbrains.kotlinx.dataframe.annotations.DataSchema
class DataRowSchemaSupertype(session: FirSession) : FirSupertypeGenerationExtension(session) {
companion object {
private val PREDICATE = LookupPredicate.create {
annotated(AnnotationFqn(DataSchema::class.java.name))
}
private val dataRowSchema = ClassId(FqName("org.jetbrains.kotlinx.dataframe.api"), Name.identifier("DataRowSchema"))
}
override fun FirDeclarationPredicateRegistrar.registerPredicates() {
register(PREDICATE)
}
override fun needTransformSupertypes(declaration: FirClassLikeDeclaration): Boolean {
return declaration is FirRegularClass
&& declaration.classKind == ClassKind.CLASS
&& session.predicateBasedProvider.matches(PREDICATE, declaration)
}
override fun computeAdditionalSupertypes(
classLikeDeclaration: FirClassLikeDeclaration,
resolvedSupertypes: List<FirResolvedTypeRef>,
typeResolver: TypeResolveService
): List<FirResolvedTypeRef> {
if (resolvedSupertypes.any { it.toClassLikeSymbol(session)?.classId == dataRowSchema }) return emptyList()
return listOf(
buildResolvedTypeRef {
type = dataRowSchema.constructClassLikeType(emptyArray())
}
)
}
}
@@ -0,0 +1,259 @@
/*
* Copyright 2010-2022 JetBrains s.r.o. and Kotlin Programming Language contributors.
* Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
*/
package org.jetbrains.kotlinx.dataframe.plugin.extensions
import com.intellij.psi.PsiElement
import org.jetbrains.kotlin.KtSourceElement
import org.jetbrains.kotlin.diagnostics.AbstractSourceElementPositioningStrategy
import org.jetbrains.kotlin.diagnostics.DiagnosticFactory1DelegateProvider
import org.jetbrains.kotlin.diagnostics.DiagnosticReporter
import org.jetbrains.kotlin.diagnostics.KtDiagnosticFactory1
import org.jetbrains.kotlin.diagnostics.Severity
import org.jetbrains.kotlin.diagnostics.SourceElementPositioningStrategies
import org.jetbrains.kotlin.diagnostics.error1
import org.jetbrains.kotlin.diagnostics.reportOn
import org.jetbrains.kotlin.diagnostics.warning1
import org.jetbrains.kotlin.fir.FirSession
import org.jetbrains.kotlin.fir.analysis.checkers.MppCheckerKind
import org.jetbrains.kotlin.fir.analysis.checkers.context.CheckerContext
import org.jetbrains.kotlin.fir.analysis.checkers.declaration.DeclarationCheckers
import org.jetbrains.kotlin.fir.analysis.checkers.declaration.FirPropertyChecker
import org.jetbrains.kotlin.fir.analysis.checkers.declaration.FirSimpleFunctionChecker
import org.jetbrains.kotlin.fir.analysis.checkers.expression.ExpressionCheckers
import org.jetbrains.kotlin.fir.analysis.checkers.expression.FirFunctionCallChecker
import org.jetbrains.kotlin.fir.analysis.checkers.expression.FirPropertyAccessExpressionChecker
import org.jetbrains.kotlin.fir.analysis.extensions.FirAdditionalCheckersExtension
import org.jetbrains.kotlin.fir.declarations.FirProperty
import org.jetbrains.kotlin.fir.declarations.FirSimpleFunction
import org.jetbrains.kotlin.fir.declarations.hasAnnotation
import org.jetbrains.kotlin.fir.expressions.FirFunctionCall
import org.jetbrains.kotlin.fir.expressions.FirPropertyAccessExpression
import org.jetbrains.kotlin.fir.references.FirResolvedNamedReference
import org.jetbrains.kotlin.fir.references.toResolvedCallableSymbol
import org.jetbrains.kotlin.fir.resolve.fullyExpandedType
import org.jetbrains.kotlin.fir.types.ConeClassLikeType
import org.jetbrains.kotlin.fir.types.ConeKotlinType
import org.jetbrains.kotlin.fir.types.FirTypeProjectionWithVariance
import org.jetbrains.kotlin.fir.types.coneType
import org.jetbrains.kotlin.fir.types.isSubtypeOf
import org.jetbrains.kotlin.fir.types.renderReadable
import org.jetbrains.kotlin.fir.types.resolvedType
import org.jetbrains.kotlin.fir.types.toSymbol
import org.jetbrains.kotlin.fir.types.type
import org.jetbrains.kotlin.name.CallableId
import org.jetbrains.kotlin.name.ClassId
import org.jetbrains.kotlin.name.FqName
import org.jetbrains.kotlin.name.Name
import org.jetbrains.kotlin.psi.KtElement
import org.jetbrains.kotlinx.dataframe.plugin.impl.PluginDataFrameSchema
import org.jetbrains.kotlinx.dataframe.plugin.impl.SimpleColumnGroup
import org.jetbrains.kotlinx.dataframe.plugin.impl.SimpleDataColumn
import org.jetbrains.kotlinx.dataframe.plugin.impl.SimpleFrameColumn
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.flatten
import org.jetbrains.kotlinx.dataframe.plugin.pluginDataFrameSchema
import org.jetbrains.kotlinx.dataframe.plugin.utils.Names
import org.jetbrains.kotlinx.dataframe.plugin.utils.isDataFrame
import org.jetbrains.kotlinx.dataframe.plugin.utils.isDataRow
import org.jetbrains.kotlinx.dataframe.plugin.utils.isGroupBy
class ExpressionAnalysisAdditionalChecker(
session: FirSession,
isTest: Boolean,
dumpSchemas: Boolean
) : FirAdditionalCheckersExtension(session) {
override val expressionCheckers: ExpressionCheckers = object : ExpressionCheckers() {
override val functionCallCheckers: Set<FirFunctionCallChecker> = setOfNotNull(
Checker(isTest), FunctionCallSchemaReporter.takeIf { dumpSchemas }
)
override val propertyAccessExpressionCheckers: Set<FirPropertyAccessExpressionChecker> = setOfNotNull(
PropertyAccessSchemaReporter.takeIf { dumpSchemas }
)
}
override val declarationCheckers: DeclarationCheckers = object : DeclarationCheckers() {
override val propertyCheckers: Set<FirPropertyChecker> = setOfNotNull(PropertySchemaReporter.takeIf { dumpSchemas })
override val simpleFunctionCheckers: Set<FirSimpleFunctionChecker> = setOfNotNull(FunctionDeclarationSchemaReporter.takeIf { dumpSchemas })
}
}
private class Checker(
val isTest: Boolean,
) : FirFunctionCallChecker(mppKind = MppCheckerKind.Common) {
companion object {
val ERROR by error1<KtElement, String>(SourceElementPositioningStrategies.DEFAULT)
val CAST_ERROR by error1<KtElement, String>(SourceElementPositioningStrategies.CALL_ELEMENT_WITH_DOT)
val CAST_TARGET_WARNING by warning1<KtElement, String>(SourceElementPositioningStrategies.CALL_ELEMENT_WITH_DOT)
val CAST_ID = CallableId(FqName.fromSegments(listOf("org", "jetbrains", "kotlinx", "dataframe", "api")), Name.identifier("cast"))
val CHECK = ClassId(FqName("org.jetbrains.kotlinx.dataframe.annotations"), Name.identifier("Check"))
}
override fun check(expression: FirFunctionCall, context: CheckerContext, reporter: DiagnosticReporter) {
with(KotlinTypeFacadeImpl(context.session, isTest)) {
analyzeCast(expression, reporter, context)
// analyzeRefinedCallShape(expression, reporter = object : InterpretationErrorReporter {
// override var errorReported: Boolean = false
//
// override fun reportInterpretationError(call: FirFunctionCall, message: String) {
// reporter.reportOn(call.source, ERROR, message, context)
// errorReported = true
// }
//
// override fun doNotReportInterpretationError() {
// errorReported = true
// }
// })
}
}
private fun KotlinTypeFacadeImpl.analyzeCast(expression: FirFunctionCall, reporter: DiagnosticReporter, context: CheckerContext) {
val calleeReference = expression.calleeReference
if (calleeReference !is FirResolvedNamedReference
|| calleeReference.toResolvedCallableSymbol()?.callableId != CAST_ID
|| !calleeReference.resolvedSymbol.hasAnnotation(CHECK, session)) {
return
}
val targetProjection = expression.typeArguments.getOrNull(0) as? FirTypeProjectionWithVariance ?: return
val targetType = targetProjection.typeRef.coneType as? ConeClassLikeType ?: return
val targetSymbol = targetType.toSymbol(session)
if (targetSymbol != null && !targetSymbol.hasAnnotation(Names.DATA_SCHEMA_CLASS_ID, session)) {
val text = "Annotate ${targetType.renderReadable()} with @DataSchema to use generated properties"
reporter.reportOn(expression.source, CAST_TARGET_WARNING, text, context)
}
val coneType = expression.explicitReceiver?.resolvedType
if (coneType != null) {
val sourceType = coneType.fullyExpandedType(session).typeArguments.getOrNull(0)?.type as? ConeClassLikeType
?: return
val source = pluginDataFrameSchema(sourceType)
val target = pluginDataFrameSchema(targetType)
val sourceColumns = source.flatten(includeFrames = true)
val targetColumns = target.flatten(includeFrames = true)
val sourceMap = sourceColumns.associate { it.path.path to it.column }
val missingColumns = mutableListOf<String>()
var valid = true
for (target in targetColumns) {
val source = sourceMap[target.path.path]
val present = if (source != null) {
if (source !is SimpleDataColumn || target.column !is SimpleDataColumn) { continue }
if (source.type.type().isSubtypeOf(target.column.type.type(), session)) {
true
} else {
missingColumns += "${target.path.path} ${target.column.name}: ${source.type.type().renderReadable()} is not subtype of ${target.column.type.type()}"
false
}
} else {
missingColumns += "${target.path.path} ${target.column.name} is missing"
false
}
valid = valid && present
}
if (!valid) {
reporter.reportOn(expression.source, CAST_ERROR, "Cast cannot succeed \n ${missingColumns.joinToString("\n")}", context)
}
}
}
}
private data object PropertySchemaReporter : FirPropertyChecker(mppKind = MppCheckerKind.Common) {
val SCHEMA by info1<KtElement, String>(SourceElementPositioningStrategies.DECLARATION_NAME)
override fun check(declaration: FirProperty, context: CheckerContext, reporter: DiagnosticReporter) {
context.sessionContext {
declaration.returnTypeRef.coneType.let { type ->
reportSchema(reporter, declaration.source, SCHEMA, type, context)
}
}
}
}
private data object FunctionCallSchemaReporter : FirFunctionCallChecker(mppKind = MppCheckerKind.Common) {
val SCHEMA by info1<KtElement, String>(SourceElementPositioningStrategies.REFERENCED_NAME_BY_QUALIFIED)
override fun check(expression: FirFunctionCall, context: CheckerContext, reporter: DiagnosticReporter) {
if (expression.calleeReference.name in setOf(Name.identifier("let"), Name.identifier("run"))) return
val initializer = expression.resolvedType
context.sessionContext {
reportSchema(reporter, expression.source, SCHEMA, initializer, context)
}
}
}
private data object PropertyAccessSchemaReporter : FirPropertyAccessExpressionChecker(mppKind = MppCheckerKind.Common) {
val SCHEMA by info1<KtElement, String>(SourceElementPositioningStrategies.REFERENCED_NAME_BY_QUALIFIED)
override fun check(
expression: FirPropertyAccessExpression,
context: CheckerContext,
reporter: DiagnosticReporter
) {
val initializer = expression.resolvedType
context.sessionContext {
reportSchema(reporter, expression.source, SCHEMA, initializer, context)
}
}
}
private data object FunctionDeclarationSchemaReporter : FirSimpleFunctionChecker(mppKind = MppCheckerKind.Common) {
val SCHEMA by info1<KtElement, String>(SourceElementPositioningStrategies.DECLARATION_SIGNATURE)
override fun check(declaration: FirSimpleFunction, context: CheckerContext, reporter: DiagnosticReporter) {
val type = declaration.returnTypeRef.coneType
context.sessionContext {
reportSchema(reporter, declaration.source, SCHEMA, type, context)
}
}
}
private fun SessionContext.reportSchema(
reporter: DiagnosticReporter,
source: KtSourceElement?,
factory: KtDiagnosticFactory1<String>,
type: ConeKotlinType,
context: CheckerContext,
) {
val expandedType = type.fullyExpandedType(session)
var schema: PluginDataFrameSchema? = null
when {
expandedType.isDataFrame(session) -> {
schema = expandedType.typeArguments.getOrNull(0)?.let {
pluginDataFrameSchema(it)
}
}
expandedType.isDataRow(session) -> {
schema = expandedType.typeArguments.getOrNull(0)?.let {
pluginDataFrameSchema(it)
}
}
expandedType.isGroupBy(session) -> {
val keys = expandedType.typeArguments.getOrNull(0)
val grouped = expandedType.typeArguments.getOrNull(1)
if (keys != null && grouped != null) {
val keysSchema = pluginDataFrameSchema(keys)
val groupedSchema = pluginDataFrameSchema(grouped)
schema = PluginDataFrameSchema(
listOf(
SimpleColumnGroup("keys", keysSchema.columns()),
SimpleFrameColumn("groups", groupedSchema.columns())
)
)
}
}
}
if (schema != null && source != null) {
reporter.reportOn(source, factory, "\n" + schema.toString(), context)
}
}
fun CheckerContext.sessionContext(f: SessionContext.() -> Unit) {
SessionContext(session).f()
}
inline fun <reified P : PsiElement, A> info1(
positioningStrategy: AbstractSourceElementPositioningStrategy = SourceElementPositioningStrategies.DEFAULT
): DiagnosticFactory1DelegateProvider<A> {
return DiagnosticFactory1DelegateProvider(Severity.INFO, positioningStrategy, P::class)
}
@@ -0,0 +1,601 @@
package org.jetbrains.kotlinx.dataframe.plugin.extensions
import org.jetbrains.kotlin.cli.common.repl.replEscapeLineBreaks
import org.jetbrains.kotlin.contracts.description.EventOccurrencesRange
import org.jetbrains.kotlin.descriptors.ClassKind
import org.jetbrains.kotlin.descriptors.EffectiveVisibility
import org.jetbrains.kotlin.descriptors.Modality
import org.jetbrains.kotlin.descriptors.Visibilities
import org.jetbrains.kotlin.fir.FirAnnotationContainer
import org.jetbrains.kotlin.fir.FirElement
import org.jetbrains.kotlin.fir.FirFunctionTarget
import org.jetbrains.kotlin.fir.FirSession
import org.jetbrains.kotlin.fir.analysis.checkers.fullyExpandedClassId
import org.jetbrains.kotlinx.dataframe.plugin.InterpretationErrorReporter
import org.jetbrains.kotlinx.dataframe.plugin.extensions.impl.SchemaProperty
import org.jetbrains.kotlinx.dataframe.plugin.analyzeRefinedCallShape
import org.jetbrains.kotlinx.dataframe.plugin.utils.Names
import org.jetbrains.kotlinx.dataframe.plugin.utils.projectOverDataColumnType
import org.jetbrains.kotlin.fir.declarations.EmptyDeprecationsProvider
import org.jetbrains.kotlin.fir.declarations.FirClass
import org.jetbrains.kotlin.fir.declarations.FirDeclarationOrigin
import org.jetbrains.kotlin.fir.declarations.FirRegularClass
import org.jetbrains.kotlin.fir.declarations.FirResolvePhase
import org.jetbrains.kotlin.fir.declarations.InlineStatus
import org.jetbrains.kotlin.fir.declarations.builder.buildAnonymousFunction
import org.jetbrains.kotlin.fir.declarations.builder.buildRegularClass
import org.jetbrains.kotlin.fir.declarations.builder.buildValueParameter
import org.jetbrains.kotlin.fir.declarations.impl.FirResolvedDeclarationStatusImpl
import org.jetbrains.kotlin.fir.declarations.utils.classId
import org.jetbrains.kotlin.fir.expressions.FirFunctionCall
import org.jetbrains.kotlin.fir.expressions.FirLiteralExpression
import org.jetbrains.kotlin.fir.expressions.buildResolvedArgumentList
import org.jetbrains.kotlin.fir.expressions.builder.buildAnonymousFunctionExpression
import org.jetbrains.kotlin.fir.expressions.builder.buildBlock
import org.jetbrains.kotlin.fir.expressions.builder.buildFunctionCall
import org.jetbrains.kotlin.fir.expressions.builder.buildPropertyAccessExpression
import org.jetbrains.kotlin.fir.expressions.builder.buildReturnExpression
import org.jetbrains.kotlin.fir.extensions.FirExtensionApiInternals
import org.jetbrains.kotlin.fir.extensions.FirFunctionCallRefinementExtension
import org.jetbrains.kotlin.fir.moduleData
import org.jetbrains.kotlin.fir.references.FirResolvedNamedReference
import org.jetbrains.kotlin.fir.references.builder.buildResolvedNamedReference
import org.jetbrains.kotlin.fir.resolve.calls.candidate.CallInfo
import org.jetbrains.kotlin.fir.resolve.defaultType
import org.jetbrains.kotlin.fir.resolve.fqName
import org.jetbrains.kotlin.fir.resolve.providers.symbolProvider
import org.jetbrains.kotlin.fir.scopes.FirKotlinScopeProvider
import org.jetbrains.kotlin.fir.symbols.SymbolInternals
import org.jetbrains.kotlin.fir.symbols.impl.ConeClassLikeLookupTagImpl
import org.jetbrains.kotlin.fir.symbols.impl.ConeClassLookupTagWithFixedSymbol
import org.jetbrains.kotlin.fir.symbols.impl.FirAnonymousFunctionSymbol
import org.jetbrains.kotlin.fir.symbols.impl.FirFunctionSymbol
import org.jetbrains.kotlin.fir.symbols.impl.FirNamedFunctionSymbol
import org.jetbrains.kotlin.fir.symbols.impl.FirPropertySymbol
import org.jetbrains.kotlin.fir.symbols.impl.FirRegularClassSymbol
import org.jetbrains.kotlin.fir.symbols.impl.FirValueParameterSymbol
import org.jetbrains.kotlin.fir.toFirResolvedTypeRef
import org.jetbrains.kotlin.fir.types.ConeKotlinTypeProjection
import org.jetbrains.kotlin.fir.types.ConeStarProjection
import org.jetbrains.kotlin.fir.types.ConeTypeProjection
import org.jetbrains.kotlin.fir.types.builder.buildResolvedTypeRef
import org.jetbrains.kotlin.fir.types.builder.buildTypeProjectionWithVariance
import org.jetbrains.kotlin.fir.types.classId
import org.jetbrains.kotlin.fir.types.impl.ConeClassLikeTypeImpl
import org.jetbrains.kotlin.fir.types.impl.FirImplicitAnyTypeRef
import org.jetbrains.kotlin.fir.types.resolvedType
import org.jetbrains.kotlin.fir.types.toClassSymbol
import org.jetbrains.kotlin.fir.types.toRegularClassSymbol
import org.jetbrains.kotlin.fir.visitors.FirTransformer
import org.jetbrains.kotlin.name.CallableId
import org.jetbrains.kotlin.name.ClassId
import org.jetbrains.kotlin.name.FqName
import org.jetbrains.kotlin.name.Name
import org.jetbrains.kotlin.text
import org.jetbrains.kotlin.types.Variance
import org.jetbrains.kotlinx.dataframe.plugin.extensions.impl.PropertyName
import org.jetbrains.kotlinx.dataframe.plugin.impl.PluginDataFrameSchema
import org.jetbrains.kotlinx.dataframe.plugin.impl.SimpleCol
import org.jetbrains.kotlinx.dataframe.plugin.impl.SimpleDataColumn
import org.jetbrains.kotlinx.dataframe.plugin.impl.SimpleColumnGroup
import org.jetbrains.kotlinx.dataframe.plugin.impl.SimpleFrameColumn
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.GroupBy
import kotlin.math.abs
@OptIn(FirExtensionApiInternals::class)
class FunctionCallTransformer(
session: FirSession,
override val isTest: Boolean,
) : FirFunctionCallRefinementExtension(session), KotlinTypeFacade {
companion object {
const val DEFAULT_NAME = "DataFrameType"
}
private interface CallTransformer {
fun interceptOrNull(callInfo: CallInfo, symbol: FirNamedFunctionSymbol, hash: String): CallReturnType?
/**
* must still generate let with declared class from interceptOrNull when interpretation fails.
* it should only return null if later some frontend checker fails compilation in general
*/
fun transformOrNull(call: FirFunctionCall, originalSymbol: FirNamedFunctionSymbol): FirFunctionCall?
}
// also update [ReturnTypeBasedReceiverInjector.SCHEMA_TYPES]
private val transformers = listOf(
GroupByCallTransformer(),
DataFrameCallTransformer(),
DataRowCallTransformer(),
ColumnGroupCallTransformer(),
)
override fun intercept(callInfo: CallInfo, symbol: FirNamedFunctionSymbol): CallReturnType? {
val callSiteAnnotations = (callInfo.callSite as? FirAnnotationContainer)?.annotations ?: emptyList()
if (callSiteAnnotations.any { it.fqName(session)?.shortName()?.equals(Name.identifier("DisableInterpretation")) == true }) {
return null
}
val noRefineAnnotation =
symbol.annotations.none { it.fqName(session)?.shortName()?.equals(Name.identifier("Refine")) == true }
val optIn = symbol.annotations.any { it.fqName(session)?.shortName()?.equals(Name.identifier("OptInRefine")) == true } &&
callSiteAnnotations.any { it.fqName(session)?.shortName()?.equals(Name.identifier("Import")) == true }
if (noRefineAnnotation && !optIn) {
return null
}
if (exposesLocalType(callInfo)) return null
val hash = run {
val hash = callInfo.name.hashCode() + callInfo.arguments.sumOf {
when (it) {
is FirLiteralExpression -> it.value.hashCode()
else -> it.source?.text?.hashCode() ?: 42
}
}
hashToTwoCharString(abs(hash))
}
return transformers.firstNotNullOfOrNull { it.interceptOrNull(callInfo, symbol, hash) }
}
private fun exposesLocalType(callInfo: CallInfo): Boolean {
val property = callInfo.containingDeclarations.lastOrNull()?.symbol as? FirPropertySymbol
return (property != null && !property.resolvedStatus.effectiveVisibility.privateApi)
}
private fun hashToTwoCharString(hash: Int): String {
val baseChars = "0123456789"
val base = baseChars.length
val positiveHash = abs(hash)
val char1 = baseChars[positiveHash % base]
val char2 = baseChars[(positiveHash / base) % base]
return "$char1$char2"
}
override fun transform(call: FirFunctionCall, originalSymbol: FirNamedFunctionSymbol): FirFunctionCall {
return transformers
.firstNotNullOfOrNull { it.transformOrNull(call, originalSymbol) }
?: call
}
inner class DataSchemaLikeCallTransformer(val classId: ClassId) : CallTransformer {
override fun interceptOrNull(callInfo: CallInfo, symbol: FirNamedFunctionSymbol, hash: String): CallReturnType? {
if (symbol.resolvedReturnType.fullyExpandedClassId(session) != classId) return null
// possibly null if explicit receiver type is typealias
val argument = (callInfo.explicitReceiver?.resolvedType)?.typeArguments?.getOrNull(0)
val newDataFrameArgument = buildNewTypeArgument(argument, callInfo.name, hash)
val lookupTag = ConeClassLikeLookupTagImpl(classId)
val typeRef = buildResolvedTypeRef {
type = ConeClassLikeTypeImpl(
lookupTag,
arrayOf(
ConeClassLikeTypeImpl(
ConeClassLookupTagWithFixedSymbol(newDataFrameArgument.classId, newDataFrameArgument.symbol),
emptyArray(),
isNullable = false
)
),
isNullable = false
)
}
return CallReturnType(typeRef)
}
@OptIn(SymbolInternals::class)
override fun transformOrNull(call: FirFunctionCall, originalSymbol: FirNamedFunctionSymbol): FirFunctionCall? {
val callResult = analyzeRefinedCallShape<PluginDataFrameSchema>(call, classId, InterpretationErrorReporter.DEFAULT)
val (tokens, dataFrameSchema) = callResult ?: return null
val token = tokens[0]
val firstSchema = token.toClassSymbol(session)?.resolvedSuperTypes?.get(0)!!.toRegularClassSymbol(session)?.fir!!
val dataSchemaApis = materialize(dataFrameSchema ?: PluginDataFrameSchema.EMPTY, call, firstSchema)
val tokenFir = token.toClassSymbol(session)!!.fir
tokenFir.callShapeData = CallShapeData.RefinedType(dataSchemaApis.map { it.scope.symbol })
return buildScopeFunctionCall(call, originalSymbol, dataSchemaApis, listOf(tokenFir))
}
}
inner class DataFrameCallTransformer : CallTransformer by DataSchemaLikeCallTransformer(Names.DF_CLASS_ID)
inner class DataRowCallTransformer : CallTransformer by DataSchemaLikeCallTransformer(Names.DATA_ROW_CLASS_ID)
inner class ColumnGroupCallTransformer : CallTransformer by DataSchemaLikeCallTransformer(Names.COLUM_GROUP_CLASS_ID)
inner class GroupByCallTransformer : CallTransformer {
override fun interceptOrNull(
callInfo: CallInfo,
symbol: FirNamedFunctionSymbol,
hash: String
): CallReturnType? {
if (symbol.resolvedReturnType.fullyExpandedClassId(session) != Names.GROUP_BY_CLASS_ID) return null
val keys = buildNewTypeArgument(null, Name.identifier("Key"), hash)
val group = buildNewTypeArgument(null, Name.identifier("Group"), hash)
val lookupTag = ConeClassLikeLookupTagImpl(Names.GROUP_BY_CLASS_ID)
val typeRef = buildResolvedTypeRef {
type = ConeClassLikeTypeImpl(
lookupTag,
arrayOf(
ConeClassLikeTypeImpl(
ConeClassLookupTagWithFixedSymbol(keys.classId, keys.symbol),
emptyArray<ConeTypeProjection>(),
isNullable = false
),
ConeClassLikeTypeImpl(
ConeClassLookupTagWithFixedSymbol(group.classId, group.symbol),
emptyArray<ConeTypeProjection>(),
isNullable = false
)
),
isNullable = false
)
}
return CallReturnType(typeRef)
}
@OptIn(SymbolInternals::class)
override fun transformOrNull(call: FirFunctionCall, originalSymbol: FirNamedFunctionSymbol): FirFunctionCall? {
val callResult = analyzeRefinedCallShape<GroupBy>(call, Names.GROUP_BY_CLASS_ID, InterpretationErrorReporter.DEFAULT)
val (rootMarkers, groupBy) = callResult ?: return null
val keyMarker = rootMarkers[0]
val groupMarker = rootMarkers[1]
val (keySchema, groupSchema) = if (groupBy != null) {
val keySchema = groupBy.keys
val groupSchema = groupBy.groups
keySchema to groupSchema
} else {
PluginDataFrameSchema.EMPTY to PluginDataFrameSchema.EMPTY
}
val firstSchema = keyMarker.toClassSymbol(session)?.resolvedSuperTypes?.get(0)!!.toRegularClassSymbol(session)?.fir!!
val firstSchema1 = groupMarker.toClassSymbol(session)?.resolvedSuperTypes?.get(0)!!.toRegularClassSymbol(session)?.fir!!
val keyApis = materialize(keySchema, call, firstSchema, "Key")
val groupApis = materialize(groupSchema, call, firstSchema1, "Group", i = keyApis.size)
val groupToken = keyMarker.toClassSymbol(session)!!.fir
groupToken.callShapeData = CallShapeData.RefinedType(keyApis.map { it.scope.symbol })
val keyToken = groupMarker.toClassSymbol(session)!!.fir
keyToken.callShapeData = CallShapeData.RefinedType(groupApis.map { it.scope.symbol })
return buildScopeFunctionCall(call, originalSymbol, keyApis + groupApis, additionalDeclarations = listOf(groupToken, keyToken))
}
}
private fun buildNewTypeArgument(argument: ConeTypeProjection?, name: Name, hash: String): FirRegularClass {
val suggestedName = if (argument == null) {
"${name.asTokenName()}_$hash"
} else {
when (argument) {
is ConeStarProjection -> {
"${name.asTokenName()}_$hash"
}
is ConeKotlinTypeProjection -> {
val titleCase = argument.type.classId?.shortClassName
?.identifierOrNullIfSpecial?.titleCase()
?.substringBeforeLast("_")
?: DEFAULT_NAME
"${titleCase}_$hash"
}
}
}
val tokenId = nextName("${suggestedName}I")
val token = buildSchema(tokenId)
val dataFrameTypeId = nextName(suggestedName)
val dataFrameType = buildRegularClass {
moduleData = session.moduleData
resolvePhase = FirResolvePhase.BODY_RESOLVE
origin = FirDeclarationOrigin.Source
status = FirResolvedDeclarationStatusImpl(Visibilities.Local, Modality.ABSTRACT, EffectiveVisibility.Local)
deprecationsProvider = EmptyDeprecationsProvider
classKind = ClassKind.CLASS
scopeProvider = FirKotlinScopeProvider()
superTypeRefs += buildResolvedTypeRef {
type = ConeClassLikeTypeImpl(
ConeClassLookupTagWithFixedSymbol(tokenId, token.symbol),
emptyArray(),
isNullable = false
)
}
this.name = dataFrameTypeId.shortClassName
this.symbol = FirRegularClassSymbol(dataFrameTypeId)
}
return dataFrameType
}
private fun nextName(s: String) = ClassId(CallableId.PACKAGE_FQ_NAME_FOR_LOCAL, FqName(s), true)
private fun Name.asTokenName() = identifierOrNullIfSpecial?.titleCase() ?: DEFAULT_NAME
@OptIn(SymbolInternals::class)
private fun buildScopeFunctionCall(
call: FirFunctionCall,
originalSymbol: FirNamedFunctionSymbol,
dataSchemaApis: List<DataSchemaApi>,
additionalDeclarations: List<FirClass>
): FirFunctionCall {
val explicitReceiver = call.explicitReceiver
val receiverType = explicitReceiver?.resolvedType
val returnType = call.resolvedType
val scopeFunction = if (explicitReceiver != null) findLet() else findRun()
val originalSource = call.calleeReference.source
// original call is inserted later
call.transformCalleeReference(object : FirTransformer<Nothing?>() {
override fun <E : FirElement> transformElement(element: E, data: Nothing?): E {
return if (element is FirResolvedNamedReference) {
@Suppress("UNCHECKED_CAST")
buildResolvedNamedReference {
this.name = element.name
resolvedSymbol = originalSymbol
} as E
} else {
element
}
}
}, null)
val callExplicitReceiver = call.explicitReceiver
val callDispatchReceiver = call.dispatchReceiver
val callExtensionReceiver = call.extensionReceiver
val argument = buildAnonymousFunctionExpression {
isTrailingLambda = true
val fSymbol = FirAnonymousFunctionSymbol()
val target = FirFunctionTarget(null, isLambda = true)
anonymousFunction = buildAnonymousFunction {
resolvePhase = FirResolvePhase.BODY_RESOLVE
moduleData = session.moduleData
origin = FirDeclarationOrigin.Source
status = FirResolvedDeclarationStatusImpl(Visibilities.Local, Modality.FINAL, EffectiveVisibility.Local)
deprecationsProvider = EmptyDeprecationsProvider
returnTypeRef = buildResolvedTypeRef {
type = returnType
}
val parameterSymbol = receiverType?.let {
val itName = Name.identifier("it")
val parameterSymbol = FirValueParameterSymbol(itName)
valueParameters += buildValueParameter {
moduleData = session.moduleData
origin = FirDeclarationOrigin.Source
returnTypeRef = buildResolvedTypeRef {
type = receiverType
}
this.name = itName
this.symbol = parameterSymbol
containingFunctionSymbol = fSymbol
isCrossinline = false
isNoinline = false
isVararg = false
}
parameterSymbol
}
body = buildBlock {
this.coneTypeOrNull = returnType
dataSchemaApis.asReversed().forEach {
statements += it.schema
statements += it.scope
}
statements += additionalDeclarations
statements += buildReturnExpression {
if (parameterSymbol != null) {
val itPropertyAccess = buildPropertyAccessExpression {
coneTypeOrNull = receiverType
calleeReference = buildResolvedNamedReference {
name = parameterSymbol.name
resolvedSymbol = parameterSymbol
}
}
if (callDispatchReceiver != null) {
call.replaceDispatchReceiver(itPropertyAccess)
}
call.replaceExplicitReceiver(itPropertyAccess)
if (callExtensionReceiver != null) {
call.replaceExtensionReceiver(itPropertyAccess)
}
}
result = call
this.target = target
}
}
this.symbol = fSymbol
isLambda = true
hasExplicitParameterList = false
typeRef = buildResolvedTypeRef {
type = if (receiverType != null) {
ConeClassLikeTypeImpl(
ConeClassLikeLookupTagImpl(ClassId(FqName("kotlin"), Name.identifier("Function1"))),
typeArguments = arrayOf(receiverType, returnType),
isNullable = false
)
} else {
ConeClassLikeTypeImpl(
ConeClassLikeLookupTagImpl(ClassId(FqName("kotlin"), Name.identifier("Function0"))),
typeArguments = arrayOf(returnType),
isNullable = false
)
}
}
invocationKind = EventOccurrencesRange.EXACTLY_ONCE
inlineStatus = InlineStatus.Inline
}.also {
target.bind(it)
}
}
val newCall1 = buildFunctionCall {
// source = call.source makes IDE navigate to `let` declaration
source = null
this.coneTypeOrNull = returnType
if (receiverType != null) {
typeArguments += buildTypeProjectionWithVariance {
typeRef = buildResolvedTypeRef {
type = receiverType
}
variance = Variance.INVARIANT
}
}
typeArguments += buildTypeProjectionWithVariance {
typeRef = buildResolvedTypeRef {
type = returnType
}
variance = Variance.INVARIANT
}
dispatchReceiver = null
this.explicitReceiver = callExplicitReceiver
extensionReceiver = callExtensionReceiver ?: callDispatchReceiver
argumentList = buildResolvedArgumentList(
original = null,
linkedMapOf(argument to scopeFunction.valueParameterSymbols[0].fir)
)
calleeReference = buildResolvedNamedReference {
source = originalSource
this.name = scopeFunction.name
resolvedSymbol = scopeFunction
}
}
return newCall1
}
private fun materialize(
dataFrameSchema: PluginDataFrameSchema,
call: FirFunctionCall,
firstSchema: FirRegularClass,
prefix: String = "",
i: Int = 0
): List<DataSchemaApi> {
var i = i
val dataSchemaApis = mutableListOf<DataSchemaApi>()
val usedNames = mutableMapOf<String, Int>()
fun PluginDataFrameSchema.materialize(
schema: FirRegularClass? = null,
suggestedName: String? = null
): DataSchemaApi {
val schema = if (schema != null) {
schema
} else {
requireNotNull(suggestedName)
val uniqueSuffix = usedNames.compute(suggestedName) { _, i -> (i ?: 0) + 1 }
val name = nextName(suggestedName + uniqueSuffix)
buildSchema(name)
}
val scopeId = ClassId(CallableId.PACKAGE_FQ_NAME_FOR_LOCAL, FqName("Scope${i++}"), true)
val scope = buildRegularClass {
moduleData = session.moduleData
resolvePhase = FirResolvePhase.BODY_RESOLVE
origin = FirDeclarationOrigin.Source
status = FirResolvedDeclarationStatusImpl(Visibilities.Local, Modality.FINAL, EffectiveVisibility.Local)
deprecationsProvider = EmptyDeprecationsProvider
classKind = ClassKind.CLASS
scopeProvider = FirKotlinScopeProvider()
superTypeRefs += FirImplicitAnyTypeRef(null)
this.name = scopeId.shortClassName
this.symbol = FirRegularClassSymbol(scopeId)
}
val properties = columns().map {
fun PluginDataFrameSchema.materialize(column: SimpleCol): DataSchemaApi {
val text = call.source?.text ?: call.calleeReference.name
val name =
"${column.name.titleCase().replEscapeLineBreaks()}_${hashToTwoCharString(abs(text.hashCode()))}"
return materialize(suggestedName = "$prefix$name")
}
when (it) {
is SimpleColumnGroup -> {
val nestedSchema = PluginDataFrameSchema(it.columns()).materialize(it)
val columnsContainerReturnType =
ConeClassLikeTypeImpl(
ConeClassLikeLookupTagImpl(Names.COLUM_GROUP_CLASS_ID),
typeArguments = arrayOf(nestedSchema.schema.defaultType()),
isNullable = false
)
val dataRowReturnType =
ConeClassLikeTypeImpl(
ConeClassLikeLookupTagImpl(Names.DATA_ROW_CLASS_ID),
typeArguments = arrayOf(nestedSchema.schema.defaultType()),
isNullable = false
)
SchemaProperty(schema.defaultType(), PropertyName.of(it.name), dataRowReturnType, columnsContainerReturnType)
}
is SimpleFrameColumn -> {
val nestedClassMarker = PluginDataFrameSchema(it.columns()).materialize(it)
val frameColumnReturnType =
ConeClassLikeTypeImpl(
ConeClassLikeLookupTagImpl(Names.DF_CLASS_ID),
typeArguments = arrayOf(nestedClassMarker.schema.defaultType()),
isNullable = false
)
SchemaProperty(
marker = schema.defaultType(),
propertyName = PropertyName.of(it.name),
dataRowReturnType = frameColumnReturnType,
columnContainerReturnType = frameColumnReturnType.toFirResolvedTypeRef()
.projectOverDataColumnType()
)
}
is SimpleDataColumn -> SchemaProperty(
marker = schema.defaultType(),
propertyName = PropertyName.of(it.name),
dataRowReturnType = it.type.type(),
columnContainerReturnType = it.type.type().toFirResolvedTypeRef().projectOverDataColumnType()
)
}
}
schema.callShapeData = CallShapeData.Schema(properties)
scope.callShapeData = CallShapeData.Scope(properties, call.calleeReference.source)
val schemaApi = DataSchemaApi(schema, scope)
dataSchemaApis.add(schemaApi)
return schemaApi
}
dataFrameSchema.materialize(firstSchema)
return dataSchemaApis
}
data class DataSchemaApi(val schema: FirRegularClass, val scope: FirRegularClass)
private fun buildSchema(tokenId: ClassId): FirRegularClass {
val token = buildRegularClass {
moduleData = session.moduleData
resolvePhase = FirResolvePhase.BODY_RESOLVE
origin = FirDeclarationOrigin.Source
status = FirResolvedDeclarationStatusImpl(Visibilities.Local, Modality.ABSTRACT, EffectiveVisibility.Local)
deprecationsProvider = EmptyDeprecationsProvider
classKind = ClassKind.CLASS
scopeProvider = FirKotlinScopeProvider()
superTypeRefs += FirImplicitAnyTypeRef(null)
name = tokenId.shortClassName
this.symbol = FirRegularClassSymbol(tokenId)
}
return token
}
private fun findLet(): FirFunctionSymbol<*> {
return session.symbolProvider.getTopLevelFunctionSymbols(FqName("kotlin"), Name.identifier("let")).single()
}
private fun findRun(): FirFunctionSymbol<*> {
return session.symbolProvider.getTopLevelFunctionSymbols(FqName("kotlin"), Name.identifier("run")).single { it.typeParameterSymbols.size == 1 }
}
private fun String.titleCase() = replaceFirstChar { it.uppercaseChar() }
}
@@ -0,0 +1,211 @@
/*
* Copyright 2010-2022 JetBrains s.r.o. and Kotlin Programming Language contributors.
* Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
*/
package org.jetbrains.kotlinx.dataframe.plugin.extensions
import org.jetbrains.kotlin.backend.common.FileLoweringPass
import org.jetbrains.kotlin.backend.common.extensions.IrGenerationExtension
import org.jetbrains.kotlin.backend.common.extensions.IrPluginContext
import org.jetbrains.kotlin.ir.IrStatement
import org.jetbrains.kotlin.ir.backend.js.utils.valueArguments
import org.jetbrains.kotlin.ir.declarations.IrClass
import org.jetbrains.kotlin.ir.declarations.IrConstructor
import org.jetbrains.kotlin.ir.declarations.IrDeclarationOrigin
import org.jetbrains.kotlin.ir.declarations.IrDeclarationWithName
import org.jetbrains.kotlin.ir.declarations.IrFile
import org.jetbrains.kotlin.ir.declarations.IrModuleFragment
import org.jetbrains.kotlin.ir.declarations.IrPackageFragment
import org.jetbrains.kotlin.ir.declarations.IrProperty
import org.jetbrains.kotlin.ir.declarations.copyAttributes
import org.jetbrains.kotlin.ir.declarations.createBlockBody
import org.jetbrains.kotlin.ir.expressions.IrBody
import org.jetbrains.kotlin.ir.expressions.IrConst
import org.jetbrains.kotlin.ir.expressions.IrErrorCallExpression
import org.jetbrains.kotlin.ir.expressions.IrExpression
import org.jetbrains.kotlin.ir.expressions.IrTypeOperator
import org.jetbrains.kotlin.ir.expressions.IrTypeOperatorCall
import org.jetbrains.kotlin.ir.expressions.impl.IrCallImpl
import org.jetbrains.kotlin.ir.expressions.impl.IrConstImpl
import org.jetbrains.kotlin.ir.expressions.impl.IrConstructorCallImpl
import org.jetbrains.kotlin.ir.expressions.impl.IrDelegatingConstructorCallImpl
import org.jetbrains.kotlin.ir.expressions.impl.IrGetValueImpl
import org.jetbrains.kotlin.ir.expressions.impl.IrInstanceInitializerCallImpl
import org.jetbrains.kotlin.ir.expressions.impl.IrReturnImpl
import org.jetbrains.kotlin.ir.expressions.impl.IrTypeOperatorCallImpl
import org.jetbrains.kotlin.ir.symbols.IrValueSymbol
import org.jetbrains.kotlin.ir.symbols.UnsafeDuringIrConstructionAPI
import org.jetbrains.kotlin.ir.types.IrSimpleType
import org.jetbrains.kotlin.ir.types.IrType
import org.jetbrains.kotlin.ir.types.classFqName
import org.jetbrains.kotlin.ir.types.classOrFail
import org.jetbrains.kotlin.ir.types.classifierOrNull
import org.jetbrains.kotlin.ir.types.getClass
import org.jetbrains.kotlin.ir.util.constructors
import org.jetbrains.kotlin.ir.util.findAnnotation
import org.jetbrains.kotlin.ir.util.parentAsClass
import org.jetbrains.kotlin.ir.util.primaryConstructor
import org.jetbrains.kotlin.ir.util.superTypes
import org.jetbrains.kotlin.ir.visitors.IrElementTransformerVoid
import org.jetbrains.kotlin.name.CallableId
import org.jetbrains.kotlin.name.ClassId
import org.jetbrains.kotlin.name.FqName
import org.jetbrains.kotlin.name.Name
import org.jetbrains.kotlinx.dataframe.DataColumn
import org.jetbrains.kotlinx.dataframe.columns.ColumnGroup
import org.jetbrains.kotlinx.dataframe.plugin.utils.Names
class IrBodyFiller : IrGenerationExtension {
override fun generate(moduleFragment: IrModuleFragment, pluginContext: IrPluginContext) {
DataFrameFileLowering(pluginContext).lower(moduleFragment)
}
}
private class DataFrameFileLowering(val context: IrPluginContext) : FileLoweringPass, IrElementTransformerVoid() {
companion object {
val COLUMNS_CONTAINER_ID =
CallableId(ClassId(FqName("org.jetbrains.kotlinx.dataframe"), Name.identifier("ColumnsContainer")), Name.identifier("get"))
val COLUMNS_SCOPE_ID =
CallableId(ClassId(FqName("org.jetbrains.kotlinx.dataframe"), Name.identifier("ColumnsScope")), Name.identifier("get"))
val DATA_ROW_ID =
CallableId(ClassId(FqName("org.jetbrains.kotlinx.dataframe"), Name.identifier("DataRow")), Name.identifier("get"))
}
override fun lower(irFile: IrFile) {
irFile.transformChildren(this, null)
}
override fun visitConstructor(declaration: IrConstructor): IrStatement {
val origin = declaration.origin
if (!(origin is IrDeclarationOrigin.GeneratedByPlugin && origin.pluginKey is TokenGenerator.Key)) return declaration
declaration.body = generateBodyForDefaultConstructor(declaration)
return declaration
}
@OptIn(UnsafeDuringIrConstructionAPI::class)
private fun generateBodyForDefaultConstructor(declaration: IrConstructor): IrBody? {
val irType = declaration.returnType.superTypes()[0]
val symbol = irType.classOrFail.owner.primaryConstructor?.symbol ?: return null
val type = declaration.returnType as? IrSimpleType ?: return null
val delegatingAnyCall = IrDelegatingConstructorCallImpl(
-1,
-1,
irType,
symbol,
typeArgumentsCount = 0,
valueArgumentsCount = 0
).copyAttributes(declaration.parentAsClass)
val initializerCall = IrInstanceInitializerCallImpl(
-1,
-1,
(declaration.parent as? IrClass)?.symbol ?: return null,
type
)
return context.irFactory.createBlockBody(-1, -1, listOf(delegatingAnyCall, initializerCall))
}
@OptIn(UnsafeDuringIrConstructionAPI::class)
override fun visitProperty(declaration: IrProperty): IrStatement {
val origin = declaration.origin
val pluginKey = (origin as? IrDeclarationOrigin.GeneratedByPlugin)?.pluginKey as? DataFramePlugin
if (pluginKey == null) {
declaration.transformChildren(this, null)
return declaration
}
val getter = declaration.getter ?: return declaration
val constructors = context.referenceConstructors(ClassId(FqName("kotlin.jvm"), Name.identifier("JvmName")))
val jvmName = constructors.single { it.owner.valueParameters.size == 1 }
val marker =
((getter.extensionReceiverParameter!!.type as IrSimpleType).arguments.single() as IrSimpleType).classOrFail.owner
val jvmNameArg = "${marker.nestedName()}_${declaration.name.identifier}"
getter.annotations = listOf(
IrConstructorCallImpl(-1, -1, jvmName.owner.returnType, jvmName, 0, 0, 1)
.also {
it.putValueArgument(0, IrConstImpl.string(-1, -1, context.irBuiltIns.stringType, jvmNameArg))
}
)
val returnType = getter.returnType
val isDataColumn = returnType.classFqName!!.asString().let {
it == DataColumn::class.qualifiedName!! || it == ColumnGroup::class.qualifiedName!!
}
val get = if (isDataColumn) {
context
.referenceFunctions(COLUMNS_SCOPE_ID)
.single {
it.owner.valueParameters.size == 1 && it.owner.valueParameters[0].type == context.irBuiltIns.stringType
}
} else {
context
.referenceFunctions(DATA_ROW_ID)
.single {
it.owner.valueParameters.size == 1 && it.owner.valueParameters[0].type == context.irBuiltIns.stringType
}
}
val call = IrCallImpl(-1, -1, context.irBuiltIns.anyNType, get, 0, 1).also {
val thisSymbol: IrValueSymbol = getter.extensionReceiverParameter?.symbol!!
it.dispatchReceiver = IrGetValueImpl(-1, -1, thisSymbol)
val annotation = declaration.annotations.findAnnotation(Names.COLUMN_NAME_ANNOTATION.asSingleFqName())
val columnName = (annotation?.valueArguments?.get(0) as? IrConst<*>)?.value as? String
val columName = columnName ?: declaration.name.identifier
it.putValueArgument(0, IrConstImpl.string(-1, -1, context.irBuiltIns.stringType, columName))
}
val typeOp = IrTypeOperatorCallImpl(-1, -1, returnType, IrTypeOperator.CAST, returnType, call)
val returnExpression = IrReturnImpl(-1, -1, returnType, getter.symbol, typeOp)
getter.apply {
body = factory.createBlockBody(-1, -1, listOf(returnExpression))
}
return declaration
}
private fun IrDeclarationWithName.nestedName() = buildString { computeNestedName(this@nestedName, this) }
private fun computeNestedName(declaration: IrDeclarationWithName, result: StringBuilder): Boolean {
when (val parent = declaration.parent) {
is IrClass -> {
if (!computeNestedName(parent, result)) return false
}
is IrPackageFragment -> {}
else -> return false
}
if (result.isNotEmpty()) result.append('_')
result.append(declaration.name.asString())
return true
}
// org.jetbrains.kotlin.fir.backend.generators.CallAndReferenceGenerator#applyReceivers
override fun visitTypeOperator(expression: IrTypeOperatorCall): IrExpression {
if (isScope(expression.typeOperand)) {
return expression.replaceWithConstructorCall()
}
return super.visitTypeOperator(expression)
}
override fun visitErrorCallExpression(expression: IrErrorCallExpression): IrExpression {
if (!isScope(expression.type)) {
return expression
}
return expression.replaceWithConstructorCall()
}
@OptIn(UnsafeDuringIrConstructionAPI::class)
private fun isScope(type: IrType): Boolean {
val origin = (type.classifierOrNull?.owner as? IrClass)?.origin ?: return false
val fromPlugin = origin is IrDeclarationOrigin.GeneratedByPlugin && origin.pluginKey is DataFramePlugin
val scopeReference = type.classFqName?.shortName()?.asString()?.startsWith("Scope") ?: false
return fromPlugin || scopeReference
}
@OptIn(UnsafeDuringIrConstructionAPI::class)
private fun IrExpression.replaceWithConstructorCall(): IrConstructorCallImpl {
val constructor = type.getClass()!!.constructors.toList().single()
return IrConstructorCallImpl(-1, -1, type, constructor.symbol, 0, 0, 0)
}
}
@@ -0,0 +1,96 @@
package org.jetbrains.kotlinx.dataframe.plugin.extensions
import org.jetbrains.kotlin.fir.FirSession
import org.jetbrains.kotlin.fir.types.ConeClassLikeType
import org.jetbrains.kotlin.fir.types.ConeFlexibleType
import org.jetbrains.kotlin.fir.types.ConeKotlinType
import org.jetbrains.kotlin.fir.types.ConeNullability
import org.jetbrains.kotlin.fir.types.isNullable
import org.jetbrains.kotlin.fir.types.typeContext
import org.jetbrains.kotlin.fir.types.withNullability
import org.jetbrains.kotlin.name.ClassId
import org.jetbrains.kotlin.name.Name
import org.jetbrains.kotlin.name.StandardClassIds
interface KotlinTypeFacade : SessionContext {
val isTest: Boolean
fun Marker.type() = type
fun Marker.changeNullability(map: (Boolean) -> Boolean): Marker {
val coneNullability = when (map(type.isNullable)) {
true -> ConeNullability.NULLABLE
false -> ConeNullability.NOT_NULL
}
return Marker(type = type.withNullability(coneNullability, session.typeContext))
}
fun Marker.isList(): Boolean {
return type.isBuiltinType(List, isNullable = null)
}
fun Marker.typeArgument(): Marker {
val argument = when (val argument = type.typeArguments[0]) {
is ConeKotlinType -> argument
else -> error("${argument::class} ${argument}")
}
return Marker(argument)
}
}
interface SessionContext {
val session: FirSession
}
fun SessionContext(session: FirSession) = object : SessionContext {
override val session: FirSession = session
}
private val List = "List".collectionsId()
private fun ConeKotlinType.isBuiltinType(classId: ClassId, isNullable: Boolean?): Boolean {
if (this !is ConeClassLikeType) return false
return lookupTag.classId == classId && (isNullable == null || type.isNullable == isNullable)
}
private fun String.collectionsId() = ClassId(StandardClassIds.BASE_COLLECTIONS_PACKAGE, Name.identifier(this))
class KotlinTypeFacadeImpl(
override val session: FirSession,
override val isTest: Boolean
) : KotlinTypeFacade
class Marker private constructor(internal val type: ConeKotlinType) {
companion object {
operator fun invoke(type: ConeKotlinType): Marker {
val type = if (type is ConeFlexibleType) {
type.lowerBound
} else {
type
}
return Marker(type)
}
}
override fun toString(): String {
return "Marker(type=$type (${type::class}))"
}
override fun equals(other: Any?): Boolean {
if (this === other) return true
if (javaClass != other?.javaClass) return false
other as Marker
return type == other.type
}
override fun hashCode(): Int {
return type.hashCode()
}
}
fun ConeKotlinType.wrap(): Marker = Marker(this)
@@ -0,0 +1,50 @@
package org.jetbrains.kotlinx.dataframe.plugin.extensions
import org.jetbrains.kotlin.fir.FirSession
import org.jetbrains.kotlin.fir.declarations.FirResolvePhase
import org.jetbrains.kotlin.fir.declarations.getAnnotationByClassId
import org.jetbrains.kotlin.fir.expressions.FirFunctionCall
import org.jetbrains.kotlin.fir.extensions.FirExpressionResolutionExtension
import org.jetbrains.kotlin.fir.scopes.collectAllProperties
import org.jetbrains.kotlin.fir.scopes.impl.declaredMemberScope
import org.jetbrains.kotlin.fir.symbols.SymbolInternals
import org.jetbrains.kotlin.fir.symbols.impl.FirPropertySymbol
import org.jetbrains.kotlin.fir.types.ConeKotlinType
import org.jetbrains.kotlin.fir.types.classId
import org.jetbrains.kotlin.fir.types.resolvedType
import org.jetbrains.kotlin.fir.types.toRegularClassSymbol
import org.jetbrains.kotlinx.dataframe.plugin.utils.Names
class ReturnTypeBasedReceiverInjector(session: FirSession) : FirExpressionResolutionExtension(session) {
companion object {
private val SCHEMA_TYPES = setOf(
Names.DF_CLASS_ID,
Names.GROUP_BY_CLASS_ID,
Names.DATA_ROW_CLASS_ID,
Names.COLUM_GROUP_CLASS_ID,
)
}
@OptIn(SymbolInternals::class)
override fun addNewImplicitReceivers(functionCall: FirFunctionCall): List<ConeKotlinType> {
val callReturnType = functionCall.resolvedType
return if (callReturnType.classId in SCHEMA_TYPES) {
val typeArguments = callReturnType.typeArguments
typeArguments
.mapNotNull {
val symbol = (it as? ConeKotlinType)?.toRegularClassSymbol(session)
symbol?.takeIf { it.fir.callShapeData != null }
}
.takeIf { it.size == typeArguments.size }
.orEmpty()
.flatMap { marker ->
marker.declaredMemberScope(session, FirResolvePhase.DECLARATIONS).collectAllProperties()
.filterIsInstance<FirPropertySymbol>()
.filter { it.getAnnotationByClassId(Names.SCOPE_PROPERTY_ANNOTATION, session) != null }
.map { it.resolvedReturnType }
}
} else {
emptyList()
}
}
}
@@ -0,0 +1,151 @@
package org.jetbrains.kotlinx.dataframe.plugin.extensions
import org.jetbrains.kotlin.GeneratedDeclarationKey
import org.jetbrains.kotlin.descriptors.EffectiveVisibility
import org.jetbrains.kotlin.descriptors.Modality
import org.jetbrains.kotlin.descriptors.Visibilities
import org.jetbrains.kotlin.fir.FirSession
import org.jetbrains.kotlin.fir.caches.FirCache
import org.jetbrains.kotlin.fir.caches.createCache
import org.jetbrains.kotlin.fir.caches.firCachesFactory
import org.jetbrains.kotlin.fir.caches.getValue
import org.jetbrains.kotlinx.dataframe.plugin.utils.Names
import org.jetbrains.kotlinx.dataframe.plugin.utils.generateExtensionProperty
import org.jetbrains.kotlin.fir.declarations.FirProperty
import org.jetbrains.kotlin.fir.expressions.FirAnnotation
import org.jetbrains.kotlin.fir.expressions.builder.buildAnnotation
import org.jetbrains.kotlin.fir.expressions.builder.buildAnnotationArgumentMapping
import org.jetbrains.kotlin.fir.expressions.builder.buildLiteralExpression
import org.jetbrains.kotlin.fir.extensions.FirDeclarationGenerationExtension
import org.jetbrains.kotlin.fir.extensions.MemberGenerationContext
import org.jetbrains.kotlin.fir.plugin.createConstructor
import org.jetbrains.kotlin.fir.plugin.createMemberProperty
import org.jetbrains.kotlin.fir.resolve.defaultType
import org.jetbrains.kotlin.fir.symbols.SymbolInternals
import org.jetbrains.kotlin.fir.symbols.impl.ConeClassLikeLookupTagImpl
import org.jetbrains.kotlin.fir.symbols.impl.FirClassSymbol
import org.jetbrains.kotlin.fir.symbols.impl.FirConstructorSymbol
import org.jetbrains.kotlin.fir.symbols.impl.FirPropertySymbol
import org.jetbrains.kotlin.fir.toFirResolvedTypeRef
import org.jetbrains.kotlin.fir.types.FirResolvedTypeRef
import org.jetbrains.kotlin.fir.types.builder.buildResolvedTypeRef
import org.jetbrains.kotlin.fir.types.impl.ConeClassLikeTypeImpl
import org.jetbrains.kotlin.name.CallableId
import org.jetbrains.kotlin.name.Name
import org.jetbrains.kotlin.name.SpecialNames
import org.jetbrains.kotlin.types.ConstantValueKind
class TokenGenerator(session: FirSession) : FirDeclarationGenerationExtension(session) {
object Key : GeneratedDeclarationKey()
@OptIn(SymbolInternals::class)
private val propertiesCache: FirCache<FirClassSymbol<*>, Map<Name, List<FirProperty>>?, Nothing?> =
session.firCachesFactory.createCache { k ->
val callShapeData = k.fir.callShapeData ?: return@createCache null
when (callShapeData) {
is CallShapeData.Schema -> callShapeData.columns.withIndex().associate { (index, property) ->
val resolvedTypeRef = buildResolvedTypeRef {
type = property.dataRowReturnType
}
val identifier = property.propertyName.identifier
identifier to listOf(buildProperty(resolvedTypeRef, identifier, k, property.propertyName.columnNameAnnotation, order = index))
}
is CallShapeData.RefinedType -> callShapeData.scopes.associate {
val identifier = Name.identifier(it.name.identifier.replaceFirstChar { it.lowercaseChar() })
identifier to listOf(buildProperty(it.defaultType().toFirResolvedTypeRef(), identifier, k, isScopeProperty = true))
}
is CallShapeData.Scope -> callShapeData.columns.associate { schemaProperty ->
val propertyName = schemaProperty.propertyName
val callableId = CallableId(k.classId, propertyName.identifier)
val dataRowExtension = generateExtensionProperty(
callableId = callableId,
receiverType = ConeClassLikeTypeImpl(
ConeClassLikeLookupTagImpl(Names.DATA_ROW_CLASS_ID),
typeArguments = arrayOf(schemaProperty.marker),
isNullable = false
),
propertyName = propertyName,
returnTypeRef = schemaProperty.dataRowReturnType.toFirResolvedTypeRef(),
symbol = k,
effectiveVisibility = EffectiveVisibility.Local,
source = callShapeData.source
)
val columnContainerExtension = generateExtensionProperty(
callableId = callableId,
receiverType = ConeClassLikeTypeImpl(
ConeClassLikeLookupTagImpl(Names.COLUMNS_SCOPE_CLASS_ID),
typeArguments = arrayOf(schemaProperty.marker),
isNullable = false
),
propertyName = propertyName,
returnTypeRef = schemaProperty.columnContainerReturnType.toFirResolvedTypeRef(),
symbol = k,
effectiveVisibility = EffectiveVisibility.Local,
source = callShapeData.source
)
propertyName.identifier to listOf(dataRowExtension, columnContainerExtension)
}
}
}
@OptIn(SymbolInternals::class)
override fun getCallableNamesForClass(classSymbol: FirClassSymbol<*>, context: MemberGenerationContext): Set<Name> {
val destination = mutableSetOf<Name>()
when (classSymbol.fir.callShapeData) {
is CallShapeData.RefinedType -> destination.add(SpecialNames.INIT)
is CallShapeData.Schema -> destination.add(SpecialNames.INIT)
is CallShapeData.Scope -> destination.add(SpecialNames.INIT)
null -> Unit
}
return propertiesCache.getValue(classSymbol)?.values?.flatten()?.mapTo(destination) { it.name } ?: emptySet()
}
override fun generateProperties(callableId: CallableId, context: MemberGenerationContext?): List<FirPropertySymbol> {
val owner = context?.owner ?: return emptyList()
val properties = propertiesCache.getValue(owner)?.get(callableId.callableName) ?: return emptyList()
return properties.map { it.symbol }
}
private fun buildProperty(
resolvedTypeRef: FirResolvedTypeRef,
propertyName: Name,
k: FirClassSymbol<*>,
columnNameAnnotation: FirAnnotation? = null,
isScopeProperty: Boolean = false,
order: Int? = null,
): FirProperty {
return createMemberProperty(k, Key, propertyName, resolvedTypeRef.type) {
modality = Modality.ABSTRACT
visibility = Visibilities.Public
}.apply {
val annotations = mutableListOf<FirAnnotation>()
if (order != null) {
annotations += buildAnnotation {
annotationTypeRef = buildResolvedTypeRef {
type = Names.ORDER_ANNOTATION.defaultType(emptyList())
}
argumentMapping = buildAnnotationArgumentMapping {
mapping[Names.ORDER_ARGUMENT] = buildLiteralExpression(null, ConstantValueKind.Int, order, setType = true)
}
}
}
if (isScopeProperty) {
annotations += buildAnnotation {
annotationTypeRef = buildResolvedTypeRef {
type = Names.SCOPE_PROPERTY_ANNOTATION.defaultType(emptyList())
}
argumentMapping = buildAnnotationArgumentMapping()
}
}
columnNameAnnotation?.let {
annotations += it
}
replaceAnnotations(annotations)
}
}
override fun generateConstructors(context: MemberGenerationContext): List<FirConstructorSymbol> {
return listOf(createConstructor(context.owner, Key, isPrimary = true).symbol)
}
}
@@ -0,0 +1,155 @@
package org.jetbrains.kotlinx.dataframe.plugin.extensions
import org.jetbrains.kotlin.fir.FirSession
import org.jetbrains.kotlin.fir.analysis.checkers.toClassLikeSymbol
import org.jetbrains.kotlin.fir.declarations.getAnnotationByClassId
import org.jetbrains.kotlinx.dataframe.plugin.utils.Names
import org.jetbrains.kotlinx.dataframe.plugin.utils.generateExtensionProperty
import org.jetbrains.kotlinx.dataframe.plugin.utils.projectOverDataColumnType
import org.jetbrains.kotlin.fir.declarations.hasAnnotation
import org.jetbrains.kotlin.fir.declarations.utils.isLocal
import org.jetbrains.kotlin.fir.expressions.FirLiteralExpression
import org.jetbrains.kotlin.fir.extensions.ExperimentalTopLevelDeclarationsGenerationApi
import org.jetbrains.kotlin.fir.extensions.FirDeclarationGenerationExtension
import org.jetbrains.kotlin.fir.extensions.FirDeclarationPredicateRegistrar
import org.jetbrains.kotlin.fir.extensions.MemberGenerationContext
import org.jetbrains.kotlin.fir.extensions.predicate.LookupPredicate
import org.jetbrains.kotlin.fir.extensions.predicateBasedProvider
import org.jetbrains.kotlin.fir.symbols.impl.ConeClassLikeLookupTagImpl
import org.jetbrains.kotlin.fir.symbols.impl.FirPropertySymbol
import org.jetbrains.kotlin.fir.symbols.impl.FirRegularClassSymbol
import org.jetbrains.kotlin.fir.toFirResolvedTypeRef
import org.jetbrains.kotlin.fir.types.ConeClassLikeType
import org.jetbrains.kotlin.fir.types.ConeTypeProjection
import org.jetbrains.kotlin.fir.types.classId
import org.jetbrains.kotlin.fir.types.coneType
import org.jetbrains.kotlin.fir.types.constructType
import org.jetbrains.kotlin.fir.types.impl.ConeClassLikeTypeImpl
import org.jetbrains.kotlin.fir.types.toSymbol
import org.jetbrains.kotlin.fir.types.toTypeProjection
import org.jetbrains.kotlin.name.CallableId
import org.jetbrains.kotlin.name.FqName
import org.jetbrains.kotlin.types.Variance
import org.jetbrains.kotlinx.dataframe.annotations.DataSchema
import org.jetbrains.kotlinx.dataframe.plugin.extensions.impl.PropertyName
/**
* extensions inside scope classes are generated here:
* @see org.jetbrains.kotlinx.dataframe.plugin.extensions.TokenGenerator
*/
class TopLevelExtensionsGenerator(session: FirSession) : FirDeclarationGenerationExtension(session) {
private companion object {
val dataSchema = FqName(DataSchema::class.qualifiedName!!)
}
private val predicateBasedProvider = session.predicateBasedProvider
private val matchedClasses by lazy {
predicateBasedProvider.getSymbolsByPredicate(predicate).filterIsInstance<FirRegularClassSymbol>()
}
private val predicate: LookupPredicate = LookupPredicate.BuilderContext.annotated(dataSchema)
override fun FirDeclarationPredicateRegistrar.registerPredicates() {
register(predicate)
}
private val fields by lazy {
matchedClasses.filterNot { it.isLocal }.flatMap { classSymbol ->
classSymbol.declarationSymbols.filterIsInstance<FirPropertySymbol>().map { propertySymbol ->
val callableId = propertySymbol.callableId
DataSchemaField(
classSymbol,
propertySymbol,
CallableId(packageName = callableId.packageName, className = null, callableName = callableId.callableName)
)
}
}
}
private data class DataSchemaField(
val classSymbol: FirRegularClassSymbol,
val propertySymbol: FirPropertySymbol,
val callableId: CallableId
)
@OptIn(ExperimentalTopLevelDeclarationsGenerationApi::class)
override fun getTopLevelCallableIds(): Set<CallableId> {
return buildSet {
fields.mapTo(this) { it.callableId }
}
}
override fun generateProperties(callableId: CallableId, context: MemberGenerationContext?): List<FirPropertySymbol> {
val owner = context?.owner
return when (owner) {
null -> fields.filter { it.callableId == callableId }.flatMap { (owner, property, callableId) ->
var resolvedReturnTypeRef = property.resolvedReturnTypeRef
val columnName = property.getAnnotationByClassId(Names.COLUMN_NAME_ANNOTATION, session)?.let { annotation ->
(annotation.argumentMapping.mapping[Names.COLUMN_NAME_ARGUMENT] as? FirLiteralExpression)?.value as? String?
}
val name = property.name
val marker = owner.constructType(arrayOf(), isNullable = false).toTypeProjection(Variance.INVARIANT)
val columnGroupProjection: ConeTypeProjection? = if (resolvedReturnTypeRef.coneType.classId?.equals(
Names.DATA_ROW_CLASS_ID) == true) {
resolvedReturnTypeRef.coneType.typeArguments[0]
} else if (resolvedReturnTypeRef.toClassLikeSymbol(session)?.hasAnnotation(Names.DATA_SCHEMA_CLASS_ID, session) == true) {
resolvedReturnTypeRef.coneType
} else {
null
}
if (
resolvedReturnTypeRef.type.classId?.equals(Names.LIST) == true &&
(resolvedReturnTypeRef.type.typeArguments[0] as? ConeClassLikeType)?.toSymbol(session)?.hasAnnotation(
Names.DATA_SCHEMA_CLASS_ID, session) == true
) {
require(columnGroupProjection == null)
resolvedReturnTypeRef = ConeClassLikeTypeImpl(
ConeClassLikeLookupTagImpl(Names.DF_CLASS_ID),
typeArguments = arrayOf(resolvedReturnTypeRef.type.typeArguments[0]),
isNullable = false
).toFirResolvedTypeRef()
}
val rowExtension = generateExtensionProperty(
callableId = callableId,
receiverType = ConeClassLikeTypeImpl(
ConeClassLikeLookupTagImpl(Names.DATA_ROW_CLASS_ID),
typeArguments = arrayOf(marker),
isNullable = false
),
propertyName = PropertyName.of(name, columnName?.let { PropertyName.buildAnnotation(it) }),
returnTypeRef = resolvedReturnTypeRef,
source = owner.source
)
val columnReturnType = when {
columnGroupProjection != null -> {
ConeClassLikeTypeImpl(
ConeClassLikeLookupTagImpl(Names.COLUM_GROUP_CLASS_ID),
typeArguments = arrayOf(columnGroupProjection),
isNullable = false
).toFirResolvedTypeRef()
}
else -> resolvedReturnTypeRef.projectOverDataColumnType().toFirResolvedTypeRef()
}
val columnsContainerExtension = generateExtensionProperty(
callableId = callableId,
receiverType = ConeClassLikeTypeImpl(
ConeClassLikeLookupTagImpl(Names.COLUMNS_CONTAINER_CLASS_ID),
typeArguments = arrayOf(marker),
isNullable = false
),
propertyName = PropertyName.of(name, columnName?.let { PropertyName.buildAnnotation(it) }),
returnTypeRef = columnReturnType,
source = owner.source
)
listOf(rowExtension.symbol, columnsContainerExtension.symbol)
}
else -> emptyList()
}
}
}
@@ -0,0 +1,52 @@
package org.jetbrains.kotlinx.dataframe.plugin.extensions.impl
import org.jetbrains.kotlin.fir.expressions.FirAnnotation
import org.jetbrains.kotlin.fir.expressions.builder.buildAnnotationArgumentMapping
import org.jetbrains.kotlin.fir.expressions.builder.buildLiteralExpression
import org.jetbrains.kotlin.fir.resolve.defaultType
import org.jetbrains.kotlin.fir.types.builder.buildResolvedTypeRef
import org.jetbrains.kotlin.name.Name
import org.jetbrains.kotlin.types.ConstantValueKind
import org.jetbrains.kotlinx.dataframe.codeGen.ValidFieldName
import org.jetbrains.kotlinx.dataframe.plugin.utils.Names
data class PropertyName(val identifier: Name, val columnNameAnnotation: FirAnnotation?) {
companion object {
fun of(name: String): PropertyName {
val valid = ValidFieldName.of(name)
var columnName = false
val identifier = if (valid.unquoted != name) {
columnName = true
Name.identifier(valid.unquoted)
} else {
Name.identifier(name)
}
val columnNameAnnotation: FirAnnotation? = if (columnName) {
buildAnnotation(name)
} else {
null
}
return PropertyName(identifier, columnNameAnnotation)
}
fun buildAnnotation(name: String): FirAnnotation {
return org.jetbrains.kotlin.fir.expressions.builder.buildAnnotation {
annotationTypeRef = buildResolvedTypeRef {
type = Names.COLUMN_NAME_ANNOTATION.defaultType(emptyList())
}
argumentMapping = buildAnnotationArgumentMapping {
mapping[Names.COLUMN_NAME_ARGUMENT] = buildLiteralExpression(
source = null,
kind = ConstantValueKind.String,
value = name,
setType = true
)
}
}
}
fun of(identifier: Name, columnNameAnnotation: FirAnnotation?): PropertyName {
return PropertyName(identifier, columnNameAnnotation)
}
}
}
@@ -0,0 +1,12 @@
package org.jetbrains.kotlinx.dataframe.plugin.extensions.impl
import org.jetbrains.kotlin.fir.types.ConeKotlinType
import org.jetbrains.kotlin.fir.types.ConeTypeProjection
data class SchemaProperty(
val marker: ConeTypeProjection,
val propertyName: PropertyName,
val dataRowReturnType: ConeKotlinType,
val columnContainerReturnType: ConeKotlinType,
val override: Boolean = false
)
@@ -0,0 +1,54 @@
package org.jetbrains.kotlinx.dataframe.plugin.impl
import org.jetbrains.kotlinx.dataframe.AnyCol
import org.jetbrains.kotlinx.dataframe.DataColumn
import org.jetbrains.kotlinx.dataframe.DataFrame
import org.jetbrains.kotlinx.dataframe.api.asDataColumn
import org.jetbrains.kotlinx.dataframe.api.cast
import org.jetbrains.kotlinx.dataframe.api.dataFrameOf
import org.jetbrains.kotlinx.dataframe.columns.ColumnGroup
import org.jetbrains.kotlinx.dataframe.columns.FrameColumn
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.TypeApproximation
fun PluginDataFrameSchema.asDataFrame(): DataFrame<ConeTypesAdapter> {
val df = columns().map()
return df
}
fun DataFrame<ConeTypesAdapter>.toPluginDataFrameSchema() = PluginDataFrameSchema(columns().mapBack())
interface ConeTypesAdapter
private fun List<SimpleCol>.map(): DataFrame<ConeTypesAdapter> {
val columns = map {
it.asDataColumn()
}
return dataFrameOf(columns).cast()
}
fun SimpleCol.asDataColumn(): DataColumn<*> {
val column = when (this) {
is SimpleDataColumn -> DataColumn.createByType(this.name, listOf(this.type))
is SimpleColumnGroup -> DataColumn.createColumnGroup(this.name, this.columns().map()).asDataColumn()
is SimpleFrameColumn -> DataColumn.createFrameColumn(this.name, listOf(this.columns().map()))
}
return column
}
private fun List<AnyCol>.mapBack(): List<SimpleCol> = map { it.asSimpleColumn() }
fun AnyCol.asSimpleColumn(): SimpleCol {
return when (this) {
is ColumnGroup<*> -> {
SimpleColumnGroup(name(), columns().mapBack())
}
is FrameColumn<*> -> {
SimpleFrameColumn(name(), this[0].columns().mapBack())
}
else -> {
SimpleDataColumn(name(), this[0] as TypeApproximation)
}
}
}
@@ -0,0 +1,42 @@
package org.jetbrains.kotlinx.dataframe.plugin.impl
import org.jetbrains.kotlinx.dataframe.plugin.impl.AbstractInterpreter.*
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.GroupBy
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.TypeApproximation
import org.jetbrains.kotlinx.dataframe.plugin.impl.data.DataFrameCallableId
import kotlin.properties.PropertyDelegateProvider
import kotlin.properties.ReadOnlyProperty
typealias ExpectedArgumentProvider<T> = PropertyDelegateProvider<Any?, ReadOnlyProperty<Arguments, T>>
fun <T> AbstractInterpreter<T>.dataFrame(
name: ArgumentName? = null
): ExpectedArgumentProvider<PluginDataFrameSchema> = arg(name, lens = Interpreter.Schema)
fun <T> AbstractInterpreter<T>.type(
name: ArgumentName? = null
): ExpectedArgumentProvider<TypeApproximation> = arg(name, lens = Interpreter.ReturnType)
fun <T, E : Enum<E>> AbstractInterpreter<T>.enum(
name: ArgumentName? = null,
defaultValue: DefaultValue<E> = Absent
): ExpectedArgumentProvider<E> = argConvert(name = name, defaultValue = defaultValue) { it: DataFrameCallableId ->
val forName: Class<*> = Class.forName("${it.packageName}.${it.className}")
@Suppress("UNCHECKED_CAST")
java.lang.Enum.valueOf(forName as Class<out Enum<*>>, it.callableName) as E
}
internal fun <T> AbstractInterpreter<T>.dsl(
name: ArgumentName? = null
): ExpectedArgumentProvider<(Any, Map<String, Interpreter.Success<Any?>>) -> Unit> =
arg(name, lens = Interpreter.Dsl, defaultValue = Present(value = {_, _ -> }))
internal fun <T> AbstractInterpreter<T>.ignore(
name: ArgumentName? = null
): ExpectedArgumentProvider<Nothing?> =
arg(name, lens = Interpreter.Id, defaultValue = Present(null))
internal fun <T> AbstractInterpreter<T>.groupBy(
name: ArgumentName? = null
): ExpectedArgumentProvider<GroupBy> = arg(name, lens = Interpreter.GroupBy)
@@ -0,0 +1,150 @@
package org.jetbrains.kotlinx.dataframe.plugin.impl
import org.jetbrains.kotlinx.dataframe.plugin.extensions.KotlinTypeFacade
import kotlin.properties.PropertyDelegateProvider
import kotlin.properties.ReadOnlyProperty
import kotlin.reflect.KType
import kotlin.reflect.typeOf
interface Interpreter<T> {
val expectedArguments: List<ExpectedArgument>
data class ExpectedArgument(
val name: String,
val klass: KType,
val lens: Lens,
val defaultValue: DefaultValue<*>
)
sealed interface Lens
data object Value : Lens
data object ReturnType : Lens
data object Dsl : Lens
data object Schema : Lens
data object GroupBy : Lens
data object Id : Lens
// required to compute whether resulting schema should be inheritor of previous class or a new class
fun startingSchema(arguments: Map<String, Success<Any?>>, kotlinTypeFacade: KotlinTypeFacade): PluginDataFrameSchema?
fun interpret(arguments: Map<String, Success<Any?>>, kotlinTypeFacade: KotlinTypeFacade): InterpretationResult<T>
sealed interface InterpretationResult<out T>
class Success<out T>(val value: T) : InterpretationResult<T> {
override fun equals(other: Any?): Boolean {
if (this === other) return true
if (javaClass != other?.javaClass) return false
other as Success<*>
return value == other.value
}
override fun hashCode(): Int {
return value?.hashCode() ?: 0
}
}
class Error(val message: String?) : InterpretationResult<Nothing>
}
sealed interface DefaultValue<out T>
class Present<T>(val value: T) : DefaultValue<T>
data object Absent : DefaultValue<Nothing>
open class Arguments(private val arguments: Map<String, Interpreter.Success<Any?>>, kotlinTypeFacade: KotlinTypeFacade): KotlinTypeFacade by kotlinTypeFacade {
operator fun get(s: String): Any? = (arguments[s] ?: error("")).value
operator fun contains(key: String): Boolean {
return arguments.contains(key)
}
}
abstract class AbstractInterpreter<T> : Interpreter<T> {
@PublishedApi
internal val _expectedArguments: MutableList<Interpreter.ExpectedArgument> = mutableListOf()
override val expectedArguments: List<Interpreter.ExpectedArgument> = _expectedArguments
protected open val Arguments.startingSchema: PluginDataFrameSchema? get() = null
final override fun startingSchema(arguments: Map<String, Interpreter.Success<Any?>>, kotlinTypeFacade: KotlinTypeFacade): PluginDataFrameSchema? {
return Arguments(arguments, kotlinTypeFacade).startingSchema
}
inline fun <Value, reified CompileTimeValue> argConvert(
defaultValue: DefaultValue<Value> = Absent,
name: ArgumentName? = null,
lens: Interpreter.Lens = Interpreter.Value,
crossinline converter: (CompileTimeValue) -> Value
): PropertyDelegateProvider<Any?, ReadOnlyProperty<Arguments, Value>> = PropertyDelegateProvider { thisRef: Any?, property ->
val name = name?.value ?: property.name
_expectedArguments.add(Interpreter.ExpectedArgument(name, typeOf<CompileTimeValue>(), lens, defaultValue))
ReadOnlyProperty { args, _ ->
if (name !in args && defaultValue is Present) {
defaultValue.value
} else {
converter(args[name] as CompileTimeValue)
}
}
}
fun <Value> arg(
name: ArgumentName? = null,
expectedType: KType? = null,
defaultValue: DefaultValue<Value> = Absent,
lens: Interpreter.Lens = Interpreter.Value
): PropertyDelegateProvider<Any?, ReadOnlyProperty<Arguments, Value>> = PropertyDelegateProvider { thisRef: Any?, property ->
val name = name?.value ?: property.name
_expectedArguments.add(
Interpreter.ExpectedArgument(
name,
expectedType ?: property.returnType,
lens,
defaultValue
)
)
ReadOnlyProperty { args, _ ->
if (name !in args && defaultValue is Present) {
defaultValue.value
} else {
@Suppress("UNCHECKED_CAST")
args[name] as Value
}
}
}
class ArgumentName private constructor(val value: String) {
companion object {
fun of(name: String): ArgumentName = ArgumentName(name)
}
}
fun name(name: String): ArgumentName = ArgumentName.of(name)
final override fun interpret(arguments: Map<String, Interpreter.Success<Any?>>, kotlinTypeFacade: KotlinTypeFacade): Interpreter.InterpretationResult<T> {
return try {
Arguments(arguments, kotlinTypeFacade).interpret().let { Interpreter.Success(it) }
} catch (e: Exception) {
Interpreter.Error(e.message + e.stackTrace.contentToString())
}
}
abstract fun Arguments.interpret(): T
}
interface SchemaModificationInterpreter : Interpreter<PluginDataFrameSchema> {
override fun interpret(arguments: Map<String, Interpreter.Success<Any?>>, kotlinTypeFacade: KotlinTypeFacade): Interpreter.InterpretationResult<PluginDataFrameSchema>
}
abstract class AbstractSchemaModificationInterpreter :
AbstractInterpreter<PluginDataFrameSchema>(),
SchemaModificationInterpreter
@@ -0,0 +1,137 @@
package org.jetbrains.kotlinx.dataframe.plugin.impl
import org.jetbrains.kotlin.fir.analysis.checkers.fullyExpandedClassId
import org.jetbrains.kotlin.fir.types.ConeKotlinType
import org.jetbrains.kotlin.fir.types.ConeNullability
import org.jetbrains.kotlin.fir.types.isNullable
import org.jetbrains.kotlin.fir.types.renderReadable
import org.jetbrains.kotlinx.dataframe.plugin.extensions.KotlinTypeFacade
import org.jetbrains.kotlinx.dataframe.plugin.extensions.wrap
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.TypeApproximation
import org.jetbrains.kotlinx.dataframe.plugin.pluginDataFrameSchema
import org.jetbrains.kotlinx.dataframe.plugin.utils.Names
data class PluginDataFrameSchema(
private val columns: List<SimpleCol>
) {
companion object {
val EMPTY = PluginDataFrameSchema(emptyList())
}
fun columns(): List<SimpleCol> {
return columns
}
override fun toString(): String {
return columns.asString()
}
}
fun PluginDataFrameSchema.add(name: String, type: ConeKotlinType, context: KotlinTypeFacade): PluginDataFrameSchema {
return PluginDataFrameSchema(columns() + context.simpleColumnOf(name, type))
}
private fun List<SimpleCol>.asString(indent: String = ""): String {
if (isEmpty()) return "$indent<empty compile time schema>"
return joinToString("\n") {
val col = when (it) {
is SimpleFrameColumn -> {
"${it.name}: *\n" + it.columns().asString("$indent ")
}
is SimpleColumnGroup -> {
"${it.name}:\n" + it.columns().asString("$indent ")
}
is SimpleDataColumn -> {
"${it.name}: ${it.type.type.renderReadable()}"
}
}
"$indent$col"
}
}
sealed interface SimpleCol {
val name: String
fun name(): String {
return name
}
fun rename(s: String): SimpleCol
}
data class SimpleDataColumn(
override val name: String,
val type: TypeApproximation
) : SimpleCol {
override fun name(): String {
return name
}
override fun rename(s: String): SimpleDataColumn {
return SimpleDataColumn(s, type)
}
fun changeType(type: TypeApproximation): SimpleDataColumn {
return SimpleDataColumn(name, type)
}
}
data class SimpleFrameColumn(
override val name: String,
private val columns: List<SimpleCol>
) : SimpleCol {
fun columns(): List<SimpleCol> {
return columns
}
override fun rename(s: String): SimpleFrameColumn {
return SimpleFrameColumn(s, columns)
}
}
data class SimpleColumnGroup(
override val name: String,
private val columns: List<SimpleCol>
) : SimpleCol {
fun columns(): List<SimpleCol> {
return columns
}
override fun rename(s: String): SimpleColumnGroup {
return SimpleColumnGroup(s, columns)
}
}
fun KotlinTypeFacade.simpleColumnOf(name: String, type: ConeKotlinType): SimpleCol {
return if (type.fullyExpandedClassId(session) == Names.DATA_ROW_CLASS_ID) {
val schema = pluginDataFrameSchema(type.typeArguments[0])
val group = SimpleColumnGroup(name, schema.columns())
val column = if (type.isNullable) {
makeNullable(group)
} else {
group
}
column
} else if (type.fullyExpandedClassId(session) == Names.DF_CLASS_ID && type.nullability == ConeNullability.NOT_NULL) {
val schema = pluginDataFrameSchema(type.typeArguments[0])
SimpleFrameColumn(name, schema.columns())
} else {
SimpleDataColumn(name, type.wrap())
}
}
internal fun KotlinTypeFacade.makeNullable(column: SimpleCol): SimpleCol {
return when (column) {
is SimpleColumnGroup -> {
SimpleColumnGroup(column.name, column.columns().map { makeNullable(it) })
}
is SimpleFrameColumn -> column
is SimpleDataColumn -> SimpleDataColumn(column.name, column.type.changeNullability { true })
}
}
@@ -0,0 +1,64 @@
package org.jetbrains.kotlinx.dataframe.plugin.impl.api
import org.jetbrains.kotlin.fir.types.ConeNullability
import org.jetbrains.kotlin.fir.types.ConeSimpleKotlinType
import org.jetbrains.kotlin.fir.types.typeContext
import org.jetbrains.kotlin.fir.types.withNullability
import org.jetbrains.kotlinx.dataframe.plugin.extensions.KotlinTypeFacade
import org.jetbrains.kotlinx.dataframe.plugin.extensions.Marker
import org.jetbrains.kotlinx.dataframe.plugin.impl.AbstractSchemaModificationInterpreter
import org.jetbrains.kotlinx.dataframe.plugin.impl.Arguments
import org.jetbrains.kotlinx.dataframe.plugin.impl.PluginDataFrameSchema
import org.jetbrains.kotlinx.dataframe.plugin.impl.Present
import org.jetbrains.kotlinx.dataframe.plugin.impl.SimpleCol
import org.jetbrains.kotlinx.dataframe.plugin.impl.SimpleColumnGroup
import org.jetbrains.kotlinx.dataframe.plugin.impl.SimpleDataColumn
import org.jetbrains.kotlinx.dataframe.plugin.impl.dataFrame
class DropNulls0 : AbstractSchemaModificationInterpreter() {
val Arguments.receiver: PluginDataFrameSchema by dataFrame()
val Arguments.columns: ColumnsResolver by arg()
override fun Arguments.interpret(): PluginDataFrameSchema {
return PluginDataFrameSchema(fillNullsImpl(receiver.columns(), columns.resolve(receiver).mapTo(mutableSetOf()) { it.path.path }, emptyList()))
}
}
class DropNa0 : AbstractSchemaModificationInterpreter() {
val Arguments.receiver: PluginDataFrameSchema by dataFrame()
val Arguments.whereAllNA: Boolean by arg(defaultValue = Present(false))
val Arguments.columns: ColumnsResolver by arg()
override fun Arguments.interpret(): PluginDataFrameSchema {
if (whereAllNA) return receiver
return PluginDataFrameSchema(fillNullsImpl(receiver.columns(), columns.resolve(receiver).mapTo(mutableSetOf()) { it.path.path }, emptyList()))
}
}
fun KotlinTypeFacade.fillNullsImpl(
columns: List<SimpleCol>,
paths: Set<List<String>>,
p: List<String>
): List<SimpleCol> {
return columns.map {
// else report?
if (p + it.name() in paths && it is SimpleDataColumn) {
val coneType = it.type.type as? ConeSimpleKotlinType
if (coneType != null) {
val type = coneType.withNullability(ConeNullability.NOT_NULL, session.typeContext)
it.changeType(Marker.invoke(type))
} else {
// report?
it
}
} else {
if (it is SimpleColumnGroup) {
val updatedColumns = fillNullsImpl(it.columns(), paths, p + it.name())
SimpleColumnGroup(it.name(), updatedColumns)
} else {
it
}
}
}
}
@@ -0,0 +1,35 @@
package org.jetbrains.kotlinx.dataframe.plugin.impl.api
import org.jetbrains.kotlinx.dataframe.plugin.impl.AbstractInterpreter
import org.jetbrains.kotlinx.dataframe.plugin.impl.AbstractSchemaModificationInterpreter
import org.jetbrains.kotlinx.dataframe.plugin.impl.Arguments
import org.jetbrains.kotlinx.dataframe.plugin.impl.PluginDataFrameSchema
import org.jetbrains.kotlinx.dataframe.plugin.impl.SimpleColumnGroup
import org.jetbrains.kotlinx.dataframe.plugin.impl.groupBy
import org.jetbrains.kotlinx.dataframe.plugin.impl.ignore
import org.jetbrains.kotlinx.dataframe.plugin.impl.makeNullable
class GroupByReducePredicate : AbstractInterpreter<GroupBy>() {
val Arguments.receiver by groupBy()
val Arguments.predicate by ignore()
override fun Arguments.interpret(): GroupBy {
return receiver
}
}
class GroupByReduceExpression : AbstractInterpreter<GroupBy>() {
val Arguments.receiver by groupBy()
val Arguments.rowExpression by ignore()
override fun Arguments.interpret(): GroupBy {
return receiver
}
}
class GroupByReduceInto : AbstractSchemaModificationInterpreter() {
val Arguments.receiver by groupBy()
val Arguments.columnName: String by arg()
override fun Arguments.interpret(): PluginDataFrameSchema {
val group = makeNullable(SimpleColumnGroup(columnName, receiver.groups.columns()))
return PluginDataFrameSchema(receiver.keys.columns() + group)
}
}
@@ -0,0 +1,84 @@
package org.jetbrains.kotlinx.dataframe.plugin.impl.api
import org.jetbrains.kotlinx.dataframe.plugin.extensions.Marker
import org.jetbrains.kotlinx.dataframe.plugin.impl.AbstractInterpreter
import org.jetbrains.kotlinx.dataframe.plugin.impl.AbstractSchemaModificationInterpreter
import org.jetbrains.kotlinx.dataframe.plugin.impl.Arguments
import org.jetbrains.kotlinx.dataframe.plugin.impl.PluginDataFrameSchema
import org.jetbrains.kotlinx.dataframe.plugin.impl.SimpleCol
import org.jetbrains.kotlinx.dataframe.plugin.impl.SimpleColumnGroup
import org.jetbrains.kotlinx.dataframe.plugin.impl.dataFrame
import org.jetbrains.kotlinx.dataframe.plugin.impl.simpleColumnOf
import org.jetbrains.kotlinx.dataframe.plugin.impl.dsl
import org.jetbrains.kotlinx.dataframe.plugin.impl.ignore
import org.jetbrains.kotlinx.dataframe.plugin.impl.type
typealias TypeApproximation = Marker
class Add : AbstractSchemaModificationInterpreter() {
val Arguments.receiver: PluginDataFrameSchema by dataFrame()
val Arguments.name: String by arg()
val Arguments.infer by ignore()
val Arguments.type: TypeApproximation by type(name("expression"))
override fun Arguments.interpret(): PluginDataFrameSchema {
return PluginDataFrameSchema(receiver.columns() + simpleColumnOf(name, type.type))
}
}
class From : AbstractInterpreter<Unit>() {
val Arguments.dsl: AddDslApproximation by arg()
val Arguments.receiver: String by arg()
val Arguments.type: TypeApproximation by type(name("expression"))
override fun Arguments.interpret() {
dsl.columns += simpleColumnOf(receiver, type.type)
}
}
class Into : AbstractInterpreter<Unit>() {
val Arguments.dsl: AddDslApproximation by arg()
val Arguments.receiver: TypeApproximation by type()
val Arguments.name: String by arg()
override fun Arguments.interpret() {
dsl.columns += simpleColumnOf(name, receiver.type)
}
}
class AddDslApproximation(val columns: MutableList<SimpleCol>)
class AddWithDsl : AbstractSchemaModificationInterpreter() {
val Arguments.receiver: PluginDataFrameSchema by dataFrame()
val Arguments.body by dsl()
override fun Arguments.interpret(): PluginDataFrameSchema {
val addDsl = AddDslApproximation(receiver.columns().toMutableList())
body(addDsl, emptyMap())
return PluginDataFrameSchema(addDsl.columns)
}
}
class AddDslStringInvoke : AbstractInterpreter<Unit>() {
val Arguments.dsl: AddDslApproximation by arg()
val Arguments.receiver: String by arg()
val Arguments.body by dsl()
override fun Arguments.interpret() {
val addDsl = AddDslApproximation(mutableListOf())
body(addDsl, emptyMap())
dsl.columns.add(SimpleColumnGroup(receiver, addDsl.columns))
}
}
class AddDslNamedGroup : AbstractInterpreter<Unit>() {
val Arguments.dsl: AddDslApproximation by arg()
val Arguments.name: String by arg()
val Arguments.body by dsl()
override fun Arguments.interpret() {
val addDsl = AddDslApproximation(mutableListOf())
body(addDsl, emptyMap())
dsl.columns.add(SimpleColumnGroup(name, addDsl.columns))
}
}
@@ -0,0 +1,21 @@
package org.jetbrains.kotlinx.dataframe.plugin.impl.api
import org.jetbrains.kotlinx.dataframe.plugin.impl.AbstractSchemaModificationInterpreter
import org.jetbrains.kotlinx.dataframe.plugin.impl.Arguments
import org.jetbrains.kotlinx.dataframe.plugin.impl.PluginDataFrameSchema
import org.jetbrains.kotlinx.dataframe.plugin.impl.Present
import org.jetbrains.kotlinx.dataframe.plugin.impl.dataFrame
import org.jetbrains.kotlinx.dataframe.plugin.impl.simpleColumnOf
class AddId : AbstractSchemaModificationInterpreter() {
val Arguments.receiver: PluginDataFrameSchema by dataFrame()
val Arguments.columnName: String by arg(defaultValue = Present("id"))
override fun Arguments.interpret(): PluginDataFrameSchema {
val columns = buildList {
add(simpleColumnOf(columnName, session.builtinTypes.intType.type))
addAll(receiver.columns())
}
return PluginDataFrameSchema(columns)
}
}
@@ -0,0 +1,79 @@
package org.jetbrains.kotlinx.dataframe.plugin.impl.api
import org.jetbrains.kotlinx.dataframe.ColumnsSelector
import org.jetbrains.kotlinx.dataframe.api.cast
import org.jetbrains.kotlinx.dataframe.api.getColumnsWithPaths
import org.jetbrains.kotlinx.dataframe.columns.ColumnReference
import org.jetbrains.kotlinx.dataframe.columns.ColumnResolutionContext
import org.jetbrains.kotlinx.dataframe.columns.ColumnWithPath
import org.jetbrains.kotlinx.dataframe.columns.SingleColumn
import org.jetbrains.kotlinx.dataframe.columns.UnresolvedColumnsPolicy
import org.jetbrains.kotlinx.dataframe.plugin.impl.AbstractInterpreter
import org.jetbrains.kotlinx.dataframe.plugin.impl.Arguments
import org.jetbrains.kotlinx.dataframe.plugin.impl.ConeTypesAdapter
import org.jetbrains.kotlinx.dataframe.plugin.impl.PluginDataFrameSchema
import org.jetbrains.kotlinx.dataframe.plugin.impl.asDataColumn
import org.jetbrains.kotlinx.dataframe.plugin.impl.asDataFrame
import org.jetbrains.kotlinx.dataframe.plugin.impl.asSimpleColumn
import org.jetbrains.kotlinx.dataframe.plugin.impl.data.ColumnWithPathApproximation
import org.jetbrains.kotlinx.dataframe.plugin.impl.toPluginDataFrameSchema
internal class And10 : AbstractInterpreter<ColumnsResolver>() {
val Arguments.other: ColumnsResolver by arg()
val Arguments.receiver: ColumnsResolver by arg()
override fun Arguments.interpret(): ColumnsResolver {
return object : ColumnsResolver {
override fun resolve(df: PluginDataFrameSchema): List<ColumnWithPathApproximation> {
return receiver.resolve(df) + other.resolve(df)
}
}
}
}
class SingleColumnApproximation(val col: ColumnWithPathApproximation) : ColumnsResolver, SingleColumn<Any?>, ColumnReference<Any?> {
override fun name(): String {
return col.column.name
}
override fun rename(newName: String): ColumnReference<Any?> {
return SingleColumnApproximation(ColumnWithPathApproximation(col.path, col.column.rename(newName)))
}
override fun resolve(df: PluginDataFrameSchema): List<ColumnWithPathApproximation> {
return listOf(col)
}
override fun resolve(context: ColumnResolutionContext): List<ColumnWithPath<Any?>> {
return listOf(resolveSingle(context))
}
override fun resolveSingle(context: ColumnResolutionContext): ColumnWithPath<Any?> {
return ColumnWithPath(col.column.asDataColumn(), col.path)
}
}
interface ColumnsResolver : org.jetbrains.kotlinx.dataframe.columns.ColumnSet<Any?> {
fun resolve(df: PluginDataFrameSchema): List<ColumnWithPathApproximation>
override fun resolve(context: ColumnResolutionContext): List<ColumnWithPath<Any?>> {
val schema = context.df.cast<ConeTypesAdapter>().toPluginDataFrameSchema()
return resolve(schema).map { ColumnWithPath(it.column.asDataColumn(), it.path) }
}
}
abstract class ColumnsResolverAdapter : org.jetbrains.kotlinx.dataframe.columns.ColumnSet<Any?>, ColumnsResolver {
override fun resolve(df: PluginDataFrameSchema): List<ColumnWithPathApproximation> {
return resolve(ColumnResolutionContext(df.asDataFrame(), UnresolvedColumnsPolicy.Skip))
.map { ColumnWithPathApproximation(it.path, it.data.asSimpleColumn()) }
}
}
fun columnsResolver(f: ColumnsSelector<*, *>): ColumnsResolver {
return object : ColumnsResolverAdapter() {
override fun resolve(context: ColumnResolutionContext): List<ColumnWithPath<Any?>> {
return context.df.getColumnsWithPaths(f)
}
}
}
@@ -0,0 +1,200 @@
package org.jetbrains.kotlinx.dataframe.plugin.impl.api
import org.jetbrains.kotlin.fir.declarations.getAnnotationByClassId
import org.jetbrains.kotlin.fir.declarations.getBooleanArgument
import org.jetbrains.kotlin.fir.declarations.getKClassArgument
import org.jetbrains.kotlin.fir.expressions.FirFunctionCall
import org.jetbrains.kotlin.fir.references.toResolvedFunctionSymbol
import org.jetbrains.kotlin.fir.types.ConeNullability
import org.jetbrains.kotlin.fir.types.typeContext
import org.jetbrains.kotlin.fir.types.withNullability
import org.jetbrains.kotlin.name.Name
import org.jetbrains.kotlinx.dataframe.api.asColumn
import org.jetbrains.kotlinx.dataframe.api.convert
import org.jetbrains.kotlinx.dataframe.api.toPath
import org.jetbrains.kotlinx.dataframe.columns.toColumnSet
import org.jetbrains.kotlinx.dataframe.plugin.extensions.KotlinTypeFacade
import org.jetbrains.kotlinx.dataframe.plugin.extensions.wrap
import org.jetbrains.kotlinx.dataframe.plugin.impl.Absent
import org.jetbrains.kotlinx.dataframe.plugin.impl.AbstractInterpreter
import org.jetbrains.kotlinx.dataframe.plugin.impl.AbstractSchemaModificationInterpreter
import org.jetbrains.kotlinx.dataframe.plugin.impl.Arguments
import org.jetbrains.kotlinx.dataframe.plugin.impl.Interpreter
import org.jetbrains.kotlinx.dataframe.plugin.impl.PluginDataFrameSchema
import org.jetbrains.kotlinx.dataframe.plugin.impl.Present
import org.jetbrains.kotlinx.dataframe.plugin.impl.SimpleCol
import org.jetbrains.kotlinx.dataframe.plugin.impl.SimpleColumnGroup
import org.jetbrains.kotlinx.dataframe.plugin.impl.SimpleDataColumn
import org.jetbrains.kotlinx.dataframe.plugin.impl.SimpleFrameColumn
import org.jetbrains.kotlinx.dataframe.plugin.impl.asDataColumn
import org.jetbrains.kotlinx.dataframe.plugin.impl.asDataFrame
import org.jetbrains.kotlinx.dataframe.plugin.impl.dataFrame
import org.jetbrains.kotlinx.dataframe.plugin.impl.ignore
import org.jetbrains.kotlinx.dataframe.plugin.impl.simpleColumnOf
import org.jetbrains.kotlinx.dataframe.plugin.impl.toPluginDataFrameSchema
import org.jetbrains.kotlinx.dataframe.plugin.impl.type
import org.jetbrains.kotlinx.dataframe.plugin.utils.Names
internal class Convert0 : AbstractInterpreter<ConvertApproximation>() {
val Arguments.columns: ColumnsResolver by arg()
val Arguments.receiver: PluginDataFrameSchema by dataFrame()
override val Arguments.startingSchema get() = receiver
override fun Arguments.interpret(): ConvertApproximation {
return ConvertApproximation(receiver, columns.resolve(receiver).map { it.path.path })
}
}
class Convert2 : AbstractInterpreter<ConvertApproximation>() {
val Arguments.receiver: PluginDataFrameSchema by dataFrame()
val Arguments.columns: List<String> by arg(defaultValue = Absent)
override fun Arguments.interpret(): ConvertApproximation {
return ConvertApproximation(receiver, columns.map { listOf(it) })
}
}
class ConvertApproximation(val schema: PluginDataFrameSchema, val columns: List<List<String>>)
internal class Convert6 : AbstractInterpreter<PluginDataFrameSchema>() {
val Arguments.firstCol: String by arg()
val Arguments.cols: List<String> by arg(defaultValue = Present(emptyList()))
val Arguments.infer by ignore()
val Arguments.expression: TypeApproximation by type()
val Arguments.receiver: PluginDataFrameSchema by dataFrame()
override val Arguments.startingSchema get() = receiver
override fun Arguments.interpret(): PluginDataFrameSchema {
val columns = (listOf(firstCol) + cols).map { listOf(it) }
return convertImpl(receiver, columns, expression)
}
}
class With0 : AbstractSchemaModificationInterpreter() {
val Arguments.receiver: ConvertApproximation by arg()
val Arguments.infer by ignore()
val Arguments.type: TypeApproximation by type(name("rowConverter"))
override fun Arguments.interpret(): PluginDataFrameSchema {
return convertImpl(receiver.schema, receiver.columns, type)
}
}
class PerRowCol : AbstractSchemaModificationInterpreter() {
val Arguments.receiver: ConvertApproximation by arg()
val Arguments.infer by ignore()
val Arguments.type: TypeApproximation by type(name("expression"))
override fun Arguments.interpret(): PluginDataFrameSchema {
return convertImpl(receiver.schema, receiver.columns, type)
}
}
internal fun KotlinTypeFacade.convertImpl(
pluginDataFrameSchema: PluginDataFrameSchema,
columns: List<List<String>>,
type: TypeApproximation
): PluginDataFrameSchema {
return pluginDataFrameSchema.map(columns.toSet()) { path, column ->
val unwrappedType = type.type
simpleColumnOf(column.name, unwrappedType)
}
}
internal fun PluginDataFrameSchema.map(selected: ColumnsSet, transform: ColumnMapper): PluginDataFrameSchema {
return PluginDataFrameSchema(
f(columns(), transform, selected, emptyList())
)
}
internal typealias ColumnsSet = Set<List<String>>
internal typealias ColumnMapper = (List<String>, SimpleCol) -> SimpleCol
internal fun f(columns: List<SimpleCol>, transform: ColumnMapper, selected: ColumnsSet, path: List<String>): List<SimpleCol> {
return columns.map {
val fullPath = path + listOf(it.name)
when (it) {
is SimpleColumnGroup -> if (fullPath in selected) {
transform(fullPath, it)
} else {
it.map(transform, selected, fullPath)
}
is SimpleFrameColumn -> if (fullPath in selected) {
transform(fullPath, it)
} else {
it.map(transform, selected, fullPath)
}
is SimpleDataColumn -> if (fullPath in selected) {
transform(path, it)
} else {
it
}
}
}
}
internal fun SimpleColumnGroup.map(transform: ColumnMapper, selected: ColumnsSet, path: List<String>): SimpleColumnGroup {
return SimpleColumnGroup(
name,
f(columns(), transform, selected, path)
)
}
internal fun SimpleFrameColumn.map(transform: ColumnMapper, selected: ColumnsSet, path: List<String>): SimpleFrameColumn {
return SimpleFrameColumn(
name,
f(columns(), transform, selected, path)
)
}
internal class To0 : AbstractInterpreter<PluginDataFrameSchema>() {
val Arguments.receiver: ConvertApproximation by arg()
val Arguments.typeArg0: TypeApproximation by arg()
override val Arguments.startingSchema get() = receiver.schema
override fun Arguments.interpret(): PluginDataFrameSchema {
return convertImpl(receiver.schema, receiver.columns, typeArg0)
}
}
internal class ConvertAsColumn : AbstractSchemaModificationInterpreter() {
val Arguments.receiver: ConvertApproximation by arg()
val Arguments.typeArg2: TypeApproximation by arg()
val Arguments.type: TypeApproximation by type(name("columnConverter"))
override fun Arguments.interpret(): PluginDataFrameSchema {
return receiver.schema.asDataFrame()
.convert { receiver.columns.map { it.toPath() }.toColumnSet() }
.asColumn { simpleColumnOf("", typeArg2.type).asDataColumn() }
.toPluginDataFrameSchema()
}
}
internal abstract class AbstractToSpecificType : AbstractInterpreter<PluginDataFrameSchema>() {
val Arguments.functionCall: FirFunctionCall by arg(lens = Interpreter.Id)
val Arguments.receiver: ConvertApproximation by arg()
override fun Arguments.interpret(): PluginDataFrameSchema {
val converterAnnotation = functionCall.calleeReference.toResolvedFunctionSymbol()?.getAnnotationByClassId(Names.CONVERTER_ANNOTATION, session)
val to = converterAnnotation?.getKClassArgument(Name.identifier("klass"), session)
val nullable = converterAnnotation?.getBooleanArgument(Name.identifier("nullable"), session)
return if (to != null && nullable != null) {
val targetType = to.withNullability(ConeNullability.create(nullable), session.typeContext)
convertImpl(receiver.schema, receiver.columns, targetType.wrap())
} else {
PluginDataFrameSchema.EMPTY
}
}
}
internal class ToSpecificType : AbstractToSpecificType()
internal class ToSpecificTypeZone : AbstractToSpecificType() {
val Arguments.zone by ignore()
}
internal class ToSpecificTypePattern : AbstractToSpecificType() {
val Arguments.pattern by ignore()
val Arguments.locale by ignore()
}
@@ -0,0 +1,19 @@
package org.jetbrains.kotlinx.dataframe.plugin.impl.api
import org.jetbrains.kotlinx.dataframe.plugin.impl.AbstractSchemaModificationInterpreter
import org.jetbrains.kotlinx.dataframe.plugin.impl.Arguments
import org.jetbrains.kotlinx.dataframe.plugin.impl.PluginDataFrameSchema
import org.jetbrains.kotlinx.dataframe.plugin.impl.Present
import org.jetbrains.kotlinx.dataframe.plugin.impl.add
import org.jetbrains.kotlinx.dataframe.plugin.impl.groupBy
import org.jetbrains.kotlinx.dataframe.plugin.impl.ignore
class GroupByCount0 : AbstractSchemaModificationInterpreter() {
val Arguments.receiver by groupBy()
val Arguments.resultName: String by arg(defaultValue = Present("count"))
val Arguments.predicate by ignore()
override fun Arguments.interpret(): PluginDataFrameSchema {
return receiver.keys.add(resultName, session.builtinTypes.intType.type, context = this)
}
}
@@ -0,0 +1,85 @@
package org.jetbrains.kotlinx.dataframe.plugin.impl.api
import org.jetbrains.kotlin.fir.expressions.FirExpression
import org.jetbrains.kotlin.fir.expressions.FirLiteralExpression
import org.jetbrains.kotlin.fir.expressions.FirVarargArgumentsExpression
import org.jetbrains.kotlin.fir.plugin.createConeType
import org.jetbrains.kotlin.fir.types.ConeKotlinType
import org.jetbrains.kotlin.fir.types.classId
import org.jetbrains.kotlin.fir.types.commonSuperTypeOrNull
import org.jetbrains.kotlin.fir.types.resolvedType
import org.jetbrains.kotlin.fir.types.type
import org.jetbrains.kotlin.fir.types.typeContext
import org.jetbrains.kotlinx.dataframe.plugin.impl.AbstractInterpreter
import org.jetbrains.kotlinx.dataframe.plugin.impl.AbstractSchemaModificationInterpreter
import org.jetbrains.kotlinx.dataframe.plugin.impl.Arguments
import org.jetbrains.kotlinx.dataframe.plugin.impl.Interpreter
import org.jetbrains.kotlinx.dataframe.plugin.impl.PluginDataFrameSchema
import org.jetbrains.kotlinx.dataframe.plugin.impl.simpleColumnOf
import org.jetbrains.kotlinx.dataframe.impl.api.withValuesImpl
import org.jetbrains.kotlinx.dataframe.plugin.utils.Names
class DataFrameOf0 : AbstractInterpreter<DataFrameBuilderApproximation>() {
val Arguments.header: List<String> by arg()
override fun Arguments.interpret(): DataFrameBuilderApproximation {
return DataFrameBuilderApproximation(header)
}
}
class DataFrameBuilderApproximation(val header: List<String>)
class DataFrameBuilderInvoke0 : AbstractSchemaModificationInterpreter() {
val Arguments.receiver: DataFrameBuilderApproximation by arg()
val Arguments.values: FirVarargArgumentsExpression by arg(lens = Interpreter.Id)
override fun Arguments.interpret(): PluginDataFrameSchema {
val columns = (receiver.header to values.arguments).withValuesImpl().map { (name, values) ->
val type = session.typeContext.commonSuperTypeOrNull(values.map { it.resolvedType }) ?: error("$name $values")
simpleColumnOf(name, type)
}
return PluginDataFrameSchema(columns)
}
}
class DataFrameOf3 : AbstractSchemaModificationInterpreter() {
val Arguments.columns: List<Interpreter.Success<Pair<*, *>>> by arg()
override fun Arguments.interpret(): PluginDataFrameSchema {
val res = columns.map {
val it = it.value
val name = (it.first as? FirLiteralExpression)?.value as? String
val type = (it.second as? FirExpression)?.resolvedType?.typeArguments?.getOrNull(0)?.type
if (name == null || type == null) return PluginDataFrameSchema(emptyList())
simpleColumnOf(name, type)
}
return PluginDataFrameSchema(res)
}
}
abstract class SchemaConstructor : AbstractSchemaModificationInterpreter() {
val Arguments.columns: List<Interpreter.Success<Pair<*, *>>> by arg()
override fun Arguments.interpret(): PluginDataFrameSchema {
val res = columns.map {
val it = it.value
val name = (it.first as? FirLiteralExpression)?.value as? String
val resolvedType = (it.second as? FirExpression)?.resolvedType
val type: ConeKotlinType? = when (resolvedType?.classId) {
Names.COLUM_GROUP_CLASS_ID -> Names.DATA_ROW_CLASS_ID.createConeType(session, arrayOf(resolvedType.typeArguments[0]))
Names.FRAME_COLUMN_CLASS_ID -> Names.DF_CLASS_ID.createConeType(session, arrayOf(resolvedType.typeArguments[0]))
Names.DATA_COLUMN_CLASS_ID -> resolvedType.typeArguments[0] as? ConeKotlinType
Names.BASE_COLUMN_CLASS_ID -> resolvedType.typeArguments[0] as? ConeKotlinType
else -> null
}
if (name == null || type == null) return PluginDataFrameSchema(emptyList())
simpleColumnOf(name, type)
}
return PluginDataFrameSchema(res)
}
}
class DataFrameOfPairs : SchemaConstructor()
class ColumnOfPairs : SchemaConstructor()
@@ -0,0 +1,85 @@
package org.jetbrains.kotlinx.dataframe.plugin.impl.api
import org.jetbrains.kotlinx.dataframe.plugin.extensions.KotlinTypeFacade
import org.jetbrains.kotlinx.dataframe.plugin.impl.AbstractInterpreter
import org.jetbrains.kotlinx.dataframe.plugin.impl.Arguments
import org.jetbrains.kotlinx.dataframe.plugin.impl.PluginDataFrameSchema
import org.jetbrains.kotlinx.dataframe.plugin.impl.Present
import org.jetbrains.kotlinx.dataframe.plugin.impl.SimpleCol
import org.jetbrains.kotlinx.dataframe.plugin.impl.SimpleColumnGroup
import org.jetbrains.kotlinx.dataframe.plugin.impl.SimpleDataColumn
import org.jetbrains.kotlinx.dataframe.plugin.impl.SimpleFrameColumn
import org.jetbrains.kotlinx.dataframe.plugin.impl.data.ColumnWithPathApproximation
import org.jetbrains.kotlinx.dataframe.plugin.impl.dataFrame
import org.jetbrains.kotlinx.dataframe.plugin.impl.simpleColumnOf
internal class Explode0 : AbstractInterpreter<PluginDataFrameSchema>() {
val Arguments.dropEmpty: Boolean by arg(defaultValue = Present(true))
val Arguments.receiver: PluginDataFrameSchema by dataFrame()
val Arguments.columns: ColumnsResolver? by arg(defaultValue = Present(null))
override val Arguments.startingSchema get() = receiver
override fun Arguments.interpret(): PluginDataFrameSchema {
val columns = columns ?: object : ColumnsResolver {
override fun resolve(df: PluginDataFrameSchema): List<ColumnWithPathApproximation> {
return df.flatten(includeFrames = false).filter {
val column = it.column
column is SimpleFrameColumn || column is SimpleDataColumn && column.type.isList()
}
}
}
return receiver.explodeImpl(dropEmpty, columns.resolve(receiver))
}
}
val KotlinTypeFacade.explodeImpl: PluginDataFrameSchema.(dropEmpty: Boolean, selector: List<ColumnWithPathApproximation>) -> PluginDataFrameSchema
get() = { dropEmpty, selector ->
val columns = selector
val selected = columns.associateBy { it.path }
fun makeNullable(column: SimpleCol): SimpleCol {
return when (column) {
is SimpleColumnGroup -> SimpleColumnGroup(column.name, column.columns().map { makeNullable(it) })
is SimpleFrameColumn -> column
is SimpleDataColumn -> {
column.changeType(type = column.type.changeNullability { nullable -> selector.size > 1 || !dropEmpty || nullable })
}
}
}
fun explode(column: SimpleCol, path: List<String>): SimpleCol {
val fullPath = path + listOf(column.name)
return when (column) {
is SimpleColumnGroup -> {
SimpleColumnGroup(column.name, column.columns().map { explode(it, fullPath) })
}
is SimpleFrameColumn -> {
val s = selected[fullPath]
if (s != null) {
SimpleColumnGroup(s.column.name, column.columns().map { makeNullable(it) })
} else {
column
}
}
is SimpleDataColumn -> {
val s = selected[fullPath]
if (s != null) {
val newType = when {
column.type.isList() -> column.type.typeArgument()
else -> column.type
}
makeNullable(simpleColumnOf(s.column.name, newType.type))
} else {
column
}
}
}
}
PluginDataFrameSchema(
columns().map { column ->
explode(column, emptyList())
}
)
}
@@ -0,0 +1,31 @@
package org.jetbrains.kotlinx.dataframe.plugin.impl.api
import org.jetbrains.kotlinx.dataframe.plugin.impl.AbstractInterpreter
import org.jetbrains.kotlinx.dataframe.plugin.impl.AbstractSchemaModificationInterpreter
import org.jetbrains.kotlinx.dataframe.plugin.impl.Arguments
import org.jetbrains.kotlinx.dataframe.plugin.impl.PluginDataFrameSchema
import org.jetbrains.kotlinx.dataframe.plugin.impl.dataFrame
import org.jetbrains.kotlinx.dataframe.plugin.impl.type
class FillNulls0 : AbstractInterpreter<FillNullsApproximation>() {
val Arguments.receiver: PluginDataFrameSchema by dataFrame()
val Arguments.columns: ColumnsResolver by arg()
override fun Arguments.interpret(): FillNullsApproximation {
return FillNullsApproximation(receiver, columns)
}
}
class FillNullsApproximation(val schema: PluginDataFrameSchema, val columns: ColumnsResolver) : UpdateApproximation
class UpdateWith0 : AbstractSchemaModificationInterpreter() {
val Arguments.receiver: UpdateApproximation by arg()
val Arguments.expression: TypeApproximation by type()
override fun Arguments.interpret(): PluginDataFrameSchema {
return when (val receiver = receiver) {
is FillNullsApproximation -> convertImpl(receiver.schema, receiver.columns.resolve(receiver.schema).map { it.path.path }, expression)
is UpdateApproximationImpl -> convertImpl(receiver.schema, receiver.columns.resolve(receiver.schema).map { it.path.path }, expression)
}
}
}
@@ -0,0 +1,38 @@
package org.jetbrains.kotlinx.dataframe.plugin.impl.api
import org.jetbrains.kotlinx.dataframe.api.flatten
import org.jetbrains.kotlinx.dataframe.api.pathOf
import org.jetbrains.kotlinx.dataframe.columns.toColumnSet
import org.jetbrains.kotlinx.dataframe.plugin.impl.AbstractSchemaModificationInterpreter
import org.jetbrains.kotlinx.dataframe.plugin.impl.Arguments
import org.jetbrains.kotlinx.dataframe.plugin.impl.PluginDataFrameSchema
import org.jetbrains.kotlinx.dataframe.plugin.impl.Present
import org.jetbrains.kotlinx.dataframe.plugin.impl.asDataFrame
import org.jetbrains.kotlinx.dataframe.plugin.impl.dataFrame
import org.jetbrains.kotlinx.dataframe.plugin.impl.toPluginDataFrameSchema
class FlattenDefault : AbstractSchemaModificationInterpreter() {
val Arguments.receiver by dataFrame()
val Arguments.keepParentNameForColumns: Boolean by arg(defaultValue = Present(false))
val Arguments.separator: String by arg(defaultValue = Present("_"))
override fun Arguments.interpret(): PluginDataFrameSchema {
return receiver.asDataFrame().flatten(keepParentNameForColumns, separator).toPluginDataFrameSchema()
}
}
class Flatten0 : AbstractSchemaModificationInterpreter() {
val Arguments.receiver by dataFrame()
val Arguments.keepParentNameForColumns: Boolean by arg(defaultValue = Present(false))
val Arguments.separator: String by arg(defaultValue = Present("_"))
val Arguments.columns: ColumnsResolver by arg()
override fun Arguments.interpret(): PluginDataFrameSchema {
val columns = columns.resolve(receiver).map { it.path }
return receiver
.asDataFrame()
.flatten(keepParentNameForColumns, separator) { columns.toColumnSet() }
.toPluginDataFrameSchema()
}
}
@@ -0,0 +1,35 @@
package org.jetbrains.kotlinx.dataframe.plugin.impl.api
import org.jetbrains.kotlinx.dataframe.plugin.impl.PluginDataFrameSchema
import org.jetbrains.kotlinx.dataframe.plugin.impl.SimpleCol
import org.jetbrains.kotlinx.dataframe.plugin.impl.SimpleDataColumn
import org.jetbrains.kotlinx.dataframe.plugin.impl.SimpleColumnGroup
import org.jetbrains.kotlinx.dataframe.plugin.impl.SimpleFrameColumn
import org.jetbrains.kotlinx.dataframe.plugin.impl.data.ColumnPathApproximation
import org.jetbrains.kotlinx.dataframe.plugin.impl.data.ColumnWithPathApproximation
fun PluginDataFrameSchema.flatten(includeFrames: Boolean): List<ColumnWithPathApproximation> {
if (columns().isEmpty()) return emptyList()
val columns = mutableListOf<ColumnWithPathApproximation>()
flattenImpl(columns(), emptyList(), columns, includeFrames)
return columns
}
fun flattenImpl(columns: List<SimpleCol>, path: List<String>, flatList: MutableList<ColumnWithPathApproximation>, includeFrames: Boolean) {
columns.forEach { column ->
val fullPath = path + listOf(column.name)
when (column) {
is SimpleColumnGroup -> {
flatList.add(ColumnWithPathApproximation(ColumnPathApproximation(fullPath), column))
flattenImpl(column.columns(), fullPath, flatList, includeFrames)
}
is SimpleFrameColumn -> {
flatList.add(ColumnWithPathApproximation(ColumnPathApproximation(fullPath), column))
flattenImpl(column.columns(), fullPath, flatList, includeFrames)
}
is SimpleDataColumn -> {
flatList.add(ColumnWithPathApproximation(ColumnPathApproximation(fullPath), column))
}
}
}
}
@@ -0,0 +1,31 @@
package org.jetbrains.kotlinx.dataframe.plugin.impl.api
import org.jetbrains.kotlinx.dataframe.api.group
import org.jetbrains.kotlinx.dataframe.api.into
import org.jetbrains.kotlinx.dataframe.plugin.impl.AbstractInterpreter
import org.jetbrains.kotlinx.dataframe.plugin.impl.AbstractSchemaModificationInterpreter
import org.jetbrains.kotlinx.dataframe.plugin.impl.Arguments
import org.jetbrains.kotlinx.dataframe.plugin.impl.PluginDataFrameSchema
import org.jetbrains.kotlinx.dataframe.plugin.impl.asDataFrame
import org.jetbrains.kotlinx.dataframe.plugin.impl.dataFrame
import org.jetbrains.kotlinx.dataframe.plugin.impl.toPluginDataFrameSchema
class Group0 : AbstractInterpreter<GroupClauseApproximation>() {
val Arguments.receiver: PluginDataFrameSchema by dataFrame()
val Arguments.columns: ColumnsResolver by arg()
override fun Arguments.interpret(): GroupClauseApproximation {
return GroupClauseApproximation(receiver, columns)
}
}
class GroupClauseApproximation(val df: PluginDataFrameSchema, val columns: ColumnsResolver)
class Into0 : AbstractSchemaModificationInterpreter() {
val Arguments.receiver: GroupClauseApproximation by arg()
val Arguments.column: String by arg()
override fun Arguments.interpret(): PluginDataFrameSchema {
return receiver.df.asDataFrame().group { receiver.columns }.into(column).toPluginDataFrameSchema()
}
}

Some files were not shown because too many files have changed in this diff Show More