make cells pause-able

githubnext · May 7, 2024 · 361bd12 · 361bd12
1 parent f7fc5a8
commit 361bd12
Show file tree

Hide file tree

Showing 7 changed files with 115 additions and 46 deletions.
diff --git a/packages/server/src/server.ts b/packages/server/src/server.ts
@@ -207,15 +207,25 @@ class VitaleDevServer {
     });
   }
 
-  private async executeCell(id: string, path: string, cellId: string) {
+  private async executeCell(
+    id: string,
+    path: string,
+    cellId: string,
+    force: boolean
+  ) {
     // TODO(jaked)
     // await so client finishes startCellExecution before we send endCellExecution
     // would be better for client to lock around startCellExecution
-    await Promise.all(
-      Array.from(this.clients.values()).map((client) =>
-        client.startCellExecution(path, cellId)
+    const startOK = (
+      await Promise.all(
+        Array.from(this.clients.values()).map((client) =>
+          client.startCellExecution(path, cellId, force)
+        )
       )
-    );
+    ).every((ok) => ok);
+    if (!startOK) {
+      return false;
+    }
 
     let data;
     let mime;
@@ -289,11 +299,12 @@ class VitaleDevServer {
           : [{ data: [...Buffer.from(data, "utf8").values()], mime }],
     };
 
-    return await Promise.all(
+    await Promise.all(
       Array.from(this.clients.values()).map((client) =>
         client.endCellExecution(path, cellId, cellOutput)
       )
     );
+    return true;
   }
 
   private invalidateModule(
@@ -343,6 +354,7 @@ class VitaleDevServer {
       language: string;
       code?: string;
     }[],
+    force: boolean,
     executeDirtyCells: boolean
   ) {
     let dirtyCells: { path: string; cellId: string; ext: string }[] = [];
@@ -376,22 +388,23 @@ class VitaleDevServer {
         path,
         cellId,
         ext: extOfLanguage(language),
+        force,
       })),
-      ...(executeDirtyCells ? dirtyCells : []),
+      ...(executeDirtyCells
+        ? dirtyCells.map((cell) => ({ ...cell, force: false }))
+        : []),
     ];
 
-    await Promise.all(
-      cellsToExecute.map(async ({ path, cellId, ext }) => {
-        const id = `${path}-cellId=${cellId}.${ext}`;
-        return this.executeCell(id, path, cellId).catch((e) => {
-          console.error(e);
-        });
-      })
+    const executed = await Promise.all(
+      cellsToExecute.map(({ path, cellId, ext, force }) =>
+        this.executeCell(`${path}-cellId=${cellId}.${ext}`, path, cellId, force)
+      )
     );
 
-    if (!executeDirtyCells) {
-      this.markCellsDirty(dirtyCells);
-    }
+    const cellsToMarkDirty = cellsToExecute.filter(
+      ({ force }, i) => !force && !executed[i]
+    );
+    this.markCellsDirty(cellsToMarkDirty);
   }
 
   private removeCellsRPC(
@@ -436,9 +449,9 @@ class VitaleDevServer {
           console.log("ping");
           return "pong";
         },
-        async executeCells(cells, executeDirtyCells) {
+        async executeCells(cells, force, executeDirtyCells) {
           try {
-            return self.executeCellsRPC(cells, executeDirtyCells);
+            return self.executeCellsRPC(cells, force, executeDirtyCells);
           } catch (e) {
             console.error(e);
           }

diff --git a/packages/server/src/types.ts b/packages/server/src/types.ts
@@ -27,13 +27,18 @@ export type ServerFunctions = {
       language: string;
       code?: string;
     }[],
+    force: boolean,
     executeDirtyCells: boolean
   ) => void;
 };
 
 export type ClientFunctions = {
   markCellsDirty: (cells: { path: string; cellId: string }[]) => void;
-  startCellExecution: (path: string, cellId: string) => void;
+  startCellExecution: (
+    path: string,
+    cellId: string,
+    force: boolean
+  ) => Promise<boolean>;
   endCellExecution: (
     path: string,
     cellId: string,

diff --git a/packages/vscode/src/cellStatusBarItemProvider.ts b/packages/vscode/src/cellStatusBarItemProvider.ts
@@ -27,6 +27,19 @@ export class NotebookCellStatusBarItemProvider
       dirtyItem.command = "notebook.cell.execute";
       items.push(dirtyItem);
     }
+
+    const pauseItem = new vscode.NotebookCellStatusBarItem(
+      cell.metadata.paused ? "$(debug-start)" : "$(debug-pause)",
+      vscode.NotebookCellStatusBarAlignment.Right
+    );
+    pauseItem.tooltip = cell.metadata.paused ? "Unpause" : "Pause";
+    pauseItem.command = {
+      title: "Pause",
+      command: "vitale.pauseCell",
+      arguments: [cell.notebook.uri, cell.metadata.id],
+    };
+    items.push(pauseItem);
+
     return items;
   }
 }
diff --git a/packages/vscode/src/controller.ts b/packages/vscode/src/controller.ts
@@ -327,31 +327,36 @@ export class NotebookController {
     );
   }
 
-  private async startCellExecution(path: string, id: string) {
+  private async startCellExecution(path: string, id: string, force: boolean) {
     const notebook = await vscode.workspace.openNotebookDocument(
       vscode.Uri.file(path)
     );
     const cell = notebook.getCells().find((cell) => cell.metadata.id === id);
-    if (cell) {
-      const key = `${path}-${id}`;
-      if (this._executions.has(key)) {
-        // this can happen when the user edits the cell
-        // thens executes it
-        // because handleDidChangeNotebookEditorSelection fires
-        // and calls runDirty
-        log.info(`already executing ${key}`);
-      } else {
-        const execution = this._controller.createNotebookCellExecution(cell);
-        execution.token.onCancellationRequested(() => {
-          this.cancelCellExecution(path, id);
-        });
-
-        execution.executionOrder = ++this._executionOrder;
-        execution.start(Date.now());
-
-        this._executions.set(key, execution);
-      }
+
+    if (!cell || (cell.metadata.paused && !force)) {
+      return false;
+    }
+
+    const key = `${path}-${id}`;
+    if (this._executions.has(key)) {
+      // this can happen when the user edits the cell
+      // thens executes it
+      // because handleDidChangeNotebookEditorSelection fires
+      // and calls runDirty
+      log.info(`already executing ${key}`);
+      return false;
     }
+
+    const execution = this._controller.createNotebookCellExecution(cell);
+    execution.token.onCancellationRequested(() => {
+      this.cancelCellExecution(path, id);
+    });
+
+    execution.executionOrder = ++this._executionOrder;
+    execution.start(Date.now());
+
+    this._executions.set(key, execution);
+    return true;
   }
 
   private cancelCellExecution(path: string, id: string) {
@@ -383,10 +388,13 @@ export class NotebookController {
   }
 
   private executeHandler(notebookCells: vscode.NotebookCell[]) {
-    this.executeCells(notebookCells);
+    this.executeCells(notebookCells, true);
   }
 
-  private async executeCells(notebookCells: vscode.NotebookCell[]) {
+  private async executeCells(
+    notebookCells: vscode.NotebookCell[],
+    force: boolean = false
+  ) {
     if (notebookCells.length === 0) {
       return;
     }
@@ -402,7 +410,7 @@ export class NotebookController {
     );
 
     const client = await this.getClient();
-    client.executeCells(cells, getRerunCellsWhenDirty());
+    client.executeCells(cells, force, getRerunCellsWhenDirty());
   }
 
   async removeCells(notebookCells: vscode.NotebookCell[]) {

diff --git a/packages/vscode/src/extension.ts b/packages/vscode/src/extension.ts
@@ -24,6 +24,28 @@ export function activate(context: vscode.ExtensionContext) {
     vscode.commands.registerCommand("vitale.copyToClipboard", (s: string) => {
       vscode.env.clipboard.writeText(s);
     }),
+    vscode.commands.registerCommand(
+      "vitale.pauseCell",
+      async (notebookUri: vscode.Uri, cellId: string) => {
+        const notebook = await vscode.workspace.openNotebookDocument(
+          notebookUri
+        );
+        const cell = notebook
+          .getCells()
+          .find((cell) => cell.metadata.id === cellId);
+        if (!cell) {
+          return false;
+        }
+
+        const paused = !cell.metadata.paused;
+        const metadata = { ...cell.metadata, paused };
+        const edit = new vscode.WorkspaceEdit();
+        edit.set(notebookUri, [
+          vscode.NotebookEdit.updateCellMetadata(cell.index, metadata),
+        ]);
+        return vscode.workspace.applyEdit(edit);
+      }
+    ),
     vscode.workspace.registerNotebookSerializer(
       "vitale-notebook",
       new NotebookSerializer(),

diff --git a/packages/vscode/src/handleDidChangeNotebookEditorSelection.ts b/packages/vscode/src/handleDidChangeNotebookEditorSelection.ts
@@ -11,8 +11,15 @@ export function makeHandleDidChangeNotebookEditorSelection(
     // so instead we run when the notebook selection changes
     // (i.e. you switch to a different cell)
     // but unfortunately this doesn't fire when you click outside any cell
-    if (getRerunCellsWhenDirty()) {
-      controller.runDirty(e.notebookEditor.notebook.uri.toString());
-    }
+
+    // this event fires when you click on cell status bar items
+    // so wait a bit to let the click handler fire (and possibly update metadata)
+    // before running dirty cells
+    // TODO(jaked) ugh
+    setTimeout(() => {
+      if (getRerunCellsWhenDirty()) {
+        controller.runDirty(e.notebookEditor.notebook.uri.toString());
+      }
+    }, 100);
   };
 }
diff --git a/packages/vscode/src/serializer.ts b/packages/vscode/src/serializer.ts
@@ -5,6 +5,7 @@ interface NotebookCellMetadata {
   id: string;
   dirty: boolean; // a dependency of cell has changed
   docDirty: boolean; // cell document has changed
+  paused: boolean;
 }
 
 interface NotebookCell {