Add new message type for vars serialization

Author: nikolay-egorov

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

@@ -23,6 +23,7 @@ import kotlinx.serialization.json.decodeFromJsonElement
 import kotlinx.serialization.json.encodeToJsonElement
 import kotlinx.serialization.json.jsonObject
 import kotlinx.serialization.serializer
+import org.jetbrains.kotlinx.jupyter.compiler.util.SerializedVariablesState
 import org.jetbrains.kotlinx.jupyter.config.LanguageInfo
 import org.jetbrains.kotlinx.jupyter.exceptions.ReplException
 import kotlin.reflect.KClass
@@ -87,7 +88,10 @@ enum class MessageType(val contentClass: KClass<out MessageContent>) {
     COMM_CLOSE(CommClose::class),
 
     LIST_ERRORS_REQUEST(ListErrorsRequest::class),
-    LIST_ERRORS_REPLY(ListErrorsReply::class);
+    LIST_ERRORS_REPLY(ListErrorsReply::class),
+
+    SERIALIZATION_REQUEST(SerializationRequest::class),
+    SERIALIZATION_REPLY(SerializationReply::class);
 
     // TODO: add custom commands
     // this custom message should be supported on client-side. either JS or Idea Plugin
@@ -555,6 +559,18 @@ class ListErrorsReply(
     val errors: List<ScriptDiagnostic>
 ) : MessageContent()
 
+@Serializable
+class SerializationRequest(
+    val cellId: Int,
+    val descriptorsState: Map<String, SerializedVariablesState>
+) : MessageContent()
+
+@Serializable
+class SerializationReply(
+    val cellId: Int,
+    val descriptorsState: Map<String, SerializedVariablesState> = emptyMap()
+) : MessageContent()
+
 @Serializable(MessageDataSerializer::class)
 data class MessageData(
     val header: MessageHeader? = null,

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

@@ -320,6 +320,13 @@ fun JupyterConnection.Socket.shellMessagesHandler(msg: Message, repl: ReplForJup
                 }
             }
         }
+        is SerializationRequest -> {
+            GlobalScope.launch(Dispatchers.Default) {
+                repl.serializeVariables(content.cellId, content.descriptorsState) { result ->
+                    sendWrapped(msg, makeReplyMessage(msg, MessageType.SERIALIZATION_REPLY, content = result))
+                }
+            }
+        }
         is IsCompleteRequest -> {
             // We are in console mode, so switch off all the loggers
             if (mainLoggerLevel() != Level.OFF) disableLogging()

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

@@ -27,6 +27,7 @@ import org.jetbrains.kotlinx.jupyter.compiler.ScriptImportsCollector
 import org.jetbrains.kotlinx.jupyter.compiler.util.Classpath
 import org.jetbrains.kotlinx.jupyter.compiler.util.EvaluatedSnippetMetadata
 import org.jetbrains.kotlinx.jupyter.compiler.util.SerializedCompiledScriptsData
+import org.jetbrains.kotlinx.jupyter.compiler.util.SerializedVariablesState
 import org.jetbrains.kotlinx.jupyter.config.catchAll
 import org.jetbrains.kotlinx.jupyter.config.getCompilationConfiguration
 import org.jetbrains.kotlinx.jupyter.dependencies.JupyterScriptDependenciesResolverImpl
@@ -128,6 +129,8 @@ interface ReplForJupyter {
 
     suspend fun listErrors(code: Code, callback: (ListErrorsResult) -> Unit)
 
+    suspend fun serializeVariables(cellId: Int, descriptorsState: Map<String, SerializedVariablesState>, callback: (SerializationReply) -> Unit)
+
     val homeDir: File?
 
     val currentClasspath: Collection<String>
@@ -533,6 +536,20 @@ class ReplForJupyterImpl(
         return ListErrorsResult(args.code, errorsList)
     }
 
+    private val serializationQueue = LockQueue<SerializationReply, SerializationArgs>()
+    override suspend fun serializeVariables(cellId: Int, descriptorsState: Map<String, SerializedVariablesState>, callback: (SerializationReply) -> Unit) {
+        doWithLock(SerializationArgs(cellId, descriptorsState, callback), serializationQueue, SerializationReply(cellId, descriptorsState), ::doSerializeVariables)
+    }
+
+    private fun doSerializeVariables(args: SerializationArgs): SerializationReply {
+        val resultMap = mutableMapOf<String, SerializedVariablesState>()
+        args.descriptorsState.forEach { (name, state) ->
+            resultMap[name] = variablesSerializer.doIncrementalSerialization(args.cellId - 1, name, state)
+        }
+        return SerializationReply(args.cellId, resultMap)
+    }
+
+
     private fun <T, Args : LockQueueArgs<T>> doWithLock(
         args: Args,
         queue: LockQueue<T, Args>,
@@ -565,6 +582,12 @@ class ReplForJupyterImpl(
     private data class ListErrorsArgs(val code: String, override val callback: (ListErrorsResult) -> Unit) :
         LockQueueArgs<ListErrorsResult>
 
+    private data class SerializationArgs(
+        val cellId: Int,
+        val descriptorsState: Map<String, SerializedVariablesState>,
+        override val callback: (SerializationReply) -> Unit
+    ) : LockQueueArgs<SerializationReply>
+
     @JvmInline
     private value class LockQueue<T, Args : LockQueueArgs<T>>(
         private val args: AtomicReference<Args?> = AtomicReference()

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

@@ -955,4 +955,47 @@ class ReplVarsSerializationTest : AbstractSingleReplTest() {
             assertEquals("${values++}", state.value)
         }
     }
+
+    @Test
+    fun testSerializationMessage() {
+        val res = eval(
+            """
+            val x = listOf(listOf(1), listOf(2), listOf(3), listOf(4))
+            """.trimIndent(),
+            jupyterId = 1
+        )
+        val varsData = res.metadata.evaluatedVariablesState
+        assertEquals(1, varsData.size)
+        val listData = varsData["x"]!!
+        assertTrue(listData.isContainer)
+        val actualContainer = listData.fieldDescriptor.entries.first().value!!
+        val propertyName = listData.fieldDescriptor.entries.first().key
+
+        runBlocking {
+            repl.serializeVariables(1, mapOf(propertyName to actualContainer)) { result ->
+                val data = result.descriptorsState
+                assertTrue(data.isNotEmpty())
+
+                val innerList = data.entries.last().value!!
+                assertTrue(innerList.isContainer)
+                var receivedDescriptor = innerList.fieldDescriptor
+                assertEquals(2, receivedDescriptor.size)
+                receivedDescriptor = receivedDescriptor.entries.last().value!!.fieldDescriptor
+
+                assertEquals(5, receivedDescriptor.size)
+                var values = 1
+                receivedDescriptor.forEach { (name, state) ->
+                    if (name == "size") {
+                        assertFalse(state!!.isContainer)
+                        assertTrue(state!!.fieldDescriptor.isEmpty())
+                        return@forEach
+                    }
+                    val fieldDescriptor = state!!.fieldDescriptor
+                    assertEquals(0, fieldDescriptor.size)
+                    assertTrue(state.isContainer)
+                    assertEquals("${values++}", state.value)
+                }
+            }
+        }
+    }
 }