Add code preprocessors as a part of library API

Fixes #223

Author: ileasile

gradle.properties

@@ -16,7 +16,7 @@ logbackVersion=1.2.3
 http4kVersion=4.4.0.1
 artifactsPath=build/artifacts
 
-baseVersion=0.9.1
+baseVersion=0.10.0
 
 projectRepoUrl=https://github.com/Kotlin/kotlin-jupyter
 docsRepo=git@github.com:ileasile/kotlin-jupyter-docs.git

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

@@ -0,0 +1,20 @@
+package org.jetbrains.kotlinx.jupyter.api
+
+import org.jetbrains.kotlinx.jupyter.api.libraries.LibraryDefinitionProducer
+
+/**
+ * Preprocesses the cell code before its execution
+ */
+interface CodePreprocessor {
+    /**
+     * Returns `true` if this preprocessor accepts the given [code]
+     */
+    fun accepts(code: String): Boolean = true
+
+    /**
+     * Performs code preprocessing
+     */
+    fun process(code: String, host: KotlinKernelHost): Result
+
+    data class Result(val code: Code, val libraries: List<LibraryDefinitionProducer> = emptyList())
+}

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

@@ -3,7 +3,9 @@ package org.jetbrains.kotlinx.jupyter.api.libraries
 import org.jetbrains.kotlinx.jupyter.api.AfterCellExecutionCallback
 import org.jetbrains.kotlinx.jupyter.api.ClassAnnotationHandler
 import org.jetbrains.kotlinx.jupyter.api.ClassDeclarationsCallback
+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.ExecutionCallback
 import org.jetbrains.kotlinx.jupyter.api.FieldHandler
 import org.jetbrains.kotlinx.jupyter.api.FieldHandlerByClass
@@ -54,6 +56,8 @@ abstract class JupyterIntegration : LibraryDefinitionProducer {
 
         private val repositories = mutableListOf<String>()
 
+        private val codePreprocessors = mutableListOf<CodePreprocessor>()
+
         fun addRenderer(handler: RendererTypeHandler) {
             renderers.add(handler)
         }
@@ -70,6 +74,10 @@ abstract class JupyterIntegration : LibraryDefinitionProducer {
             fileAnnotations.add(handler)
         }
 
+        fun addCodePreprocessor(preprocessor: CodePreprocessor) {
+            codePreprocessors.add(preprocessor)
+        }
+
         inline fun <reified T : Any> render(noinline renderer: CodeCell.(T) -> Any) {
             return renderWithHost { _, value: T -> renderer(this, value) }
         }
@@ -152,6 +160,18 @@ abstract class JupyterIntegration : LibraryDefinitionProducer {
             addFileAnnotationHanlder(FileAnnotationHandler(T::class, callback))
         }
 
+        fun preprocessCodeWithLibraries(callback: KotlinKernelHost.(Code) -> CodePreprocessor.Result) {
+            addCodePreprocessor(object : CodePreprocessor {
+                override fun process(code: String, host: KotlinKernelHost): CodePreprocessor.Result {
+                    return host.callback(code)
+                }
+            })
+        }
+
+        fun preprocessCode(callback: KotlinKernelHost.(Code) -> Code) {
+            preprocessCodeWithLibraries { CodePreprocessor.Result(this.callback(it)) }
+        }
+
         internal fun getDefinition() =
             LibraryDefinitionImpl(
                 init = init,
@@ -166,6 +186,7 @@ abstract class JupyterIntegration : LibraryDefinitionProducer {
                 classAnnotations = classAnnotations,
                 fileAnnotations = fileAnnotations,
                 resources = resources,
+                codePreprocessors = codePreprocessors,
             )
     }
 

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

@@ -2,6 +2,7 @@ package org.jetbrains.kotlinx.jupyter.api.libraries
 
 import org.jetbrains.kotlinx.jupyter.api.AfterCellExecutionCallback
 import org.jetbrains.kotlinx.jupyter.api.ClassAnnotationHandler
+import org.jetbrains.kotlinx.jupyter.api.CodePreprocessor
 import org.jetbrains.kotlinx.jupyter.api.ExecutionCallback
 import org.jetbrains.kotlinx.jupyter.api.FieldHandler
 import org.jetbrains.kotlinx.jupyter.api.FileAnnotationHandler
@@ -85,6 +86,12 @@ interface LibraryDefinition {
     val resources: List<LibraryResource>
         get() = emptyList()
 
+    /**
+     * List of code preprocessors
+     */
+    val codePreprocessors: List<CodePreprocessor>
+        get() = emptyList()
+
     /**
      * Minimal kernel version that is supported by this library
      */

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

@@ -2,6 +2,7 @@ package org.jetbrains.kotlinx.jupyter.api.libraries
 
 import org.jetbrains.kotlinx.jupyter.api.AfterCellExecutionCallback
 import org.jetbrains.kotlinx.jupyter.api.ClassAnnotationHandler
+import org.jetbrains.kotlinx.jupyter.api.CodePreprocessor
 import org.jetbrains.kotlinx.jupyter.api.ExecutionCallback
 import org.jetbrains.kotlinx.jupyter.api.FieldHandler
 import org.jetbrains.kotlinx.jupyter.api.FileAnnotationHandler
@@ -26,6 +27,7 @@ class LibraryDefinitionImpl(
     override val classAnnotations: List<ClassAnnotationHandler> = emptyList(),
     override val fileAnnotations: List<FileAnnotationHandler> = emptyList(),
     override val resources: List<LibraryResource> = emptyList(),
+    override val codePreprocessors: List<CodePreprocessor> = emptyList(),
     override val minKernelVersion: KotlinKernelVersion? = null,
     override val originalDescriptorText: String? = null,
 ) : LibraryDefinition

jupyter-lib/shared-compiler/src/main/kotlin/org/jetbrains/kotlinx/jupyter/magics/CompoundCodePreprocessor.kt

@@ -0,0 +1,31 @@
+package org.jetbrains.kotlinx.jupyter.magics
+
+import org.jetbrains.kotlinx.jupyter.api.CodePreprocessor
+import org.jetbrains.kotlinx.jupyter.api.KotlinKernelHost
+
+/**
+ * Containing [preprocessors]' [process] are run in reversed order: last added processors
+ * are run first
+ */
+class CompoundCodePreprocessor(
+    private val preprocessors: MutableList<CodePreprocessor>
+) : CodePreprocessor {
+    constructor(vararg preprocessors: CodePreprocessor) : this(preprocessors.toMutableList())
+
+    override fun process(code: String, host: KotlinKernelHost): CodePreprocessor.Result {
+        return preprocessors.foldRight(CodePreprocessor.Result(code, emptyList())) { preprocessor, result ->
+            if (preprocessor.accepts(result.code)) {
+                val newResult = preprocessor.process(result.code, host)
+                CodePreprocessor.Result(newResult.code, result.libraries + newResult.libraries)
+            } else result
+        }
+    }
+
+    fun add(preprocessor: CodePreprocessor) {
+        preprocessors.add(preprocessor)
+    }
+
+    fun addAll(preprocessors: Iterable<CodePreprocessor>) {
+        this.preprocessors.addAll(preprocessors)
+    }
+}

jupyter-lib/shared-compiler/src/main/kotlin/org/jetbrains/kotlinx/jupyter/magics/MagicsProcessor.kt

@@ -1,7 +1,7 @@
 package org.jetbrains.kotlinx.jupyter.magics
 
-import org.jetbrains.kotlinx.jupyter.api.Code
-import org.jetbrains.kotlinx.jupyter.api.libraries.LibraryDefinitionProducer
+import org.jetbrains.kotlinx.jupyter.api.CodePreprocessor
+import org.jetbrains.kotlinx.jupyter.api.KotlinKernelHost
 import org.jetbrains.kotlinx.jupyter.common.ReplLineMagic
 import org.jetbrains.kotlinx.jupyter.compiler.util.CodeInterval
 import org.jetbrains.kotlinx.jupyter.exceptions.ReplPreprocessingException
@@ -10,8 +10,8 @@ import kotlin.script.experimental.jvm.util.determineSep
 class MagicsProcessor(
     private val handler: MagicsHandler,
     private val parseOutCellMarker: Boolean = false,
-) {
-    fun processMagics(code: String, parseOnly: Boolean = false, tryIgnoreErrors: Boolean = false): MagicProcessingResult {
+) : CodePreprocessor {
+    fun processMagics(code: String, parseOnly: Boolean = false, tryIgnoreErrors: Boolean = false): CodePreprocessor.Result {
         val magics = magicsIntervals(code)
 
         for (magicRange in magics) {
@@ -38,7 +38,7 @@ class MagicsProcessor(
 
         val codes = codeIntervals(code, magics, true)
         val preprocessedCode = codes.joinToString("") { code.substring(it.from, it.to) }
-        return MagicProcessingResult(preprocessedCode, handler.getLibraries())
+        return CodePreprocessor.Result(preprocessedCode, handler.getLibraries())
     }
 
     fun codeIntervals(
@@ -79,7 +79,9 @@ class MagicsProcessor(
         }
     }
 
-    data class MagicProcessingResult(val code: Code, val libraries: List<LibraryDefinitionProducer>)
+    override fun process(code: String, host: KotlinKernelHost): CodePreprocessor.Result {
+        return processMagics(code)
+    }
 
     companion object {
         private const val MAGICS_SIGN = '%'

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

@@ -7,6 +7,7 @@ import jupyter.kotlin.KotlinKernelHostProvider
 import jupyter.kotlin.Repository
 import org.jetbrains.kotlin.config.KotlinCompilerVersion
 import org.jetbrains.kotlinx.jupyter.api.Code
+import org.jetbrains.kotlinx.jupyter.api.CodePreprocessor
 import org.jetbrains.kotlinx.jupyter.api.ExecutionCallback
 import org.jetbrains.kotlinx.jupyter.api.KotlinKernelHost
 import org.jetbrains.kotlinx.jupyter.api.KotlinKernelVersion
@@ -40,6 +41,7 @@ import org.jetbrains.kotlinx.jupyter.libraries.LibrariesScanner
 import org.jetbrains.kotlinx.jupyter.libraries.LibraryResourcesProcessorImpl
 import org.jetbrains.kotlinx.jupyter.libraries.ResolutionInfoProvider
 import org.jetbrains.kotlinx.jupyter.libraries.ResolutionInfoSwitcher
+import org.jetbrains.kotlinx.jupyter.magics.CompoundCodePreprocessor
 import org.jetbrains.kotlinx.jupyter.magics.FullMagicsHandler
 import org.jetbrains.kotlinx.jupyter.magics.MagicsProcessor
 import org.jetbrains.kotlinx.jupyter.repl.CellExecutor
@@ -231,6 +233,8 @@ class ReplForJupyterImpl(
         )
     )
 
+    private val codePreprocessor = CompoundCodePreprocessor(magics)
+
     private val importsCollector: ScriptImportsCollector = ScriptImportsCollectorImpl()
 
     // Used for various purposes, i.e. completion and listing errors
@@ -328,7 +332,7 @@ class ReplForJupyterImpl(
         fileAnnotationsProcessor,
         fieldsProcessor,
         typeRenderersProcessor,
-        magics,
+        codePreprocessor,
         resourcesProcessor,
         librariesScanner,
         notebook,

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

@@ -45,7 +45,7 @@ internal class CellExecutorImpl(private val replContext: SharedReplContext) : Ce
 
             log.debug("Executing code:\n$code")
             val preprocessedCode = if (processMagics) {
-                val processedMagics = magicsProcessor.processMagics(code)
+                val processedMagics = codePreprocessor.process(code, context)
 
                 log.debug("Adding ${processedMagics.libraries.size} libraries")
                 processedMagics.libraries.getDefinitions(notebook).forEach {
@@ -128,6 +128,7 @@ internal class CellExecutorImpl(private val replContext: SharedReplContext) : Ce
             library.classAnnotations.forEach(sharedContext.classAnnotationsProcessor::register)
             library.fileAnnotations.forEach(sharedContext.fileAnnotationsProcessor::register)
             sharedContext.afterCellExecution.addAll(library.afterCellExecution)
+            sharedContext.codePreprocessor.addAll(library.codePreprocessors)
 
             val classLoader = sharedContext.evaluator.lastClassLoader
             rethrowAsLibraryException(LibraryProblemPart.RESOURCES) {

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

@@ -9,15 +9,15 @@ import org.jetbrains.kotlinx.jupyter.codegen.FileAnnotationsProcessor
 import org.jetbrains.kotlinx.jupyter.codegen.ResultsTypeRenderersProcessor
 import org.jetbrains.kotlinx.jupyter.libraries.LibrariesScanner
 import org.jetbrains.kotlinx.jupyter.libraries.LibraryResourcesProcessor
-import org.jetbrains.kotlinx.jupyter.magics.MagicsProcessor
+import org.jetbrains.kotlinx.jupyter.magics.CompoundCodePreprocessor
 import org.jetbrains.kotlinx.jupyter.repl.InternalEvaluator
 
 internal data class SharedReplContext(
     val classAnnotationsProcessor: ClassAnnotationsProcessor,
     val fileAnnotationsProcessor: FileAnnotationsProcessor,
     val fieldsProcessor: FieldsProcessor,
     val typeRenderersProcessor: ResultsTypeRenderersProcessor,
-    val magicsProcessor: MagicsProcessor,
+    val codePreprocessor: CompoundCodePreprocessor,
     val resourcesProcessor: LibraryResourcesProcessor,
     val librariesScanner: LibrariesScanner,
     val notebook: Notebook,

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

@@ -116,7 +116,7 @@ class IntegrationApiTests {
         val repl = makeRepl()
         repl.eval(
             """
-            @file:DependsOn("src/test/testData/notebook-api-test-0.0.15.jar")
+            @file:DependsOn("src/test/testData/kotlin-jupyter-api-test-0.0.16.jar")
             """.trimIndent()
         )
 
@@ -166,4 +166,19 @@ class IntegrationApiTests {
         val result = repl.eval("B(A())")
         assertEquals("iB: iA", result.resultValue)
     }
+
+    @Test
+    fun `code preprocessing`() {
+        val repl = makeRepl()
+        repl.eval(
+            """
+            USE {
+                preprocessCode { it.replace('b', 'x') }
+            }
+            """.trimIndent()
+        )
+
+        val result = repl.eval("\"abab\"")
+        assertEquals("axax", result.resultValue)
+    }
 }