Improved output capturing and added tests

README.md

@@ -92,6 +92,9 @@ The following line magics are supported:
  - `%use <lib1>, <lib2> ...` - injects code for supported libraries: artifact resolution, default imports, initialization code, type renderers
  - `%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
+ - `%output [options]` - output capturing settings.
+
+ See detailed info about line magics [here](doc/magics.md).
 
 ### Supported Libraries
 

build.gradle

@@ -37,10 +37,10 @@ allprojects {
     }
 
     dependencies {
-        compile "org.jetbrains.kotlin:kotlin-stdlib:$kotlinVersion"
+        implementation "org.jetbrains.kotlin:kotlin-stdlib:$kotlinVersion"
 
-        testCompile 'junit:junit:4.12'
-        testCompile "org.jetbrains.kotlin:kotlin-test:$kotlinVersion"
+        testImplementation 'junit:junit:4.12'
+        testImplementation "org.jetbrains.kotlin:kotlin-test:$kotlinVersion"
     }
 
     ext {
@@ -166,6 +166,7 @@ dependencies {
     compile 'khttp:khttp:1.0.0'
     compile 'org.zeromq:jeromq:0.3.5'
     compile 'com.beust:klaxon:5.2'
+    compile 'com.github.ajalt:clikt:2.3.0'
     runtime 'org.slf4j:slf4j-simple:1.7.25'
     runtime "org.jetbrains.kotlin:jcabi-aether:1.0-dev-3"
     runtime "org.sonatype.aether:aether-api:1.13.1"

doc/magics.md

@@ -0,0 +1,14 @@
+# Line magics
+The following line magics are supported:
+ - `%use <lib1>, <lib2> ...` - injects code for supported libraries: artifact resolution, default imports, initialization code, type renderers
+ - `%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
+ - `%output [--max-cell-size=N] [--max-buffer=N] [--max-buffer-newline=N] [--max-time=N] [--no-stdout] [--reset-to-defaults]` - 
+ output capturing settings.
+     - `max-cell-size` specifies the characters count which may be printed to stdout. Default is 100000.
+     - `max-buffer` - max characters count stored in internal buffer before being sent to client. Default is 10000.
+     - `max-buffer-newline` - same as above, but trigger happens only if newline character was encountered. Default is 100.
+     - `max-time` - max time in milliseconds before the buffer is sent to client. Default is 100.
+     - `no-stdout` - don't capture output. Default is false.
+     - `reset-to-defaults` - reset all output settings that were set with magics to defaults
+

src/main/kotlin/org/jetbrains/kotlin/jupyter/config.kt

@@ -40,6 +40,22 @@ enum class JupyterSockets {
     iopub
 }
 
+data class OutputConfig(
+        var captureOutput: Boolean = true,
+        var captureBufferTimeLimitMs: Long = 100,
+        var captureBufferMaxSize: Int = 1000,
+        var cellOutputMaxSize: Int = 100000,
+        var captureNewlineBufferSize: Int = 100
+) {
+    fun assign(other: OutputConfig) {
+        captureOutput = other.captureOutput
+        captureBufferTimeLimitMs = other.captureBufferTimeLimitMs
+        captureBufferMaxSize = other.captureBufferMaxSize
+        cellOutputMaxSize = other.cellOutputMaxSize
+        captureNewlineBufferSize = other.captureNewlineBufferSize
+    }
+}
+
 data class RuntimeKernelProperties(val map: Map<String, String>) {
     val version: String
         get() = map["version"] ?: "unspecified"

src/main/kotlin/org/jetbrains/kotlin/jupyter/connection.kt

@@ -5,9 +5,7 @@ import com.beust.klaxon.Parser
 import org.jetbrains.kotlin.com.intellij.openapi.Disposable
 import org.jetbrains.kotlin.com.intellij.openapi.util.Disposer
 import org.zeromq.ZMQ
-import java.io.ByteArrayOutputStream
 import java.io.Closeable
-import java.io.PrintStream
 import java.security.SignatureException
 import java.util.*
 import javax.crypto.Mac
@@ -20,6 +18,14 @@ class JupyterConnection(val config: KernelConfig): Closeable {
         init {
             val port = config.ports[socket.ordinal]
             bind("${config.transport}://*:$port")
+            if (type == ZMQ.PUB) {
+                // Workaround to prevent losing few first messages on kernel startup
+                // For more information on losing messages see this scheme:
+                // http://zguide.zeromq.org/page:all#Missing-Message-Problem-Solver
+                // It seems we cannot do correct sync because messaging protocol
+                // doesn't support this. Value of 500 ms was chosen experimentally.
+                Thread.sleep(500)
+            }
             log.debug("[$name] listen: ${config.transport}://*:$port")
         }
 
@@ -150,13 +156,9 @@ class HMAC(algo: String, key: String?) {
     operator fun invoke(vararg data: ByteArray): String? = invoke(data.asIterable())
 }
 
-fun JupyterConnection.Socket.logWireMessage(msg: ByteArray) {
-    log.debug("[$name] >in: ${String(msg)}")
-}
-
 fun ByteArray.toHexString(): String = joinToString("", transform = { "%02x".format(it) })
 
-fun ZMQ.Socket.sendMessage(msg: Message, hmac: HMAC): Unit {
+fun ZMQ.Socket.sendMessage(msg: Message, hmac: HMAC) {
     msg.id.forEach { sendMore(it) }
     sendMore(DELIM)
     val signableMsg = listOf(msg.header, msg.parentHeader, msg.metadata, msg.content)

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

@@ -1,12 +1,19 @@
 package org.jetbrains.kotlin.jupyter
 
+import com.github.ajalt.clikt.core.CliktCommand
+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.int
+import com.github.ajalt.clikt.parameters.types.long
 import org.jetbrains.kotlin.jupyter.repl.spark.ClassWriter
 
 enum class ReplLineMagics(val desc: String, val argumentsUsage: String? = null, val visibleInHelp: Boolean = true) {
     use("include supported libraries", "klaxon(5.0.1), lets-plot"),
     trackClasspath("log current classpath changes"),
     trackExecution("log code that is going to be executed in repl", visibleInHelp = false),
-    dumpClassesForSpark("stores compiled repl classes in special folder for Spark integration", visibleInHelp = false)
+    dumpClassesForSpark("stores compiled repl classes in special folder for Spark integration", visibleInHelp = false),
+    output("setup output settings", "--max-cell-size=1000 --no-stdout --max-time=100 --max-buffer=400")
 }
 
 fun processMagics(repl: ReplForJupyter, code: String): String {
@@ -15,6 +22,35 @@ fun processMagics(repl: ReplForJupyter, code: String): String {
     var nextSearchIndex = 0
     var nextCopyIndex = 0
 
+    val outputParser = repl.outputConfig.let { conf ->
+        object : CliktCommand() {
+            val defaultConfig = OutputConfig()
+
+            val max: Int by option("--max-cell-size", help = "Maximum cell output").int().default(conf.cellOutputMaxSize)
+            val maxBuffer: Int by option("--max-buffer", help = "Maximum buffer size").int().default(conf.captureBufferMaxSize)
+            val maxBufferNewline: Int by option("--max-buffer-newline", help = "Maximum buffer size when newline got").int().default(conf.captureNewlineBufferSize)
+            val maxTimeInterval: Long by option("--max-time", help = "Maximum time wait for output to accumulate").long().default(conf.captureBufferTimeLimitMs)
+            val dontCaptureStdout: Boolean by option("--no-stdout", help = "Don't capture output").flag(default = !conf.captureOutput)
+            val reset: Boolean by option("--reset-to-defaults", help = "Reset to defaults").flag()
+
+            override fun run() {
+                if (reset) {
+                    conf.assign(defaultConfig)
+                    return
+                }
+                conf.assign(
+                        OutputConfig(
+                                !dontCaptureStdout,
+                                maxTimeInterval,
+                                maxBuffer,
+                                max,
+                                maxBufferNewline
+                        )
+                )
+            }
+        }
+    }
+
     while (true) {
 
         var magicStart: Int
@@ -55,6 +91,9 @@ fun processMagics(repl: ReplForJupyter, code: String): String {
                     if (arg == null) throw ReplCompilerException("Need some arguments for 'use' command")
                     repl.librariesCodeGenerator.processNewLibraries(repl, arg)
                 }
+                ReplLineMagics.output -> {
+                    outputParser.parse((arg ?: "").split(" "))
+                }
             }
             nextCopyIndex = magicEnd
             nextSearchIndex = magicEnd