Add code preprocessors as a part of library API

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,21 +160,34 @@ 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,
-                renderers = renderers,
-                converters = converters,
-                imports = imports,
-                dependencies = dependencies,
-                repositories = repositories,
-                initCell = beforeCellExecution,
-                afterCellExecution = afterCellExecution,
-                shutdown = shutdownCallbacks,
-                classAnnotations = classAnnotations,
-                fileAnnotations = fileAnnotations,
-                resources = resources,
-            )
+            libraryDefinition {
+                it.init = init
+                it.renderers = renderers
+                it.converters = converters
+                it.imports = imports
+                it.dependencies = dependencies
+                it.repositories = repositories
+                it.initCell = beforeCellExecution
+                it.afterCellExecution = afterCellExecution
+                it.shutdown = shutdownCallbacks
+                it.classAnnotations = classAnnotations
+                it.fileAnnotations = fileAnnotations
+                it.resources = resources
+                it.codePreprocessors = codePreprocessors
+            }
     }
 
     override fun getDefinitions(notebook: Notebook): List<LibraryDefinition> {

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,30 +2,45 @@ 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
 import org.jetbrains.kotlinx.jupyter.api.KotlinKernelVersion
 import org.jetbrains.kotlinx.jupyter.api.RendererTypeHandler
 
 /**
- * Trivial implementation of [LibraryDefinition].
- * You may use it in simple cases instead of overriding [LibraryDefinition]
- * to avoid additional anonymous classes creation
+ * Trivial implementation of [LibraryDefinition] - simple container.
  */
-class LibraryDefinitionImpl(
-    override val dependencies: List<String> = emptyList(),
-    override val repositories: List<String> = emptyList(),
-    override val imports: List<String> = emptyList(),
-    override val init: List<ExecutionCallback<*>> = emptyList(),
-    override val initCell: List<ExecutionCallback<*>> = emptyList(),
-    override val afterCellExecution: List<AfterCellExecutionCallback> = emptyList(),
-    override val shutdown: List<ExecutionCallback<*>> = emptyList(),
-    override val renderers: List<RendererTypeHandler> = emptyList(),
-    override val converters: List<FieldHandler> = emptyList(),
-    override val classAnnotations: List<ClassAnnotationHandler> = emptyList(),
-    override val fileAnnotations: List<FileAnnotationHandler> = emptyList(),
-    override val resources: List<LibraryResource> = emptyList(),
-    override val minKernelVersion: KotlinKernelVersion? = null,
-    override val originalDescriptorText: String? = null,
-) : LibraryDefinition
+class LibraryDefinitionImpl private constructor() : LibraryDefinition {
+    override var dependencies: List<String> = emptyList()
+    override var repositories: List<String> = emptyList()
+    override var imports: List<String> = emptyList()
+    override var init: List<ExecutionCallback<*>> = emptyList()
+    override var initCell: List<ExecutionCallback<*>> = emptyList()
+    override var afterCellExecution: List<AfterCellExecutionCallback> = emptyList()
+    override var shutdown: List<ExecutionCallback<*>> = emptyList()
+    override var renderers: List<RendererTypeHandler> = emptyList()
+    override var converters: List<FieldHandler> = emptyList()
+    override var classAnnotations: List<ClassAnnotationHandler> = emptyList()
+    override var fileAnnotations: List<FileAnnotationHandler> = emptyList()
+    override var resources: List<LibraryResource> = emptyList()
+    override var codePreprocessors: List<CodePreprocessor> = emptyList()
+    override var minKernelVersion: KotlinKernelVersion? = null
+    override var originalDescriptorText: String? = null
+
+    companion object {
+        internal fun build(buildAction: (LibraryDefinitionImpl) -> Unit): LibraryDefinition {
+            return LibraryDefinitionImpl().also(buildAction)
+        }
+    }
+}
+
+/**
+ * Builds an instance of [LibraryDefinition].
+ * Build action receives [LibraryDefinitionImpl] as an explicit argument
+ * because of problems with names clashing that may arise.
+ */
+fun libraryDefinition(buildAction: (LibraryDefinitionImpl) -> Unit): LibraryDefinition {
+    return LibraryDefinitionImpl.build(buildAction)
+}

jupyter-lib/shared-compiler/src/main/kotlin/org/jetbrains/kotlinx/jupyter/libraries/LibraryDescriptor.kt

@@ -8,8 +8,8 @@ import org.jetbrains.kotlinx.jupyter.api.ExactRendererTypeHandler
 import org.jetbrains.kotlinx.jupyter.api.KotlinKernelVersion
 import org.jetbrains.kotlinx.jupyter.api.libraries.CodeExecution
 import org.jetbrains.kotlinx.jupyter.api.libraries.LibraryDefinition
-import org.jetbrains.kotlinx.jupyter.api.libraries.LibraryDefinitionImpl
 import org.jetbrains.kotlinx.jupyter.api.libraries.LibraryResource
+import org.jetbrains.kotlinx.jupyter.api.libraries.libraryDefinition
 import org.jetbrains.kotlinx.jupyter.config.currentKernelVersion
 import org.jetbrains.kotlinx.jupyter.exceptions.ReplPreprocessingException
 import org.jetbrains.kotlinx.jupyter.util.KotlinKernelVersionSerializer
@@ -78,18 +78,18 @@ class LibraryDescriptor(
     }
 
     private fun processDescriptor(mapping: Map<String, String>): LibraryDefinition {
-        return LibraryDefinitionImpl(
-            dependencies = dependencies.replaceVariables(mapping),
-            repositories = repositories.replaceVariables(mapping),
-            imports = imports.replaceVariables(mapping),
-            init = init.replaceVariables(mapping),
-            shutdown = shutdown.replaceVariables(mapping),
-            initCell = initCell.replaceVariables(mapping),
-            renderers = renderers.replaceVariables(mapping),
-            resources = resources.replaceVariables(mapping),
-            minKernelVersion = minKernelVersion,
-            originalDescriptorText = Json.encodeToString(this),
-        )
+        return libraryDefinition {
+            it.dependencies = dependencies.replaceVariables(mapping)
+            it.repositories = repositories.replaceVariables(mapping)
+            it.imports = imports.replaceVariables(mapping)
+            it.init = init.replaceVariables(mapping)
+            it.shutdown = shutdown.replaceVariables(mapping)
+            it.initCell = initCell.replaceVariables(mapping)
+            it.renderers = renderers.replaceVariables(mapping)
+            it.resources = resources.replaceVariables(mapping)
+            it.minKernelVersion = minKernelVersion
+            it.originalDescriptorText = Json.encodeToString(this)
+        }
     }
 
     companion object {

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,