org.jetbrains.kotlinx.jupyter.codegen.FileAnnotationsProcessorImpl.kt Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of kotlin-jupyter-kernel Show documentation
Kotlin Jupyter kernel published as artifact
There is a newer version: 0.12.0-356
Show newest version
package org.jetbrains.kotlinx.jupyter.codegen

import jupyter.kotlin.CompilerArgs
import jupyter.kotlin.DependsOn
import jupyter.kotlin.KotlinKernelHostProvider
import jupyter.kotlin.Repository
import org.jetbrains.kotlinx.jupyter.api.FileAnnotationCallback
import org.jetbrains.kotlinx.jupyter.api.FileAnnotationHandler
import org.jetbrains.kotlinx.jupyter.api.KotlinKernelHost
import org.jetbrains.kotlinx.jupyter.compiler.CompilerArgsConfigurator
import org.jetbrains.kotlinx.jupyter.dependencies.ScriptDependencyAnnotationHandler
import org.jetbrains.kotlinx.jupyter.exceptions.LibraryProblemPart
import org.jetbrains.kotlinx.jupyter.exceptions.rethrowAsLibraryException
import org.jetbrains.kotlinx.jupyter.repl.impl.JupyterCompiler
import kotlin.script.experimental.api.ResultWithDiagnostics
import kotlin.script.experimental.api.ScriptCollectedData
import kotlin.script.experimental.api.ScriptCompilationConfiguration
import kotlin.script.experimental.api.ScriptConfigurationRefinementContext
import kotlin.script.experimental.api.asDiagnostics
import kotlin.script.experimental.api.asSuccess
import kotlin.script.experimental.api.foundAnnotations
import kotlin.script.experimental.api.valueOr

class FileAnnotationsProcessorImpl(
    val dependencyAnnotationsHandler: ScriptDependencyAnnotationHandler,
    val compilerArgsConfigurator: CompilerArgsConfigurator,
    val compiler: JupyterCompiler,
    val kernelHostProvider: KotlinKernelHostProvider
) : FileAnnotationsProcessor {

    private val handlers = mutableMapOf()

    override fun register(handler: FileAnnotationHandler) {
        handlers[handler.annotation.qualifiedName!!] = handler.callback
        compiler.updateCompilationConfigOnAnnotation(handler) { context ->
            process(context, kernelHostProvider.host!!)
        }
    }

    override fun process(
        context: ScriptConfigurationRefinementContext,
        host: KotlinKernelHost
    ): ResultWithDiagnostics {
        val collected = mutableListOf()

        fun processCollected(conf: ScriptCompilationConfiguration): ResultWithDiagnostics {
            if (collected.isEmpty()) return conf.asSuccess()
            val clazz = collected[0].annotationClass
            return try {
                when (clazz) {
                    DependsOn::class, Repository::class -> dependencyAnnotationsHandler.configure(
                        conf,
                        collected
                    )
                    CompilerArgs::class -> compilerArgsConfigurator.configure(conf, collected)
                    else -> {
                        handlers[clazz.qualifiedName!!]?.let { handler ->
                            rethrowAsLibraryException(LibraryProblemPart.FILE_ANNOTATIONS) {
                                handler.invoke(host, collected)
                            }
                        }
                        conf.asSuccess()
                    }
                }
            } catch (e: Throwable) {
                ResultWithDiagnostics.Failure(e.asDiagnostics(path = context.script.locationId))
            } finally {
                collected.clear()
            }
        }

        var config = context.compilationConfiguration

        val foundAnnotations = context.collectedData?.get(ScriptCollectedData.foundAnnotations)

        foundAnnotations?.forEach {
            if (collected.isEmpty() || collected[0].annotationClass == it.annotationClass) collected.add(it)
            else {
                config = processCollected(config).valueOr { return it }
                collected.add(it)
            }
        }
        return processCollected(config)
    }
}