Merge pull request #160 from davidbrochart/stdin

Support stdin channel
jupyter-server · Feb 14, 2022 · 9f284b0 · 9f284b0
2 parents 780b85c + f37d199
commit 9f284b0
Show file tree

Hide file tree

Showing 2 changed files with 25 additions and 5 deletions.
diff --git a/plugins/kernels/fps_kernels/kernel_server/connect.py b/plugins/kernels/fps_kernels/kernel_server/connect.py
@@ -87,19 +87,23 @@ async def launch_kernel(
     return p
 
 
-def create_socket(channel: str, cfg: cfg_t) -> Socket:
+def create_socket(channel: str, cfg: cfg_t, identity: Optional[bytes] = None) -> Socket:
     ip = cfg["ip"]
     port = cfg[f"{channel}_port"]
     url = f"tcp://{ip}:{port}"
     socket_type = channel_socket_types[channel]
     sock = context.socket(socket_type)
     sock.linger = 1000  # set linger to 1s to prevent hangs at exit
+    if identity:
+        sock.identity = identity
     sock.connect(url)
     return sock
 
 
-def connect_channel(channel_name: str, cfg: cfg_t) -> Socket:
-    sock = create_socket(channel_name, cfg)
+def connect_channel(
+    channel_name: str, cfg: cfg_t, identity: Optional[bytes] = None
+) -> Socket:
+    sock = create_socket(channel_name, cfg, identity)
     if channel_name == "iopub":
         sock.setsockopt(zmq.SUBSCRIBE, b"")
     return sock

diff --git a/plugins/kernels/fps_kernels/kernel_server/server.py b/plugins/kernels/fps_kernels/kernel_server/server.py
@@ -1,4 +1,5 @@
 import os
+import uuid
 import json
 import asyncio
 import signal
@@ -105,12 +106,21 @@ async def start(self) -> None:
             self.kernelspec_path, self.connection_file_path, self.capture_kernel_output
         )
         assert self.connection_cfg is not None
-        self.shell_channel = connect_channel("shell", self.connection_cfg)
-        self.control_channel = connect_channel("control", self.connection_cfg)
+        identity = uuid.uuid4().hex.encode("ascii")
+        self.shell_channel = connect_channel(
+            "shell", self.connection_cfg, identity=identity
+        )
+        self.stdin_channel = connect_channel(
+            "stdin", self.connection_cfg, identity=identity
+        )
+        self.control_channel = connect_channel(
+            "control", self.connection_cfg, identity=identity
+        )
         self.iopub_channel = connect_channel("iopub", self.connection_cfg)
         await self._wait_for_ready()
         self.channel_tasks += [
             asyncio.create_task(self.listen("shell")),
+            asyncio.create_task(self.listen("stdin")),
             asyncio.create_task(self.listen("control")),
             asyncio.create_task(self.listen("iopub")),
         ]
@@ -154,6 +164,8 @@ async def listen(self, channel_name: str):
             channel = self.control_channel
         elif channel_name == "iopub":
             channel = self.iopub_channel
+        elif channel_name == "stdin":
+            channel = self.stdin_channel
 
         while True:
             parts = await get_zmq_parts(channel)
@@ -196,6 +208,8 @@ async def send_to_zmq(self, websocket):
                     send_message(msg, self.shell_channel, self.key)
                 elif channel == "control":
                     send_message(msg, self.control_channel, self.key)
+                elif channel == "stdin":
+                    send_message(msg, self.stdin_channel, self.key)
         elif websocket.accepted_subprotocol == "v1.kernel.websocket.jupyter.org":
             while True:
                 msg = await websocket.websocket.receive_bytes()
@@ -213,6 +227,8 @@ async def send_to_zmq(self, websocket):
                     send_raw_message(parts, self.shell_channel, self.key)
                 elif channel == "control":
                     send_raw_message(parts, self.control_channel, self.key)
+                elif channel == "stdin":
+                    send_raw_message(parts, self.stdin_channel, self.key)
 
     async def send_to_ws(self, websocket, parts, parent_header, channel_name):
         if not websocket.accepted_subprotocol: