Add support for variables declarations in library API

Fixes #206

Author: ileasile

jupyter-lib/api/src/main/kotlin/org/jetbrains/kotlinx/jupyter/api/KotlinKernelHost.kt

@@ -34,4 +34,9 @@ interface KotlinKernelHost {
      * Adds a new library via its definition. Fully interchangeable with `%use` approach
      */
     fun addLibrary(library: LibraryDefinition)
+
+    /**
+     * Declares the given variables. They may be used in the notebook cells
+     */
+    fun declareVariables(variables: Iterable<VariableDeclaration>)
 }

jupyter-lib/api/src/main/kotlin/org/jetbrains/kotlinx/jupyter/api/fieldsHandling.kt

@@ -3,6 +3,7 @@ package org.jetbrains.kotlinx.jupyter.api
 import kotlin.reflect.KClass
 import kotlin.reflect.KProperty
 import kotlin.reflect.KType
+import kotlin.reflect.full.createType
 import kotlin.reflect.full.isSubtypeOf
 import kotlin.reflect.full.starProjectedType
 
@@ -36,3 +37,25 @@ class FieldHandlerByClass(
 ) : FieldHandler {
     override fun acceptsType(type: KType) = type.isSubtypeOf(kClass.starProjectedType)
 }
+
+data class VariableDeclaration(
+    val name: VariableName,
+    val value: Any?,
+    val type: KType,
+    val isMutable: Boolean = false,
+) {
+    constructor(
+        name: VariableName,
+        value: Any?,
+        isMutable: Boolean = false
+    ) : this(
+        name,
+        value,
+        value?.let { it::class.starProjectedType } ?: Any::class.createType(nullable = true),
+        isMutable
+    )
+}
+
+fun KotlinKernelHost.declareVariable(variable: VariableDeclaration) = declareVariables(listOf(variable))
+fun KotlinKernelHost.declareVariable(name: VariableName, value: Any?) = declareVariable(VariableDeclaration(name, value))
+fun KotlinKernelHost.declareVariables(variables: Map<VariableName, Any?>) = declareVariables(variables.map { VariableDeclaration(it.key, it.value) })

src/main/kotlin/org/jetbrains/kotlinx/jupyter/repl/impl/CellExecutorImpl.kt

@@ -6,6 +6,7 @@ import org.jetbrains.kotlinx.jupyter.api.ExecutionCallback
 import org.jetbrains.kotlinx.jupyter.api.FieldValue
 import org.jetbrains.kotlinx.jupyter.api.HTML
 import org.jetbrains.kotlinx.jupyter.api.KotlinKernelHost
+import org.jetbrains.kotlinx.jupyter.api.VariableDeclaration
 import org.jetbrains.kotlinx.jupyter.api.libraries.CodeExecution
 import org.jetbrains.kotlinx.jupyter.api.libraries.ExecutionHost
 import org.jetbrains.kotlinx.jupyter.api.libraries.LibraryDefinition
@@ -20,6 +21,9 @@ import org.jetbrains.kotlinx.jupyter.repl.CellExecutor
 import org.jetbrains.kotlinx.jupyter.repl.ExecutionStartedCallback
 import org.jetbrains.kotlinx.jupyter.repl.InternalEvalResult
 import java.util.LinkedList
+import kotlin.reflect.KMutableProperty1
+import kotlin.reflect.full.declaredMemberProperties
+import kotlin.reflect.full.withNullability
 
 interface BaseKernelHost {
     fun <T> withHost(currentHost: KotlinKernelHost, callback: () -> T): T
@@ -161,5 +165,36 @@ internal class CellExecutorImpl(private val replContext: SharedReplContext) : Ce
         override fun <T> execute(callback: KotlinKernelHost.() -> T): T {
             return callback(ExecutionContext(sharedContext, displayHandler, executor))
         }
+
+        override fun declareVariables(variables: Iterable<VariableDeclaration>) {
+            val tempDeclarations = variables.joinToString(
+                "\n",
+                "object $TEMP_OBJECT_NAME {\n",
+                "\n}\n$TEMP_OBJECT_NAME"
+            ) {
+                it.tempDeclaration
+            }
+            val result = execute(tempDeclarations).value as Any
+            val resultClass = result::class
+            val properties = resultClass.declaredMemberProperties.associate {
+                @Suppress("UNCHECKED_CAST")
+                it.name to (it as KMutableProperty1<Any, Any?>)
+            }
+
+            val declarations = variables.joinToString("\n") {
+                val prop = properties[it.name]
+                prop?.set(result, it.value)
+                it.declaration
+            }
+            execute(declarations)
+        }
+
+        companion object {
+            private const val TEMP_OBJECT_NAME = "___temp_declarations"
+
+            private val VariableDeclaration.mutabilityQualifier get() = if (isMutable) "var" else "val"
+            private val VariableDeclaration.declaration get() = """$mutabilityQualifier `$name`: $type = $TEMP_OBJECT_NAME.`$name` as $type"""
+            private val VariableDeclaration.tempDeclaration get() = """var `$name`: ${type.withNullability(true)} = null"""
+        }
     }
 }

src/test/kotlin/org/jetbrains/kotlinx/jupyter/test/repl/CustomLibraryResolverTests.kt

@@ -6,6 +6,9 @@ import kotlinx.serialization.SerializationException
 import org.jetbrains.kotlinx.jupyter.ReplForJupyter
 import org.jetbrains.kotlinx.jupyter.ReplForJupyterImpl
 import org.jetbrains.kotlinx.jupyter.api.KotlinKernelVersion.Companion.toMaybeUnspecifiedString
+import org.jetbrains.kotlinx.jupyter.api.VariableDeclaration
+import org.jetbrains.kotlinx.jupyter.api.declareVariable
+import org.jetbrains.kotlinx.jupyter.api.declareVariables
 import org.jetbrains.kotlinx.jupyter.api.libraries.LibraryDefinition
 import org.jetbrains.kotlinx.jupyter.api.libraries.ResourceType
 import org.jetbrains.kotlinx.jupyter.config.defaultRepositories
@@ -25,6 +28,7 @@ import org.junit.jupiter.api.Test
 import org.junit.jupiter.api.assertDoesNotThrow
 import org.junit.jupiter.api.assertThrows
 import java.io.File
+import kotlin.reflect.typeOf
 import kotlin.test.assertEquals
 import kotlin.test.assertNull
 import kotlin.test.assertTrue
@@ -424,4 +428,37 @@ class CustomLibraryResolverTests : AbstractReplTest() {
         }
         assertEquals(LibraryProblemPart.BEFORE_CELL_CALLBACKS, e.part)
     }
+
+    @Test
+    @ExperimentalStdlibApi
+    fun testLibraryProperties() {
+        val mutProp = arrayListOf(1)
+
+        val repl = testOneLibUsage(
+            library {
+                onLoaded {
+                    declareVariables(
+                        mapOf(
+                            "x1" to 22,
+                            "x2" to 20
+                        )
+                    )
+
+                    declareVariable(
+                        VariableDeclaration("x3", mutProp, typeOf<ArrayList<Int>>())
+                    )
+                }
+            }
+        )
+
+        val result = repl.eval(
+            """
+            x3.add(2)
+            x1 + x2
+            """.trimIndent()
+        ).resultValue
+
+        assertEquals(42, result)
+        assertEquals(listOf(1, 2), mutProp)
+    }
 }