org.jetbrains.dataframe.ksp.ExtensionsGenerator.kt Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of symbol-processor-all Show documentation
Annotation preprocessor for DataFrame
There is a newer version: 1727
package org.jetbrains.dataframe.ksp

import com.google.devtools.ksp.getVisibility
import com.google.devtools.ksp.processing.CodeGenerator
import com.google.devtools.ksp.processing.Dependencies
import com.google.devtools.ksp.processing.KSPLogger
import com.google.devtools.ksp.processing.Resolver
import com.google.devtools.ksp.symbol.ClassKind
import com.google.devtools.ksp.symbol.KSAnnotated
import com.google.devtools.ksp.symbol.KSClassDeclaration
import com.google.devtools.ksp.symbol.KSClassifierReference
import com.google.devtools.ksp.symbol.KSDeclaration
import com.google.devtools.ksp.symbol.KSFile
import com.google.devtools.ksp.symbol.KSName
import com.google.devtools.ksp.symbol.KSTypeReference
import com.google.devtools.ksp.symbol.KSValueArgument
import com.google.devtools.ksp.symbol.Modifier
import com.google.devtools.ksp.symbol.Visibility
import com.google.devtools.ksp.validate
import org.jetbrains.kotlinx.dataframe.codeGen.MarkerVisibility
import java.io.IOException
import java.io.OutputStreamWriter

class ExtensionsGenerator(
    private val resolver: Resolver,
    private val codeGenerator: CodeGenerator,
    private val logger: KSPLogger,
) {
    private companion object {
        val EXPECTED_VISIBILITIES = setOf(Visibility.PUBLIC, Visibility.INTERNAL)
    }

    fun resolveDataSchemaDeclarations(): Pair, List> {
        val dataSchemaAnnotation = resolver.getKSNameFromString(DataFrameNames.DATA_SCHEMA)
        val symbols = resolver.getSymbolsWithAnnotation(dataSchemaAnnotation.asString())

        val (validDeclarations, invalidDeclarations) = symbols
            .filterIsInstance()
            .partition { it.validate() }

        val preprocessedDeclarations = validDeclarations
            .asSequence()
            .mapNotNull { it.toDataSchemaDeclarationOrNull() }

        return Pair(preprocessedDeclarations, invalidDeclarations)
    }

    class DataSchemaDeclaration(val origin: KSClassDeclaration, val properties: List)

    class KSAnnotatedWithType(
        private val declaration: KSAnnotated,
        val simpleName: KSName,
        val type: KSTypeReference,
    ) : KSAnnotated by declaration

    private fun KSClassDeclaration.toDataSchemaDeclarationOrNull(): DataSchemaDeclaration? =
        when {
            isClassOrInterface() && effectivelyPublicOrInternal() -> {
                DataSchemaDeclaration(
                    origin = this,
                    properties = getAllProperties()
                        .map { KSAnnotatedWithType(it, it.simpleName, it.type) }
                        .toList(),
                )
            }

            else -> null
        }

    private fun KSClassDeclaration.isClassOrInterface() =
        classKind == ClassKind.INTERFACE || classKind == ClassKind.CLASS

    private fun KSClassDeclaration.effectivelyPublicOrInternal(): Boolean =
        effectivelyPublicOrInternalOrNull(dataSchema = this) != null

    private fun KSDeclaration.effectivelyPublicOrInternalOrNull(dataSchema: KSClassDeclaration): Visibility? {
        val visibility = getVisibility()
        if (visibility !in EXPECTED_VISIBILITIES) {
            val message = buildString {
                append(
                    "DataSchema declaration ${dataSchema.nameString} at ${dataSchema.location} should be $EXPECTED_VISIBILITIES",
                )
                if (this@effectivelyPublicOrInternalOrNull != dataSchema) {
                    append(", but it's parent $nameString is $visibility")
                } else {
                    append("but is $visibility")
                }
            }
            logger.error(message)
            return null
        }

        return when (val parentDeclaration = parentDeclaration) {
            null -> visibility

            else -> when (parentDeclaration.effectivelyPublicOrInternalOrNull(dataSchema)) {
                Visibility.PUBLIC -> visibility
                Visibility.INTERNAL -> Visibility.INTERNAL
                null -> null
                else -> null
            }
        }
    }

    private val KSDeclaration.nameString get() = (qualifiedName ?: simpleName).asString()

    fun generateExtensions(file: KSFile, dataSchema: KSClassDeclaration, properties: List) {
        val packageName = file.packageName.asString()
        val fileName = getFileName(dataSchema)
        val generatedFile = codeGenerator.createNewFile(Dependencies(false, file), packageName, fileName)
        try {
            generatedFile.writer().use {
                it.appendLine("""@file:Suppress("UNCHECKED_CAST", "USELESS_CAST")""")
                if (packageName.isNotEmpty()) {
                    it.appendLine("package $packageName")
                }
                it.writeImports()
                val extensions = renderExtensions(
                    declaration = dataSchema,
                    interfaceName = dataSchema.getQualifiedNameOrThrow(),
                    visibility = getMarkerVisibility(dataSchema),
                    properties = properties.map { property ->
                        Property(getColumnName(property), property.simpleName.asString(), property.type)
                    },
                )
                it.appendLine(extensions)
            }
        } catch (e: IOException) {
            throw IOException("Error writing $fileName generated from declaration at ${file.location}", e)
        }
    }

    private fun OutputStreamWriter.writeImports() {
        appendLine("import org.jetbrains.kotlinx.dataframe.annotations.*")
        appendLine("import org.jetbrains.kotlinx.dataframe.ColumnsContainer")
        appendLine("import org.jetbrains.kotlinx.dataframe.DataColumn")
        appendLine("import org.jetbrains.kotlinx.dataframe.DataFrame")
        appendLine("import org.jetbrains.kotlinx.dataframe.DataRow")
        appendLine("import org.jetbrains.kotlinx.dataframe.columns.ColumnGroup")
        appendLine()
    }

    private fun getFileName(dataSchema: KSClassDeclaration, suffix: String = "Extensions") =
        if (dataSchema.isTopLevel) {
            val simpleName = dataSchema.simpleName.asString()
            "$simpleName${'$'}$suffix"
        } else {
            val fqName = dataSchema.getQualifiedNameOrThrow()
            "${fqName}${'$'}$suffix"
        }

    private val KSDeclaration.isTopLevel get() = parentDeclaration == null

    private fun getMarkerVisibility(dataSchema: KSClassDeclaration) =
        when (val visibility = dataSchema.getVisibility()) {
            Visibility.PUBLIC ->
                if (dataSchema.modifiers.contains(Modifier.PUBLIC)) {
                    MarkerVisibility.EXPLICIT_PUBLIC
                } else {
                    MarkerVisibility.IMPLICIT_PUBLIC
                }

            Visibility.INTERNAL ->
                MarkerVisibility.INTERNAL

            Visibility.PRIVATE, Visibility.PROTECTED, Visibility.LOCAL, Visibility.JAVA_PACKAGE ->
                error("DataSchema declaration should have $EXPECTED_VISIBILITIES, but was $visibility")
        }

    private fun getColumnName(property: KSAnnotatedWithType): String {
        val columnNameAnnotation = property.annotations.firstOrNull { annotation ->
            val annotationType = annotation.annotationType

            val typeIsColumnNameOrNull = (annotationType.element as? KSClassifierReference)
                ?.referencedName()
                ?.let { it == DataFrameNames.SHORT_COLUMN_NAME } != false

            val declarationIsColumnName = annotationType
                .resolve()
                .declaration
                .qualifiedName
                ?.asString() == DataFrameNames.COLUMN_NAME

            typeIsColumnNameOrNull && declarationIsColumnName
        }
        return if (columnNameAnnotation != null) {
            when (val arg = columnNameAnnotation.arguments.singleOrNull()) {
                null -> argumentMismatchError(property, columnNameAnnotation.arguments)
                else -> (arg.value as? String) ?: typeMismatchError(property, arg)
            }
        } else {
            property.simpleName.asString()
        }
    }

    private fun typeMismatchError(property: KSAnnotatedWithType, arg: KSValueArgument): Nothing {
        error(
            "Expected one argument of type String in annotation ColumnName on property ${property.simpleName}, but got ${arg.value}",
        )
    }

    private fun argumentMismatchError(property: KSAnnotatedWithType, args: List): Nothing {
        error(
            "Expected one argument of type String in annotation ColumnName on property ${property.simpleName}, but got $args",
        )
    }
}