Add socket messages callbacks and comms API

Fixes #366

build.gradle.kts

@@ -12,11 +12,6 @@ plugins {
 
 val deploy: Configuration by configurations.creating
 
-deploy.apply {
-    exclude("org.jetbrains.kotlinx", "kotlinx-serialization-json-jvm")
-    exclude("org.jetbrains.kotlinx", "kotlinx-serialization-core-jvm")
-}
-
 ktlint {
     filter {
         exclude("**/org/jetbrains/kotlinx/jupyter/repl.kt")

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

@@ -1,5 +1,6 @@
 package org.jetbrains.kotlinx.jupyter.api
 
+import org.jetbrains.kotlinx.jupyter.api.libraries.JupyterConnection
 import org.jetbrains.kotlinx.jupyter.api.libraries.LibraryResolutionRequest
 
 /**
@@ -99,4 +100,6 @@ interface Notebook {
      * All requests for libraries made during this session
      */
     val libraryRequests: Collection<LibraryResolutionRequest>
+
+    val connection: JupyterConnection
 }

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

@@ -0,0 +1,117 @@
+package org.jetbrains.kotlinx.jupyter.api.libraries
+
+import kotlinx.serialization.json.Json
+import kotlinx.serialization.json.JsonElement
+import kotlinx.serialization.json.JsonObject
+import kotlinx.serialization.json.JsonPrimitive
+import kotlinx.serialization.json.decodeFromJsonElement
+import kotlinx.serialization.json.encodeToJsonElement
+import kotlinx.serialization.json.jsonObject
+import org.jetbrains.kotlinx.jupyter.util.emptyJsonObject
+
+enum class JupyterSocket {
+    HB,
+    SHELL,
+    CONTROL,
+    STDIN,
+    IOPUB;
+}
+
+interface RawMessage {
+    val id: List<ByteArray>
+    val data: JsonElement
+}
+
+val RawMessage.header: JsonObject?
+    get() {
+        if (data !is JsonObject) return null
+        val header = (data as JsonObject)["header"]
+
+        if (header !is JsonObject) return null
+        return header
+    }
+
+val RawMessage.type: String?
+    get() {
+        val myHeader = header
+        if (myHeader !is JsonObject) return null
+        val type = myHeader["msg_type"]
+
+        if (type !is JsonPrimitive || !type.isString) return null
+        return type.content
+    }
+
+val RawMessage.content: JsonObject?
+    get() {
+        if (data !is JsonObject) return null
+        val content = (data as JsonObject)["content"]
+        return (content as? JsonObject)
+    }
+
+typealias CommOpenCallback = (Comm, JsonObject) -> Unit
+typealias CommMsgCallback = (JsonObject) -> Unit
+typealias CommCloseCallback = (JsonObject) -> Unit
+
+interface RawJupyterMessageCallback {
+    val socket: JupyterSocket
+    val messageType: String
+    val action: (RawMessage) -> Unit
+}
+
+interface JupyterConnection {
+    fun addCallback(callback: RawJupyterMessageCallback)
+    fun removeCallback(callback: RawJupyterMessageCallback)
+    fun send(socketName: JupyterSocket, message: RawMessage)
+    fun send(socketName: JupyterSocket, parentMessage: RawMessage, type: String, content: JsonObject, metadata: JsonObject?)
+
+    /**
+     * Creates a comm with a given target, generates unique ID for it. Sends comm_open request to frontend
+     *
+     * @param target Target to create comm for. Should be registered on frontend side.
+     * @param data Content of comm_open message
+     * @return Created comm
+     */
+    fun openComm(target: String, data: JsonObject = emptyJsonObject): Comm
+
+    /**
+     * Closes a comm with a given ID. Sends comm_close request to frontend
+     *
+     * @param id ID of a comm to close
+     * @param data Content of comm_close message
+     */
+    fun closeComm(id: String, data: JsonObject = emptyJsonObject)
+
+    /**
+     * Get all comms for a given target, or all opened comms if `target` is `null`
+     */
+    fun getComms(target: String? = null): Collection<Comm>
+    fun registerTarget(target: String, callback: CommOpenCallback)
+}
+
+interface Comm {
+    val target: String
+    val id: String
+    fun sendJson(data: JsonObject)
+    fun onMessage(action: CommMsgCallback)
+    fun onClose(action: CommCloseCallback)
+
+    // Closes a comm. Sends comm_close request to frontend
+    fun close(data: JsonObject = emptyJsonObject, notifyClient: Boolean = true)
+
+    fun messageReceived(data: JsonObject)
+}
+
+/**
+ * Send an object. `data` should be serializable to JSON object
+ * (generally it means that the corresponding class should be marked with @Serializable)
+ */
+inline fun <reified T> Comm.send(data: T) {
+    sendJson(Json.encodeToJsonElement(data).jsonObject)
+}
+
+inline fun <reified T> Comm.onData(crossinline action: (T) -> Unit) {
+    onMessage { json ->
+        val data = Json.decodeFromJsonElement<T>(json)
+        action(data)
+    }
+}

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

@@ -143,8 +143,12 @@ class SubtypeRendererTypeHandler(private val superType: KClass<*>, override val
     }
 }
 
-inline fun <reified T : Any> createRenderer(crossinline renderAction: (T) -> Any?): RendererTypeHandler {
-    return SubtypeRendererTypeHandler(T::class) { _, result ->
+inline fun <T : Any> createRenderer(kClass: KClass<T>, crossinline renderAction: (T) -> Any?): RendererTypeHandler {
+    return SubtypeRendererTypeHandler(kClass) { _, result ->
         FieldValue(renderAction(result.value as T), result.name)
     }
 }
+
+inline fun <reified T : Any> createRenderer(crossinline renderAction: (T) -> Any?): RendererTypeHandler {
+    return createRenderer(T::class, renderAction)
+}

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

@@ -21,6 +21,8 @@ import org.jetbrains.kotlinx.jupyter.api.libraries.ResourceFallbacksBundle
 import kotlin.reflect.KClass
 import kotlin.reflect.KProperty1
 
+val emptyJsonObject = JsonObject(mapOf())
+
 abstract class PrimitiveStringPropertySerializer<T : Any>(
     kClass: KClass<T>,
     private val prop: KProperty1<T, String>,

src/main/kotlin/org/jetbrains/kotlinx/jupyter/apiImpl.kt

@@ -12,6 +12,7 @@ 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.VariableState
+import org.jetbrains.kotlinx.jupyter.api.libraries.JupyterConnection
 import org.jetbrains.kotlinx.jupyter.api.libraries.LibraryResolutionRequest
 import org.jetbrains.kotlinx.jupyter.repl.impl.SharedReplContext
 
@@ -135,6 +136,7 @@ class EvalData(
 
 class NotebookImpl(
     private val runtimeProperties: ReplRuntimeProperties,
+    override val connection: JupyterConnection
 ) : MutableNotebook {
     private val cells = hashMapOf<Int, MutableCodeCell>()
     override var sharedReplContext: SharedReplContext? = null

src/main/kotlin/org/jetbrains/kotlinx/jupyter/config.kt

@@ -36,7 +36,7 @@ val defaultRuntimeProperties by lazy {
     RuntimeKernelProperties(readResourceAsIniFile("runtime.properties"))
 }
 
-enum class JupyterSockets(val zmqKernelType: SocketType, val zmqClientType: SocketType) {
+enum class JupyterSocketInfo(val zmqKernelType: SocketType, val zmqClientType: SocketType) {
     HB(SocketType.REP, SocketType.REQ),
     SHELL(SocketType.ROUTER, SocketType.REQ),
     CONTROL(SocketType.ROUTER, SocketType.REQ),
@@ -103,7 +103,7 @@ object KernelJupyterParamsSerializer : KSerializer<KernelJupyterParams> {
         return KernelJupyterParams(
             map["signature_scheme"]?.content,
             map["key"]?.content,
-            JupyterSockets.values().map { socket ->
+            JupyterSocketInfo.values().map { socket ->
                 val fieldName = "${socket.nameForUser}_port"
                 map[fieldName]?.let { Json.decodeFromJsonElement<Int>(it) } ?: throw RuntimeException("Cannot find $fieldName in config")
             },
@@ -117,7 +117,7 @@ object KernelJupyterParamsSerializer : KSerializer<KernelJupyterParams> {
             "key" to JsonPrimitive(value.key),
             "transport" to JsonPrimitive(value.transport)
         )
-        JupyterSockets.values().forEach {
+        JupyterSocketInfo.values().forEach {
             map["${it.nameForUser}_port"] = JsonPrimitive(value.ports[it.ordinal])
         }
         utilSerializer.serialize(encoder, map)