1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
(** {2 User-defined communication} *)
open Format
open Jupyter.Iopub
open Jupyter.Comm
type target = string
type comm = string
type receiver =
{
target_name : string;
recv_open : comm -> Yojson.Safe.t -> unit;
recv_msg : comm -> Yojson.Safe.t -> unit;
recv_close : comm -> Yojson.Safe.t -> unit;
}
module Target =
struct
type t = target
let target_tbl = Hashtbl.create 4
let to_string target_name = target_name
let default _ _ = ()
let create
?(recv_open = default)
?(recv_msg = default)
?(recv_close = default)
target_name =
let receiver = { target_name; recv_open; recv_msg; recv_close } in
Hashtbl.replace target_tbl target_name receiver ;
target_name
let close target_name = Hashtbl.remove target_tbl target_name
end
module Comm =
struct
type t = comm
let comm_tbl = Hashtbl.create 8
let to_string comm_id = comm_id
let comms ?target_name () =
let filter_by_target_name all_comms =
match target_name with
| None -> all_comms
| Some target_name ->
List.filter
(fun (_, target_name') -> target_name = target_name')
all_comms
in
Hashtbl.fold
(fun comm_id receiver acc -> (comm_id, receiver.target_name) :: acc)
comm_tbl []
|> filter_by_target_name
let register target_name comm_id =
let receiver = Hashtbl.find Target.target_tbl target_name in
Hashtbl.replace comm_tbl comm_id receiver ;
receiver
(** {2 Send} *)
let default = `Assoc []
let create ?(data = default) target =
let comm_id = Uuidm.(to_string (v `V4)) in
ignore (register target comm_id) ;
Jupyter_notebook__Unsafe.send_iopub
(IOPUB_COMM_OPEN {
comm_target = Some target;
comm_id;
comm_data = data;
}) ;
comm_id
let send comm_id data =
Jupyter_notebook__Unsafe.send_iopub
(IOPUB_COMM_MSG {
comm_target = None;
comm_id;
comm_data = data; })
let close ?(data = default) comm_id =
Jupyter_notebook__Unsafe.send_iopub
(IOPUB_COMM_CLOSE {
comm_target = None;
comm_id;
comm_data = data; })
(** {2 Receive} *)
let recv_open comm =
match comm.comm_target with
| None -> eprintf "Received a comm message without target_name.@."
| Some target_name ->
try
let receiver = register target_name comm.comm_id in
receiver.recv_open comm.comm_id comm.comm_data
with Not_found -> ()
let recv_msg comm =
try
let receiver = Hashtbl.find comm_tbl comm.comm_id in
receiver.recv_msg comm.comm_id comm.comm_data
with Not_found -> ()
let recv_close comm =
try
let receiver = Hashtbl.find comm_tbl comm.comm_id in
receiver.recv_close comm.comm_id comm.comm_data ;
Hashtbl.remove comm_tbl comm.comm_id
with Not_found -> ()
let recv =
let open Jupyter.Shell in
function
| SHELL_COMM_OPEN comm -> recv_open comm
| SHELL_COMM_MSG comm -> recv_msg comm
| SHELL_COMM_CLOSE comm -> recv_close comm
| _ -> failwith "Not comm message"
end