Use cell ID instead of cell index in execute API

jupyter-server · Nov 21, 2023 · 6a41166 · 6a41166
1 parent 757a90e
commit 6a41166
Show file tree

Hide file tree

Showing 4 changed files with 13 additions and 10 deletions.
diff --git a/jupyverse_api/jupyverse_api/kernels/models.py b/jupyverse_api/jupyverse_api/kernels/models.py
@@ -39,4 +39,4 @@ class Session(BaseModel):
 
 class Execution(BaseModel):
     document_id: str
-    cell_idx: int
+    cell_id: str
diff --git a/plugins/kernels/fps_kernels/routes.py b/plugins/kernels/fps_kernels/routes.py
@@ -259,7 +259,11 @@ async def execute_cell(
         execution = Execution(**r)
         if kernel_id in kernels:
             ynotebook = self.yjs.get_document(execution.document_id)
-            ycell = ynotebook.ycells[execution.cell_idx]
+            ycells = [ycell for ycell in ynotebook.ycells if ycell["id"] == execution.cell_id]
+            if not ycells:
+                return  # FIXME
+
+            ycell = ycells[0]
             del ycell["outputs"][:]
 
             kernel = kernels[kernel_id]

diff --git a/tests/test_execute.py b/tests/test_execute.py
@@ -113,7 +113,7 @@ def callback(aevent, events, event):
                     f"{url}/api/kernels/{kernel_id}/execute",
                     json={
                         "document_id": document_id,
-                        "cell_idx": cell_idx,
+                        "cell_id": ynb.ycells[cell_idx]["id"],
                     }
                 )
             while True:
@@ -130,7 +130,7 @@ def callback(aevent, events, event):
                 f"{url}/api/kernels/{kernel_id}/execute",
                 json={
                     "document_id": document_id,
-                    "cell_idx": 2,
+                    "cell_id": ynb.ycells[2]["id"],
                 }
             )
             await task

diff --git a/tests/test_server.py b/tests/test_server.py
@@ -88,23 +88,22 @@ async def test_rest_api(start_jupyverse):
         # connect to the shared notebook document
         # wait for file to be loaded and Y model to be created in server and client
         await asyncio.sleep(0.5)
+        ydoc["cells"] = ycells = Array()
         # execute notebook
         for cell_idx in range(3):
             response = requests.post(
                 f"{url}/api/kernels/{kernel_id}/execute",
                 data=json.dumps(
                     {
                         "document_id": document_id,
-                        "cell_idx": cell_idx,
+                        "cell_id": ycells[cell_idx]["id"],
                     }
                 ),
             )
         # wait for Y model to be updated
         await asyncio.sleep(0.5)
         # retrieve cells
-        array = Array()
-        ydoc["cells"] = array
-        cells = json.loads(str(array))
+        cells = json.loads(str(ycells))
         assert cells[0]["outputs"] == [
             {
                 "data": {"text/plain": ["3"]},
@@ -181,7 +180,7 @@ def callback(aevent, events, event):
                 data=json.dumps(
                     {
                         "document_id": document_id,
-                        "cell_idx": cell_idx,
+                        "cell_id": ynb.ycells[cell_idx]["id"],
                     }
                 ),
             )
@@ -200,7 +199,7 @@ def callback(aevent, events, event):
             data=json.dumps(
                 {
                     "document_id": document_id,
-                    "cell_idx": 2,
+                    "cell_id": ynb.ycells[2]["id"],
                 }
             ),
         )