Rework outputs handling

Author: davidbrochart

plugins/kernels/fps_kernels/kernel_driver/driver.py

@@ -1,6 +1,5 @@
 import asyncio
 import os
-import sys
 import time
 import uuid
 from typing import Any, Dict, List, Optional, Tuple, cast
@@ -39,19 +38,6 @@ async def receive_message(sock: Socket, timeout: float = float("inf")) -> Option
     return None
 
 
-def _output_hook_default(msg: Dict[str, Any]) -> None:
-    """Default hook for redisplaying plain-text output"""
-    msg_type = msg["header"]["msg_type"]
-    content = msg["content"]
-    if msg_type == "stream":
-        stream = getattr(sys, content["name"])
-        stream.write(content["text"])
-    elif msg_type in ("display_data", "execute_result"):
-        sys.stdout.write(content["data"].get("text/plain", ""))
-    elif msg_type == "error":
-        print("\n".join(content["traceback"]), file=sys.stderr)
-
-
 class KernelDriver:
     def __init__(
         self,
@@ -136,13 +122,14 @@ async def listen_shell(self):
 
     async def execute(
         self,
-        code: str,
+        cell: Dict[str, Any],
         timeout: float = float("inf"),
         msg_id: str = "",
         wait_for_executed: bool = True,
-        output_hook=_output_hook_default,
     ) -> None:
-        content = {"code": code, "silent": False}
+        if cell["cell_type"] != "code":
+            return
+        content = {"code": cell["source"], "silent": False}
         msg = create_message(
             "execute_request", content, session_id=self.session_id, msg_cnt=self.msg_cnt
         )
@@ -168,7 +155,7 @@ async def execute(
                     error_message = f"Kernel didn't respond in {timeout} seconds"
                     raise RuntimeError(error_message)
                 msg = self.execute_requests[msg_id]["iopub_msg"]
-                output_hook(msg)
+                self._handle_outputs(cell["outputs"], msg)
                 if (
                     msg["header"]["msg_type"] == "status"
                     and msg["content"]["execution_state"] == "idle"
@@ -184,6 +171,7 @@ async def execute(
                 error_message = f"Kernel didn't respond in {timeout} seconds"
                 raise RuntimeError(error_message)
             msg = self.execute_requests[msg_id]["shell_msg"]
+            cell["execution_count"] = msg["content"]["execution_count"]
             del self.execute_requests[msg_id]
 
     async def _wait_for_ready(self, timeout):
@@ -204,3 +192,31 @@ async def _wait_for_ready(self, timeout):
                 if msg is not None:
                     break
             new_timeout = deadline_to_timeout(deadline)
+
+    def _handle_outputs(self, outputs: List[Dict[str, Any]], msg: Dict[str, Any]):
+        msg_type = msg["header"]["msg_type"]
+        content = msg["content"]
+        if msg_type == "stream":
+            if (not outputs) or (outputs[-1]["name"] != content["name"]):
+                outputs.append({"name": content["name"], "output_type": msg_type, "text": []})
+            outputs[-1]["text"].append(content["text"])
+        elif msg_type in ("display_data", "execute_result"):
+            outputs.append(
+                {
+                    "data": {"text/plain": [content["data"].get("text/plain", "")]},
+                    "execution_count": content["execution_count"],
+                    "metadata": {},
+                    "output_type": msg_type,
+                }
+            )
+        elif msg_type == "error":
+            outputs.append(
+                {
+                    "ename": content["ename"],
+                    "evalue": content["evalue"],
+                    "output_type": "error",
+                    "traceback": content["traceback"],
+                }
+            )
+        else:
+            return

plugins/kernels/fps_kernels/routes.py

@@ -2,9 +2,7 @@
 import pathlib
 import sys
 import uuid
-from functools import partial
 from http import HTTPStatus
-from typing import Any, Dict
 
 from fastapi import APIRouter, Depends, Response, WebSocket, status
 from fastapi.responses import FileResponse
@@ -201,41 +199,10 @@ async def execute_cell(
             await driver.connect()
         driver = kernel["driver"]
 
-        await driver.execute(cell["source"], output_hook=partial(output_hook, cell["outputs"]))
+        await driver.execute(cell)
         room.document.source = nb
 
 
-def output_hook(outputs, msg: Dict[str, Any]):
-    # msg_id = msg["parent_header"]["msg_id"]
-    execution_count = 1  # self.msg_id_2_execution_count[msg_id]
-    msg_type = msg["header"]["msg_type"]
-    content = msg["content"]
-    if msg_type == "stream":
-        if (not outputs) or (outputs[-1]["name"] != content["name"]):
-            outputs.append({"name": content["name"], "output_type": msg_type, "text": []})
-        outputs[-1]["text"].append(content["text"])
-    elif msg_type in ("display_data", "execute_result"):
-        outputs.append(
-            {
-                "data": {"text/plain": [content["data"].get("text/plain", "")]},
-                "execution_count": execution_count,
-                "metadata": {},
-                "output_type": msg_type,
-            }
-        )
-    elif msg_type == "error":
-        outputs.append(
-            {
-                "ename": content["ename"],
-                "evalue": content["evalue"],
-                "output_type": "error",
-                "traceback": content["traceback"],
-            }
-        )
-    else:
-        return
-
-
 @router.get("/api/kernels/{kernel_id}")
 async def get_kernel(
     kernel_id,

tests/data/notebook0.ipynb

@@ -19,6 +19,16 @@
       "metadata": {
         "trusted": 0
       }
+    },
+    {
+      "source": "3 + 4",
+      "outputs": [],
+      "metadata": {
+        "trusted": 0
+      },
+      "cell_type": "code",
+      "execution_count": null,
+      "id": "a7243792-6f06-4462-a6b5-7e9ec604348f"
     }
   ],
   "metadata": {

tests/test_server.py

@@ -71,7 +71,7 @@ async def test_rest_api(start_jupyverse):
         # wait for file to be loaded and Y model to be created in server and client
         await asyncio.sleep(0.1)
         # execute notebook
-        for cell_idx in range(2):
+        for cell_idx in range(3):
             response = requests.post(
                 f"{url}/api/kernels/{kernel_id}/execute",
                 data=json.dumps(
@@ -96,3 +96,11 @@ async def test_rest_api(start_jupyverse):
         assert cells[1]["outputs"] == [
             {"name": "stdout", "output_type": "stream", "text": ["Hello World!\n"]}
         ]
+        assert cells[2]["outputs"] == [
+            {
+                "data": {"text/plain": ["7"]},
+                "execution_count": 3.0,
+                "metadata": {},
+                "output_type": "execute_result",
+            }
+        ]