fixup! Giant refactor to move all state into a Kernel struct

Author: JamesWrigley

src/IJulia.jl

@@ -73,6 +73,17 @@ mutable struct Msg
     end
 end
 
+mutable struct Comm{target}
+    id::String
+    primary::Bool
+    on_msg::Function
+    on_close::Function
+    function (::Type{Comm{target}})(id, primary, on_msg, on_close) where {target}
+        comm = new{target}(id, primary, on_msg, on_close)
+        return comm
+    end
+end
+
 @kwdef mutable struct Kernel
     verbose::Bool = IJULIA_DEBUG
     inited::Bool = false
@@ -84,10 +95,16 @@ end
     Out::Dict{Int, Any} = Dict{Int, Any}()
     ans::Any = nothing
     n::Int = 0
+
     capture_stdout::Bool = true
     capture_stderr::Bool = !IJULIA_DEBUG
     capture_stdin::Bool = true
 
+    # This dict holds a map from CommID to Comm so that we can
+    # pick out the right Comm object when messages arrive
+    # from the front-end.
+    comms = Dict{String, CommManager.Comm}()
+
     postexecute_hooks::Vector{Function} = Function[]
     preexecute_hooks::Vector{Function} = Function[]
     posterror_hooks::Vector{Function} = Function[]

src/comm_manager.jl

@@ -2,41 +2,25 @@ module CommManager
 
 using IJulia
 
-import IJulia: Msg, uuid4, send_ipython, msg_pub
+import IJulia: Msg, uuid4, send_ipython, msg_pub, Comm
 
-export Comm, comm_target, msg_comm, send_comm, close_comm,
+export comm_target, msg_comm, send_comm, close_comm,
        register_comm, comm_msg, comm_open, comm_close, comm_info_request
 
-mutable struct Comm{target}
-    id::String
-    primary::Bool
-    on_msg::Function
-    on_close::Function
-    function (::Type{Comm{target}})(id, primary, on_msg, on_close) where {target}
-        comm = new{target}(id, primary, on_msg, on_close)
-        comms[id] = comm
-        return comm
-    end
-end
-
-# This dict holds a map from CommID to Comm so that we can
-# pick out the right Comm object when messages arrive
-# from the front-end.
-const comms = Dict{String, Comm}()
-
 noop_callback(msg) = nothing
 function Comm(target,
               id=uuid4(),
               primary=true,
               on_msg=noop_callback,
               on_close=noop_callback;
+              kernel=IJulia._default_kernel,
               data=Dict(),
               metadata=Dict())
     comm = Comm{Symbol(target)}(id, primary, on_msg, on_close)
     if primary
         # Request a secondary object be created at the front end
-        send_ipython(IJulia.publish[],
-                     msg_comm(comm, IJulia.execute_msg, "comm_open",
+        send_ipython(kernel.publish[], kernel,
+                     msg_comm(comm, kernel.execute_msg, "comm_open",
                               data, metadata, target_name=string(target)))
     end
     return comm
@@ -47,10 +31,10 @@ comm_target(comm :: Comm{target}) where {target} = target
 function comm_info_request(sock, kernel, msg)
     reply = if haskey(msg.content, "target_name")
         t = Symbol(msg.content["target_name"])
-        filter(kv -> comm_target(kv.second) == t, comms)
+        filter(kv -> comm_target(kv.second) == t, kernel.comms)
     else
         # reply with all comms.
-        comms
+        kernel.comms
     end
 
     _comms = Dict{String, Dict{Symbol,Symbol}}()
@@ -75,16 +59,16 @@ function msg_comm(comm::Comm, m::IJulia.Msg, msg_type,
     return msg_pub(m, msg_type, content, metadata)
 end
 
-function send_comm(comm::Comm, kernel, data::Dict,
-                   metadata::Dict = Dict(); kwargs...)
-    msg = msg_comm(comm, IJulia.execute_msg, "comm_msg", data,
+function send_comm(comm::Comm, data::Dict,
+                   metadata::Dict = Dict(); kernel=IJulia._default_kernel, kwargs...)
+    msg = msg_comm(comm, kernel.execute_msg, "comm_msg", data,
                    metadata; kwargs...)
     send_ipython(kernel.publish[], kernel, msg)
 end
 
-function close_comm(comm::Comm, kernel, data::Dict = Dict(),
-                    metadata::Dict = Dict(); kwargs...)
-    msg = msg_comm(comm, IJulia.execute_msg, "comm_close", data,
+function close_comm(comm::Comm, data::Dict = Dict(),
+                    metadata::Dict = Dict(); kernel=IJulia._default_kernel, kwargs...)
+    msg = msg_comm(comm, kernel.execute_msg, "comm_close", data,
                    metadata; kwargs...)
     send_ipython(kernel.publish[], kernel, msg)
 end
@@ -97,32 +81,32 @@ end
 
 # handlers for incoming comm_* messages
 
-function comm_open(sock, msg, kernel)
+function comm_open(sock, kernel, msg)
     if haskey(msg.content, "comm_id")
         comm_id = msg.content["comm_id"]
         if haskey(msg.content, "target_name")
             target = msg.content["target_name"]
             if !haskey(msg.content, "data")
                 msg.content["data"] = Dict()
             end
-            comm = Comm(target, comm_id, false)
+            comm = Comm(target, comm_id, false; kernel)
             invokelatest(register_comm, comm, msg)
-            comms[comm_id] = comm
+            kernel.comms[comm_id] = comm
         else
             # Tear down comm to maintain consistency
             # if a target_name is not present
             send_ipython(kernel.publish[], kernel,
-                         msg_comm(Comm(:notarget, comm_id),
+                         msg_comm(Comm(:notarget, comm_id; kernel),
                                   msg, "comm_close"))
         end
     end
 end
 
-function comm_msg(sock, msg)
+function comm_msg(sock, kernel, msg)
     if haskey(msg.content, "comm_id")
         comm_id = msg.content["comm_id"]
-        if haskey(comms, comm_id)
-            comm = comms[comm_id]
+        if haskey(kernel.comms, comm_id)
+            comm = kernel.comms[comm_id]
         else
             # We don't have that comm open
             return
@@ -135,17 +119,17 @@ function comm_msg(sock, msg)
     end
 end
 
-function comm_close(sock, msg)
+function comm_close(sock, kernel, msg)
     if haskey(msg.content, "comm_id")
         comm_id = msg.content["comm_id"]
-        comm = comms[comm_id]
+        comm = kernel.comms[comm_id]
 
         if !haskey(msg.content, "data")
             msg.content["data"] = Dict()
         end
         comm.on_close(msg)
 
-        delete!(comms, comm.id)
+        delete!(kernel.comms, comm.id)
     end
 end
 

src/handlers.jl

@@ -331,7 +331,7 @@ function interrupt_request(socket, kernel, msg)
     send_ipython(socket, kernel, msg_reply(msg, "interrupt_reply", Dict()))
 end
 
-function unknown_request(socket, msg)
+function unknown_request(socket, kernel, msg)
     @vprintln("UNKNOWN MESSAGE TYPE $(msg.header["msg_type"])")
 end
 

test/kernel.jl

@@ -95,15 +95,19 @@ inspect(client, code) = make_request(client.inspect, client.get_shell_msg; code)
 get_stdin_msg(client) = make_request(Returns(nothing), client.get_stdin_msg)
 get_iopub_msg(client) = make_request(Returns(nothing), client.get_iopub_msg)
 
-function get_execute_result(client)
+function get_iopub_msgtype(client, msg_type)
     while true
         msg = get_iopub_msg(client)
-        if msg["header"]["msg_type"] == "execute_result"
+        if msg["header"]["msg_type"] == msg_type
             return msg
         end
     end
 end
 
+get_execute_result(client) = get_iopub_msgtype(client, "execute_result")
+get_comm_close(client) = get_iopub_msgtype(client, "comm_close")
+get_comm_msg(client) = get_iopub_msgtype(client, "comm_msg")
+
 function msg_ok(msg)
     ok = msg["content"]["status"] == "ok"
     if !ok
@@ -190,6 +194,52 @@ end
         shutdown_called = false
         Kernel(profile; capture_stdout=false, capture_stderr=false, shutdown=() -> shutdown_called = true) do kernel
             jupyter_client(profile) do client
+                @testset "Comms" begin
+                    # Try opening a Comm without a target_name, which should
+                    # only trigger a comm_close message.
+                    open_msg = IJulia.Msg(["foo"],
+                                          Dict("username" => "user",
+                                               "session" => "session1"),
+                                          Dict("comm_id" => "foo",
+                                               "data" => Dict()))
+                    IJulia.comm_open(kernel.requests[], kernel, open_msg)
+                    @test isempty(kernel.comms)
+                    @test get_comm_close(client)["content"]["comm_id"] == "foo"
+
+                    # Setting the target_name should cause the Comm to be created
+                    open_msg.content["target_name"] = "foo"
+                    IJulia.comm_open(kernel.requests[], kernel, open_msg)
+                    @test kernel.comms["foo"] isa IJulia.Comm{:foo}
+
+                    @test haskey(comm_info(client)["content"]["comms"], "foo")
+
+                    # Smoke test for comm_msg (incoming to the kernel)
+                    msg_msg = IJulia.Msg(["foo"],
+                                         Dict("username" => "user",
+                                              "session" => "session1"),
+                                         Dict("comm_id" => "foo",
+                                              "data" => Dict()))
+                    IJulia.comm_msg(kernel.requests[], kernel, msg_msg)
+
+                    # Test comm_msg (outgoing from the kernel)
+                    IJulia.send_comm(kernel.comms["foo"], Dict(1 => 2))
+                    @test get_comm_msg(client)["content"]["data"]["1"] == 2
+
+                    # Test comm_close (outgoing from the kernel)
+                    IJulia.close_comm(kernel.comms["foo"])
+                    # Should this also delete the Comm from kernel.comms?
+                    @test get_comm_close(client)["content"]["comm_id"] == "foo"
+
+                    # Test comm_close (incoming to the kernel)
+                    close_msg = IJulia.Msg(["foo"],
+                                           Dict("username" => "user",
+                                                "session" => "session1"),
+                                           Dict("comm_id" => "foo",
+                                                "data" => Dict()))
+                    IJulia.comm_close(kernel.requests[], kernel, close_msg)
+                    @test !haskey(kernel.comms, "foo")
+                end
+
                 # Test load()/load_string()
                 mktemp() do path, _
                     write(path, "42")