Add on-shutdown codes for library descriptors

ileasile · ileasile · commit 154cdbce6fff · 2020-07-17T08:36:14.000+03:00
Fixes #87
diff --git a/README.md b/README.md
@@ -196,6 +196,7 @@ Library descriptor is a `<libName>.json` file with the following fields:
 - `imports`: a list of default imports for library
 - `init`: a list of code snippets to be executed when library is included
 - `initCell`: a list of code snippets to be executed before execution of any cell
+- `shutdown`: a list of code snippets to be executed on kernel shutdown. Any cleanup code goes here
 - `renderers`: a list of type converters for special rendering of particular types
 
 *All fields are optional
diff --git a/src/main/kotlin/org/jetbrains/kotlin/jupyter/config.kt b/src/main/kotlin/org/jetbrains/kotlin/jupyter/config.kt
@@ -129,6 +129,7 @@ open class LibraryDefinition(
         val imports: List<String>,
         val repositories: List<String>,
         val init: List<String>,
+        val shutdown: List<String>,
         val renderers: List<TypeHandler>,
         val converters: List<TypeHandler>,
         val annotations: List<TypeHandler>
@@ -140,10 +141,11 @@ class LibraryDescriptor(dependencies: List<String>,
                         imports: List<String>,
                         repositories: List<String>,
                         init: List<String>,
+                        shutdown: List<String>,
                         renderers: List<TypeHandler>,
                         converters: List<TypeHandler>,
                         annotations: List<TypeHandler>,
-                        val link: String?) : LibraryDefinition(dependencies, initCell, imports, repositories, init, renderers, converters, annotations)
+                        val link: String?) : LibraryDefinition(dependencies, initCell, imports, repositories, init, shutdown, renderers, converters, annotations)
 
 data class ResolverConfig(val repositories: List<RepositoryCoordinates>,
                           val libraries: Deferred<Map<String, LibraryDescriptor>>)
@@ -339,6 +341,7 @@ fun parserLibraryDescriptors(libJsons: Map<String, JsonObject>): Map<String, Lib
                 imports = it.value.array<String>("imports")?.toList().orEmpty(),
                 repositories = it.value.array<String>("repositories")?.toList().orEmpty(),
                 init = it.value.array<String>("init")?.toList().orEmpty(),
+                shutdown = it.value.array<String>("shutdown")?.toList().orEmpty(),
                 initCell = it.value.array<String>("initCell")?.toList().orEmpty(),
                 renderers = it.value.obj("renderers")?.map {
                     TypeHandler(it.key, it.value.toString())
diff --git a/src/main/kotlin/org/jetbrains/kotlin/jupyter/libraries.kt b/src/main/kotlin/org/jetbrains/kotlin/jupyter/libraries.kt
@@ -38,6 +38,7 @@ class LibrariesProcessor(private val libraries: Deferred<Map<String, LibraryDesc
             repositories = library.repositories.map { replaceVariables(it, mapping) },
             imports = library.imports.map { replaceVariables(it, mapping) },
             init = library.init.map { replaceVariables(it, mapping) },
+            shutdown = library.shutdown.map { replaceVariables(it, mapping) },
             initCell = library.initCell.map { replaceVariables(it, mapping) },
             renderers = library.renderers.map { TypeHandler(it.className, replaceVariables(it.code, mapping)) },
             converters = library.converters.map { TypeHandler(it.className, replaceVariables(it.code, mapping)) },
diff --git a/src/main/kotlin/org/jetbrains/kotlin/jupyter/protocol.kt b/src/main/kotlin/org/jetbrains/kotlin/jupyter/protocol.kt
@@ -99,6 +99,7 @@ fun JupyterConnection.Socket.shellMessagesHandler(msg: Message, repl: ReplForJup
             send(makeReplyMessage(msg, "interrupt_reply", content = msg.content))
         }
         "shutdown_request" -> {
+            repl?.evalOnShutdown()
             send(makeReplyMessage(msg, "shutdown_reply", content = msg.content))
             exitProcess(0)
         }
diff --git a/src/main/kotlin/org/jetbrains/kotlin/jupyter/repl.kt b/src/main/kotlin/org/jetbrains/kotlin/jupyter/repl.kt
@@ -63,6 +63,8 @@ typealias Code = String
 interface ReplForJupyter {
     fun eval(code: Code, displayHandler: ((Any) -> Unit)? = null, jupyterId: Int = -1): EvalResult
 
+    fun evalOnShutdown(): List<EvalResult>
+
     fun checkComplete(code: Code): CheckResult
 
     suspend fun complete(code: String, cursor: Int, callback: (CompletionResult) -> Unit)
@@ -110,6 +112,7 @@ class ReplForJupyterImpl(private val scriptClasspath: List<File> = emptyList(),
     private val typeRenderers = mutableMapOf<String, String>()
 
     private val initCellCodes = mutableListOf<String>()
+    private val shutdownCodes = mutableListOf<String>()
 
     private fun renderResult(value: Any?, resultField: Pair<String, KotlinType>?): Any? {
         if (value == null || resultField == null) return null
@@ -119,13 +122,20 @@ class ReplForJupyterImpl(private val scriptClasspath: List<File> = emptyList(),
         return renderResult(result.value, result.resultField)
     }
 
-    data class PreprocessingResult(val code: Code, val initCodes: List<Code>, val initCellCodes: List<Code>, val typeRenderers: List<TypeHandler>)
+    data class PreprocessingResult(
+            val code: Code,
+            val initCodes: List<Code>,
+            val shutdownCodes: List<Code>,
+            val initCellCodes: List<Code>,
+            val typeRenderers: List<TypeHandler>,
+    )
 
     fun preprocessCode(code: String): PreprocessingResult {
 
         val processedMagics = magics.processMagics(code)
 
         val initCodes = mutableListOf<Code>()
+        val shutdownCodes = mutableListOf<Code>()
         val initCellCodes = mutableListOf<Code>()
         val typeRenderers = mutableListOf<TypeHandler>()
         val typeConverters = mutableListOf<TypeHandler>()
@@ -141,13 +151,16 @@ class ReplForJupyterImpl(private val scriptClasspath: List<File> = emptyList(),
             typeRenderers.addAll(libraryDefinition.renderers)
             typeConverters.addAll(libraryDefinition.converters)
             annotations.addAll(libraryDefinition.annotations)
+            initCellCodes.addAll(libraryDefinition.initCell)
+            shutdownCodes.addAll(libraryDefinition.shutdown)
             libraryDefinition.init.forEach {
 
                 // Library init code may contain other magics, so we process them recursively
                 val preprocessed = preprocessCode(it)
                 initCodes.addAll(preprocessed.initCodes)
                 typeRenderers.addAll(preprocessed.typeRenderers)
                 initCellCodes.addAll(preprocessed.initCellCodes)
+                shutdownCodes.addAll(preprocessed.shutdownCodes)
                 if (preprocessed.code.isNotBlank())
                     initCodes.add(preprocessed.code)
             }
@@ -159,7 +172,7 @@ class ReplForJupyterImpl(private val scriptClasspath: List<File> = emptyList(),
             initCodes.add(declarations)
         }
 
-        return PreprocessingResult(processedMagics.code, initCodes, initCellCodes, typeRenderers)
+        return PreprocessingResult(processedMagics.code, initCodes, shutdownCodes, initCellCodes, typeRenderers)
     }
 
     private val ctx = KotlinContext()
@@ -327,6 +340,7 @@ class ReplForJupyterImpl(private val scriptClasspath: List<File> = emptyList(),
 
     private fun registerNewLibraries(p: PreprocessingResult) {
         p.initCellCodes.filter { !initCellCodes.contains(it) }.let(initCellCodes::addAll)
+        p.shutdownCodes.filter { !shutdownCodes.contains(it) }.let(shutdownCodes::addAll)
         typeRenderers.putAll(p.typeRenderers.map { it.className to it.code })
     }
 
@@ -395,6 +409,10 @@ class ReplForJupyterImpl(private val scriptClasspath: List<File> = emptyList(),
         }
     }
 
+    override fun evalOnShutdown(): List<EvalResult> {
+        return shutdownCodes.map(::evalWithReturn)
+    }
+
     private fun updateOutputList(jupyterId: Int, result: Any?) {
         if (jupyterId >= 0) {
             while (ReplOutputs.count() <= jupyterId) ReplOutputs.add(null)
@@ -457,6 +475,17 @@ class ReplForJupyterImpl(private val scriptClasspath: List<File> = emptyList(),
         processAnnotations(lastReplLine())
     }
 
+    // Result of this function is considered to be used for testing/debug purposes
+    private fun evalWithReturn(code: String): EvalResult {
+        val result = try {
+            doEval(code)
+        } catch (e: Exception) {
+            InternalEvalResult(null, null)
+        }
+        processAnnotations(lastReplLine())
+        return EvalResult(result.value)
+    }
+
     private data class InternalEvalResult(val value: Any?, val resultField: Pair<String, KotlinType>?)
 
     private interface LockQueueArgs <T> {
diff --git a/src/test/kotlin/org/jetbrains/kotlin/jupyter/test/replTests.kt b/src/test/kotlin/org/jetbrains/kotlin/jupyter/test/replTests.kt
@@ -1,7 +1,5 @@
 package org.jetbrains.kotlin.jupyter.test
 
-import com.beust.klaxon.JsonObject
-import com.beust.klaxon.Parser
 import jupyter.kotlin.JavaRuntime
 import jupyter.kotlin.MimeTypedResult
 import jupyter.kotlin.receivers.ConstReceiver
@@ -11,10 +9,8 @@ import org.jetbrains.kotlin.jupyter.ReplCompilerException
 import org.jetbrains.kotlin.jupyter.ReplEvalRuntimeException
 import org.jetbrains.kotlin.jupyter.ReplForJupyterImpl
 import org.jetbrains.kotlin.jupyter.ResolverConfig
-import org.jetbrains.kotlin.jupyter.asAsync
 import org.jetbrains.kotlin.jupyter.defaultRepositories
 import org.jetbrains.kotlin.jupyter.generateDiagnostic
-import org.jetbrains.kotlin.jupyter.parserLibraryDescriptors
 import org.jetbrains.kotlin.jupyter.repl.completion.CompletionResult
 import org.jetbrains.kotlin.jupyter.repl.completion.ListErrorsResult
 import org.jetbrains.kotlin.jupyter.withPath
@@ -363,11 +359,10 @@ class ReplTest : AbstractReplTest() {
                                                 ]
                                             }
         """.trimIndent()
-        val parser = Parser.default()
 
-        val libJsons = arrayOf(lib1, lib2, lib3).map { it.first to parser.parse(StringBuilder(it.second)) as JsonObject }.toMap()
+        val libs = listOf(lib1, lib2, lib3).toLibrariesAsync()
 
-        val replWithResolver = ReplForJupyterImpl(classpath, ResolverConfig(defaultRepositories, parserLibraryDescriptors(libJsons).asAsync()))
+        val replWithResolver = ReplForJupyterImpl(classpath, ResolverConfig(defaultRepositories, libs))
         val res = replWithResolver.preprocessCode("%use mylib(1.0), another")
         assertEquals("", res.code)
         val inits = arrayOf(
@@ -397,6 +392,36 @@ class ReplTest : AbstractReplTest() {
         }
     }
 
+    @Test
+    fun testLibraryOnShutdown() {
+        val lib1 = "mylib" to """
+                    {
+                        "shutdown": [
+                            "14 * 3",
+                            "throw RuntimeException()",
+                            "21 + 22"
+                        ]
+                    }""".trimIndent()
+
+        val lib2 = "mylib2" to """
+                    {
+                        "shutdown": [
+                            "100"
+                        ]
+                    }""".trimIndent()
+
+        val libs = listOf(lib1, lib2).toLibrariesAsync()
+        val replWithResolver = ReplForJupyterImpl(classpath, ResolverConfig(defaultRepositories, libs))
+        replWithResolver.eval("%use mylib, mylib2")
+        val results = replWithResolver.evalOnShutdown()
+
+        assertEquals(4, results.size)
+        assertEquals(42, results[0].resultValue)
+        assertNull(results[1].resultValue)
+        assertEquals(43, results[2].resultValue)
+        assertEquals(100, results[3].resultValue)
+    }
+
     @Test
     fun testJavaRuntimeUtils() {
         val result = repl.eval("JavaRuntimeUtils.version")
diff --git a/src/test/kotlin/org/jetbrains/kotlin/jupyter/test/testUtil.kt b/src/test/kotlin/org/jetbrains/kotlin/jupyter/test/testUtil.kt
@@ -1,6 +1,10 @@
 package org.jetbrains.kotlin.jupyter.test
 
+import com.beust.klaxon.JsonObject
+import com.beust.klaxon.Parser
 import jupyter.kotlin.DependsOn
+import kotlinx.coroutines.Deferred
+import org.jetbrains.kotlin.jupyter.LibraryDescriptor
 import org.jetbrains.kotlin.jupyter.ResolverConfig
 import org.jetbrains.kotlin.jupyter.asAsync
 import org.jetbrains.kotlin.jupyter.defaultRepositories
@@ -18,3 +22,9 @@ val classpath = scriptCompilationClasspathFromContext(
 
 val testResolverConfig = ResolverConfig(defaultRepositories,
         parserLibraryDescriptors(readLibraries().toMap()).asAsync())
+
+fun Collection<Pair<String, String>>.toLibrariesAsync(): Deferred<Map<String, LibraryDescriptor>> {
+    val parser = Parser.default()
+    val libJsons = map { it.first to parser.parse(StringBuilder(it.second)) as JsonObject }.toMap()
+    return parserLibraryDescriptors(libJsons).asAsync()
+}

Original file line number	Diff line number	Diff line change
`@@ -99,6 +99,7 @@ fun JupyterConnection.Socket.shellMessagesHandler(msg: Message, repl: ReplForJup`
`99`	`99`	`send(makeReplyMessage(msg, "interrupt_reply", content = msg.content))`
`100`	`100`	`}`
`101`	`101`	`"shutdown_request" -> {`
	`102`	`+ repl?.evalOnShutdown()`
`102`	`103`	`send(makeReplyMessage(msg, "shutdown_reply", content = msg.content))`
`103`	`104`	`exitProcess(0)`
`104`	`105`	`}`