org.jetbrains.kotlinx.jupyter.test.testUtil.kt Maven / Gradle / Ivy

Go to download
package org.jetbrains.kotlinx.jupyter.test

import io.kotest.assertions.fail
import jupyter.kotlin.JavaRuntime
import org.jetbrains.kotlinx.jupyter.EvalRequestData
import org.jetbrains.kotlinx.jupyter.MutableCodeCell
import org.jetbrains.kotlinx.jupyter.MutableNotebook
import org.jetbrains.kotlinx.jupyter.ReplForJupyter
import org.jetbrains.kotlinx.jupyter.ReplRuntimeProperties
import org.jetbrains.kotlinx.jupyter.api.AfterCellExecutionCallback
import org.jetbrains.kotlinx.jupyter.api.Code
import org.jetbrains.kotlinx.jupyter.api.CodeCell
import org.jetbrains.kotlinx.jupyter.api.CodePreprocessor
import org.jetbrains.kotlinx.jupyter.api.DisplayContainer
import org.jetbrains.kotlinx.jupyter.api.DisplayResultWithCell
import org.jetbrains.kotlinx.jupyter.api.ExecutionCallback
import org.jetbrains.kotlinx.jupyter.api.ExtensionsProcessor
import org.jetbrains.kotlinx.jupyter.api.FieldsProcessor
import org.jetbrains.kotlinx.jupyter.api.HtmlData
import org.jetbrains.kotlinx.jupyter.api.InterruptionCallback
import org.jetbrains.kotlinx.jupyter.api.JREInfoProvider
import org.jetbrains.kotlinx.jupyter.api.JupyterClientType
import org.jetbrains.kotlinx.jupyter.api.KotlinKernelVersion
import org.jetbrains.kotlinx.jupyter.api.LibraryLoader
import org.jetbrains.kotlinx.jupyter.api.MimeTypedResult
import org.jetbrains.kotlinx.jupyter.api.Notebook
import org.jetbrains.kotlinx.jupyter.api.RenderersProcessor
import org.jetbrains.kotlinx.jupyter.api.ResultsAccessor
import org.jetbrains.kotlinx.jupyter.api.TextRenderersProcessor
import org.jetbrains.kotlinx.jupyter.api.VariableState
import org.jetbrains.kotlinx.jupyter.api.VariableStateImpl
import org.jetbrains.kotlinx.jupyter.api.libraries.ColorScheme
import org.jetbrains.kotlinx.jupyter.api.libraries.ColorSchemeChangedCallback
import org.jetbrains.kotlinx.jupyter.api.libraries.CommManager
import org.jetbrains.kotlinx.jupyter.api.libraries.ExecutionHost
import org.jetbrains.kotlinx.jupyter.api.libraries.JupyterConnection
import org.jetbrains.kotlinx.jupyter.api.libraries.JupyterIntegration
import org.jetbrains.kotlinx.jupyter.api.libraries.LibraryDefinition
import org.jetbrains.kotlinx.jupyter.api.libraries.LibraryReference
import org.jetbrains.kotlinx.jupyter.api.libraries.LibraryResolutionRequest
import org.jetbrains.kotlinx.jupyter.api.libraries.Variable
import org.jetbrains.kotlinx.jupyter.api.libraries.createLibrary
import org.jetbrains.kotlinx.jupyter.api.withId
import org.jetbrains.kotlinx.jupyter.defaultRepositoriesCoordinates
import org.jetbrains.kotlinx.jupyter.defaultRuntimeProperties
import org.jetbrains.kotlinx.jupyter.libraries.AbstractLibraryResolutionInfo
import org.jetbrains.kotlinx.jupyter.libraries.ChainedLibraryResolver
import org.jetbrains.kotlinx.jupyter.libraries.KERNEL_LIBRARIES
import org.jetbrains.kotlinx.jupyter.libraries.LibraryDescriptor
import org.jetbrains.kotlinx.jupyter.libraries.LibraryResolver
import org.jetbrains.kotlinx.jupyter.libraries.parseLibraryDescriptors
import org.jetbrains.kotlinx.jupyter.log
import org.jetbrains.kotlinx.jupyter.messaging.CommManagerImpl
import org.jetbrains.kotlinx.jupyter.messaging.DisplayHandler
import org.jetbrains.kotlinx.jupyter.repl.CompletionResult
import org.jetbrains.kotlinx.jupyter.repl.creating.MockJupyterConnection
import org.jetbrains.kotlinx.jupyter.repl.renderValue
import java.io.File
import kotlin.reflect.KClass
import kotlin.reflect.typeOf
import kotlin.script.experimental.jvm.util.scriptCompilationClasspathFromContext
import kotlin.test.assertEquals

val testDataDir = File("src/test/testData")

const val standardResolverBranch = "master"

val testRepositories = defaultRepositoriesCoordinates

val standardResolverRuntimeProperties = object : ReplRuntimeProperties by defaultRuntimeProperties {
    override val currentBranch: String
        get() = standardResolverBranch
}

val classpath = scriptCompilationClasspathFromContext(
    "lib",
    "api",
    "shared-compiler",
    "kotlin-stdlib",
    "kotlin-reflect",
    "kotlin-script-runtime",
    "kotlinx-serialization-json-jvm",
    "kotlinx-serialization-core-jvm",
    classLoader = TestDisplayHandler::class.java.classLoader,
)

val KClass<*>.classPathEntry: File get() {
    return File(this.java.protectionDomain.codeSource.location.toURI().path)
}

inline fun  classPathEntry(): File {
    return (typeOf().classifier as KClass<*>).classPathEntry
}

val testLibraryResolver: LibraryResolver
    get() = getResolverFromNamesMap(parseLibraryDescriptors(readLibraries()))

fun assertUnit(value: Any?) = assertEquals(Unit, value)

fun assertStartsWith(expectedPrefix: String, actual: String) {
    if (actual.startsWith(expectedPrefix)) return
    val actualStart = actual.substring(0, minOf(expectedPrefix.length, actual.length))
    throw AssertionError("Expected a string to start with '$expectedPrefix', but it starts with '$actualStart")
}

fun Collection>.toLibraries(): LibraryResolver {
    val libJsons = associate { it.first to it.second }
    return getResolverFromNamesMap(parseLibraryDescriptors(libJsons))
}

@JvmName("toLibrariesStringLibraryDefinition")
fun Collection>.toLibraries() = getResolverFromNamesMap(definitions = toMap())

fun getResolverFromNamesMap(
    descriptors: Map? = null,
    definitions: Map? = null,
): LibraryResolver {
    return InMemoryLibraryResolver(
        null,
        descriptors?.mapKeys { entry -> LibraryReference(AbstractLibraryResolutionInfo.Default(), entry.key) },
        definitions?.mapKeys { entry -> LibraryReference(AbstractLibraryResolutionInfo.Default(), entry.key) },
    )
}

fun readLibraries(basePath: String? = null): Map {
    return KERNEL_LIBRARIES.homeLibrariesDir(basePath?.let(::File))
        .listFiles()?.filter(KERNEL_LIBRARIES::isLibraryDescriptor)
        ?.map {
            log.info("Loading '${it.nameWithoutExtension}' descriptor from '${it.canonicalPath}'")
            it.nameWithoutExtension to it.readText()
        }
        .orEmpty()
        .toMap()
}

fun CompletionResult.getOrFail(): CompletionResult.Success = when (this) {
    is CompletionResult.Success -> this
    else -> fail("Result should be success")
}

fun Map.mapToStringValues(): Map {
    return mapValues { it.value.stringValue }
}

fun Map.getStringValue(variableName: String): String? {
    return get(variableName)?.stringValue
}

fun Map.getValue(variableName: String): Any? {
    return get(variableName)?.value?.getOrNull()
}

class InMemoryLibraryResolver(
    parent: LibraryResolver?,
    initialDescriptorsCache: Map? = null,
    initialDefinitionsCache: Map? = null,
) : ChainedLibraryResolver(parent) {
    private val definitionsCache = hashMapOf()
    private val descriptorsCache = hashMapOf()

    init {
        initialDescriptorsCache?.forEach { (key, value) ->
            descriptorsCache[key] = value
        }
        initialDefinitionsCache?.forEach { (key, value) ->
            definitionsCache[key] = value
        }
    }

    override fun shouldResolve(reference: LibraryReference): Boolean {
        return reference.shouldBeCachedInMemory
    }

    override fun tryResolve(reference: LibraryReference, arguments: List): LibraryDefinition? {
        return definitionsCache[reference] ?: descriptorsCache[reference]?.convertToDefinition(arguments)
    }

    override fun save(reference: LibraryReference, definition: LibraryDefinition) {
        definitionsCache[reference] = definition
    }
}

open class TestDisplayHandler(val list: MutableList = mutableListOf()) : DisplayHandler {
    override fun handleDisplay(value: Any, host: ExecutionHost, id: String?) {
        list.add(value)
    }

    override fun handleUpdate(value: Any, host: ExecutionHost, id: String?) {
        TODO("Not yet implemented")
    }
}

class TestDisplayHandlerWithRendering(
    private val notebook: MutableNotebook,
    list: MutableList = mutableListOf(),
) : TestDisplayHandler(list) {

    override fun handleDisplay(value: Any, host: ExecutionHost, id: String?) {
        super.handleDisplay(value, host, id)
        val display = renderValue(notebook, host, value)?.let { if (id != null) it.withId(id) else it } ?: return
        notebook.currentCell?.addDisplay(display)
    }

    override fun handleUpdate(value: Any, host: ExecutionHost, id: String?) {
        super.handleUpdate(value, host, id)
        val display = renderValue(notebook, host, value) ?: return
        val container = notebook.displays
        container.update(id, display)
        container.getById(id).distinctBy { it.cell.id }.forEach {
            it.cell.displays.update(id, display)
        }
    }
}

object NotebookMock : Notebook {
    private val cells = hashMapOf()

    override val cellsList: Collection
        get() = emptyList()
    override val variablesState = mutableMapOf()
    override val cellVariables = mapOf>()
    override val resultsAccessor = ResultsAccessor { getResult(it) }

    override fun getCell(id: Int): MutableCodeCell {
        return cells[id] ?: throw ArrayIndexOutOfBoundsException(
            "There is no cell with number '$id'",
        )
    }

    override fun getResult(id: Int): Any? {
        return getCell(id).result
    }

    override val displays: DisplayContainer
        get() = notImplemented()

    override fun getAllDisplays(): List {
        return displays.getAll()
    }

    override fun getDisplaysById(id: String?): List {
        return displays.getById(id)
    }

    override fun history(before: Int): CodeCell? {
        notImplemented()
    }

    override val currentColorScheme: ColorScheme?
        get() = null

    override fun changeColorScheme(newScheme: ColorScheme) {
        notImplemented()
    }

    override fun renderHtmlAsIFrame(data: HtmlData): MimeTypedResult {
        notImplemented()
    }

    override val kernelVersion: KotlinKernelVersion
        get() = defaultRuntimeProperties.version!!
    override val jreInfo: JREInfoProvider
        get() = JavaRuntime

    override val renderersProcessor: RenderersProcessor
        get() = notImplemented()

    override val textRenderersProcessor: TextRenderersProcessor
        get() = notImplemented()

    override val fieldsHandlersProcessor: FieldsProcessor
        get() = notImplemented()

    override val beforeCellExecutionsProcessor: ExtensionsProcessor>
        get() = notImplemented()

    override val afterCellExecutionsProcessor: ExtensionsProcessor
        get() = notImplemented()

    override val shutdownExecutionsProcessor: ExtensionsProcessor>
        get() = notImplemented()

    override val codePreprocessorsProcessor: ExtensionsProcessor
        get() = notImplemented()
    override val interruptionCallbacksProcessor: ExtensionsProcessor
        get() = notImplemented()
    override val colorSchemeChangeCallbacksProcessor: ExtensionsProcessor
        get() = notImplemented()

    override val libraryRequests: Collection
        get() = notImplemented()

    override val libraryLoader: LibraryLoader
        get() = notImplemented()

    override fun getLibraryFromDescriptor(descriptorText: String, options: Map): LibraryDefinition {
        notImplemented()
    }

    private fun notImplemented(): Nothing {
        error("Not supposed to be called")
    }

    override val jupyterClientType: JupyterClientType
        get() = JupyterClientType.UNKNOWN

    override val connection: JupyterConnection
        get() = MockJupyterConnection

    override val commManager: CommManager
        get() = CommManagerImpl(MockJupyterConnection)
}

fun library(builder: JupyterIntegration.Builder.() -> Unit) = createLibrary(NotebookMock, builder)

fun ReplForJupyter.evalEx(code: Code) = evalEx(EvalRequestData(code))

fun ReplForJupyter.evalRaw(code: Code): Any? {
    return evalEx(code).rawValue
}

fun ReplForJupyter.evalRendered(code: Code): Any? {
    return evalEx(code).renderedValue
}