Refactor magics arguments handling, make it more uniform

Author: ileasile

build-plugin/common-dependencies/src/org/jetbrains/kotlinx/jupyter/common/ReplLineMagic.kt

@@ -2,10 +2,10 @@ package org.jetbrains.kotlinx.jupyter.common
 
 enum class ReplLineMagic(val desc: String, val argumentsUsage: String? = null, val visibleInHelp: Boolean = true) {
     USE("injects code for supported libraries: artifact resolution, default imports, initialization code, type renderers", "klaxon(5.5), lets-plot"),
-    TRACK_CLASSPATH("logs any changes of current classpath. Useful for debugging artifact resolution failures"),
-    TRACK_EXECUTION("logs pieces of code that are going to be executed. Useful for debugging of libraries support"),
-    DUMP_CLASSES_FOR_SPARK("stores compiled repl classes in special folder for Spark integration", visibleInHelp = false),
-    USE_LATEST_DESCRIPTORS("use latest versions of library descriptors available. By default, bundled descriptors are used", "-[on|off]"),
+    TRACK_CLASSPATH("logs any changes of current classpath. Useful for debugging artifact resolution failures", "[on|off]"),
+    TRACK_EXECUTION("logs pieces of code that are going to be executed. Useful for debugging of libraries support", "[all|generated|off]"),
+    DUMP_CLASSES_FOR_SPARK("stores compiled repl classes in special folder for Spark integration", "[on|off]", visibleInHelp = false),
+    USE_LATEST_DESCRIPTORS("use latest versions of library descriptors available. By default, bundled descriptors are used", "[on|off]"),
     OUTPUT("output capturing settings", "--max-cell-size=1000 --no-stdout --max-time=100 --max-buffer=400"),
     LOG_LEVEL("set logging level", "[off|error|warn|info|debug]"),
     LOG_HANDLER("manage logging handlers", "[list | remove <name> | add <name> --<type> [... typeArgs]]", visibleInHelp = false);

docs/README.md

@@ -150,9 +150,9 @@ The following maven repositories are included by default:
 
 The following line magics are supported:
  - `%use` - injects code for supported libraries: artifact resolution, default imports, initialization code, type renderers. Usage example: `%use klaxon(5.5), lets-plot`
- - `%trackClasspath` - logs any changes of current classpath. Useful for debugging artifact resolution failures.
- - `%trackExecution` - logs pieces of code that are going to be executed. Useful for debugging of libraries support.
- - `%useLatestDescriptors` - use latest versions of library descriptors available. By default, bundled descriptors are used. Usage example: `%useLatestDescriptors -[on|off]`
+ - `%trackClasspath` - logs any changes of current classpath. Useful for debugging artifact resolution failures. Usage example: `%trackClasspath [on|off]`
+ - `%trackExecution` - logs pieces of code that are going to be executed. Useful for debugging of libraries support. Usage example: `%trackExecution [all|generated|off]`
+ - `%useLatestDescriptors` - use latest versions of library descriptors available. By default, bundled descriptors are used. Usage example: `%useLatestDescriptors [on|off]`
  - `%output` - output capturing settings. Usage example: `%output --max-cell-size=1000 --no-stdout --max-time=100 --max-buffer=400`
  - `%logLevel` - set logging level. Usage example: `%logLevel [off|error|warn|info|debug]`
 

jupyter-lib/shared-compiler/build.gradle.kts

@@ -36,6 +36,9 @@ dependencies {
 
     // Logging
     compileOnly(libs.logging.slf4j.api)
+
+    // Clikt library for parsing output magics
+    implementation(libs.clikt)
 }
 
 buildSettings {

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

@@ -1,5 +1,9 @@
 package org.jetbrains.kotlinx.jupyter.magics
 
+import com.github.ajalt.clikt.core.CliktCommand
+import com.github.ajalt.clikt.parameters.arguments.argument
+import com.github.ajalt.clikt.parameters.arguments.optional
+import com.github.ajalt.clikt.parameters.types.choice
 import org.jetbrains.kotlinx.jupyter.api.libraries.LibraryDefinitionProducer
 import org.jetbrains.kotlinx.jupyter.common.ReplLineMagic
 
@@ -8,6 +12,16 @@ abstract class AbstractMagicsHandler : MagicsHandler {
     protected var tryIgnoreErrors: Boolean = false
     protected var parseOnly: Boolean = false
 
+    protected fun argumentsList() = arg?.trim()?.takeIf { it.isNotEmpty() }?.split(" ") ?: emptyList()
+    protected fun handleSingleOptionalFlag(action: (Boolean?) -> Unit) {
+        object : CliktCommand() {
+            val arg by nullableFlag()
+            override fun run() {
+                action(arg)
+            }
+        }.parse(argumentsList())
+    }
+
     protected val newLibraries: MutableList<LibraryDefinitionProducer> = mutableListOf()
 
     private val callbackMap: Map<ReplLineMagic, () -> Unit> = mapOf(
@@ -45,4 +59,8 @@ abstract class AbstractMagicsHandler : MagicsHandler {
     open fun handleOutput() {}
     open fun handleLogLevel() {}
     open fun handleLogHandler() {}
+
+    companion object {
+        fun CliktCommand.nullableFlag() = argument().choice(mapOf("on" to true, "off" to false)).optional()
+    }
 }

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

@@ -20,10 +20,9 @@ open class UseMagicsHandler(
     }
 
     override fun handleUseLatestDescriptors() {
-        libraryResolutionInfoSwitcher.switch = when (arg?.trim()) {
-            "-on" -> DefaultInfoSwitch.GIT_REFERENCE
-            "-off" -> DefaultInfoSwitch.DIRECTORY
-            else -> DefaultInfoSwitch.GIT_REFERENCE
+        handleSingleOptionalFlag {
+            libraryResolutionInfoSwitcher.switch = if (it == false) DefaultInfoSwitch.DIRECTORY
+            else DefaultInfoSwitch.GIT_REFERENCE
         }
     }
 }

samples/Titanic.ipynb

@@ -16,7 +16,7 @@
    "outputs": [],
    "source": [
     "// to see autogenerated code, uncomment the line below:\n",
-    "//%trackExecution -generated"
+    "//%trackExecution generated"
    ]
   },
   {
@@ -1436,17 +1436,8 @@
    "name": "kotlin",
    "pygments_lexer": "kotlin",
    "version": "1.4.0-dev-7568"
-  },
-  "pycharm": {
-   "stem_cell": {
-    "cell_type": "raw",
-    "metadata": {
-     "collapsed": false
-    },
-    "source": []
-   }
   }
  },
  "nbformat": 4,
  "nbformat_minor": 2
-}
+}

src/main/kotlin/org/jetbrains/kotlinx/jupyter/magics/FullMagicsHandler.kt

@@ -5,9 +5,13 @@ import ch.qos.logback.classic.spi.ILoggingEvent
 import ch.qos.logback.core.Appender
 import ch.qos.logback.core.FileAppender
 import com.github.ajalt.clikt.core.CliktCommand
+import com.github.ajalt.clikt.parameters.arguments.argument
+import com.github.ajalt.clikt.parameters.arguments.optional
 import com.github.ajalt.clikt.parameters.options.default
 import com.github.ajalt.clikt.parameters.options.flag
 import com.github.ajalt.clikt.parameters.options.option
+import com.github.ajalt.clikt.parameters.types.choice
+import com.github.ajalt.clikt.parameters.types.enum
 import com.github.ajalt.clikt.parameters.types.int
 import com.github.ajalt.clikt.parameters.types.long
 import org.jetbrains.kotlinx.jupyter.ExecutedCodeLogging
@@ -57,36 +61,46 @@ class FullMagicsHandler(
     }
 
     override fun handleTrackExecution() {
-        repl.executedCodeLogging = when (arg?.trim()) {
-            "-all" -> ExecutedCodeLogging.All
-            "-off" -> ExecutedCodeLogging.Off
-            "-generated" -> ExecutedCodeLogging.Generated
-            else -> ExecutedCodeLogging.All
-        }
+        object : CliktCommand() {
+            val logLevel by argument().enum<ExecutedCodeLogging>(true) { it.name.lowercase() }.optional()
+            override fun run() {
+                repl.executedCodeLogging = logLevel ?: ExecutedCodeLogging.ALL
+            }
+        }.parse(argumentsList())
     }
 
     override fun handleTrackClasspath() {
-        repl.trackClasspath = true
+        handleSingleOptionalFlag {
+            repl.trackClasspath = it ?: true
+        }
     }
 
     override fun handleDumpClassesForSpark() {
-        repl.writeCompiledClasses = true
+        handleSingleOptionalFlag {
+            repl.writeCompiledClasses = it ?: true
+        }
     }
 
     override fun handleOutput() {
-        repl.outputConfig = updateOutputConfig(repl.outputConfig, (arg ?: "").split(" "))
+        repl.outputConfig = updateOutputConfig(repl.outputConfig, argumentsList())
     }
 
     override fun handleLogLevel() {
-        val level = when (val levelStr = arg?.trim()) {
-            "off" -> Level.OFF
-            "error" -> Level.ERROR
-            "warn" -> Level.WARN
-            "info" -> Level.INFO
-            "debug" -> Level.DEBUG
-            else -> throw ReplException("Unknown log level: '$levelStr'")
-        }
-        setRootLoggingLevel(level)
+        object : CliktCommand() {
+            val level by argument().choice(
+                mapOf(
+                    "off" to Level.OFF,
+                    "error" to Level.ERROR,
+                    "warn" to Level.WARN,
+                    "info" to Level.INFO,
+                    "debug" to Level.DEBUG,
+                ),
+                ignoreCase = false
+            )
+            override fun run() {
+                setRootLoggingLevel(level)
+            }
+        }.parse(argumentsList())
     }
 
     override fun handleLogHandler() {

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

@@ -95,9 +95,9 @@ class EvalRequestData(
 class ReplEvalRuntimeException(message: String, cause: Throwable? = null) : ReplException(message, cause)
 
 enum class ExecutedCodeLogging {
-    Off,
-    All,
-    Generated
+    OFF,
+    ALL,
+    GENERATED
 }
 
 interface ReplRuntimeProperties {
@@ -212,12 +212,12 @@ class ReplForJupyterImpl(
 
     override var trackClasspath: Boolean = false
 
-    private var _executedCodeLogging: ExecutedCodeLogging = ExecutedCodeLogging.Off
+    private var _executedCodeLogging: ExecutedCodeLogging = ExecutedCodeLogging.OFF
     override var executedCodeLogging: ExecutedCodeLogging
         get() = _executedCodeLogging
         set(value) {
             _executedCodeLogging = value
-            internalEvaluator.logExecution = value != ExecutedCodeLogging.Off
+            internalEvaluator.logExecution = value != ExecutedCodeLogging.OFF
         }
 
     override var writeCompiledClasses: Boolean
@@ -326,7 +326,7 @@ class ReplForJupyterImpl(
         jupyterCompiler,
         evaluator,
         contextUpdater,
-        executedCodeLogging != ExecutedCodeLogging.Off,
+        executedCodeLogging != ExecutedCodeLogging.OFF,
         internalVariablesMarkersProcessor,
     )
 

src/test/kotlin/org/jetbrains/kotlinx/jupyter/test/parseMagicsTests.kt

@@ -22,6 +22,7 @@ import org.junit.jupiter.api.Assertions.assertEquals
 import org.junit.jupiter.api.Assertions.assertNull
 import org.junit.jupiter.api.Test
 import java.io.File
+import kotlin.test.assertFalse
 import kotlin.test.assertTrue
 
 private typealias MagicsAndCodeIntervals = Pair<List<CodeInterval>, List<CodeInterval>>
@@ -127,7 +128,7 @@ class ParseMagicsTests {
             get() = standardResolverBranch
         override val librariesDir = KERNEL_LIBRARIES.localLibrariesDir
         override var trackClasspath = false
-        override var executedCodeLogging = ExecutedCodeLogging.Off
+        override var executedCodeLogging = ExecutedCodeLogging.OFF
         override var writeCompiledClasses = false
         override var outputConfig = OutputConfig()
     }
@@ -201,6 +202,15 @@ class ParseMagicsTests {
         }
 
         assertTrue(options.trackClasspath)
+
+        test(
+            """
+                %trackClasspath off
+            """.trimIndent(),
+            ""
+        )
+
+        assertFalse(options.trackClasspath)
     }
 
     @Test
@@ -212,7 +222,7 @@ class ParseMagicsTests {
                     val x = 9
                     %wrongMagic
                     fun f() = 42
-                    %trackExecution -generated
+                    %trackExecution generated
             """.trimIndent(),
             """
                     
@@ -226,7 +236,7 @@ class ParseMagicsTests {
             assertEquals(1, libs.size)
         }
 
-        assertEquals(ExecutedCodeLogging.Generated, options.executedCodeLogging)
+        assertEquals(ExecutedCodeLogging.GENERATED, options.executedCodeLogging)
     }
 
     @Test

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

@@ -61,7 +61,7 @@ class ReplWithStandardResolverTests : AbstractSingleReplTest() {
         eval("%useLatestDescriptors")
         Assertions.assertTrue(infoProvider.fallback is AbstractLibraryResolutionInfo.ByGitRef)
 
-        eval("%useLatestDescriptors -off")
+        eval("%useLatestDescriptors off")
         Assertions.assertTrue(infoProvider.fallback === initialDefaultResolutionInfo)
     }