re-run edited cells automatically when they lose focus

packages/server/src/server.ts

@@ -298,18 +298,18 @@ class VitaleDevServer {
 
   private invalidateModule(
     id: string,
-    dirtyCells: { path: string; cellId: string }[]
+    dirtyCells: { path: string; cellId: string; ext: string }[]
   ) {
     const mod = this.viteRuntime.moduleCache.get(id);
     this.viteRuntime.moduleCache.delete(id);
 
     const match = cellIdRegex.exec(id);
     if (match) {
-      const [_, path, cellId] = match;
+      const [_, path, cellId, ext] = match;
       if (
         !dirtyCells.some((cell) => cell.path === path && cell.cellId === cellId)
       ) {
-        dirtyCells.push({ path, cellId });
+        dirtyCells.push({ path, cellId, ext });
       }
     }
 
@@ -331,20 +331,21 @@ class VitaleDevServer {
   }
 
   private invalidateModuleAndDirty(id: string) {
-    const cells: { path: string; cellId: string }[] = [];
+    const cells: { path: string; cellId: string; ext: string }[] = [];
     this.invalidateModule(id, cells);
     this.markCellsDirty(cells);
   }
 
-  private executeCellsRPC(
+  private async executeCellsRPC(
     cells: {
       path: string;
       cellId: string;
       language: string;
       code?: string;
-    }[]
+    }[],
+    executeDirtyCells: boolean
   ) {
-    let dirtyCells: { path: string; cellId: string }[] = [];
+    let dirtyCells: { path: string; cellId: string; ext: string }[] = [];
 
     for (const { path, cellId, language, code } of cells) {
       const ext = extOfLanguage(language);
@@ -362,23 +363,35 @@ class VitaleDevServer {
       this.invalidateModule(id, dirtyCells);
     }
 
-    for (const { path, cellId, language } of cells) {
-      const ext = extOfLanguage(language);
-      const id = `${path}-cellId=${cellId}.${ext}`;
-      this.executeCell(id, path, cellId).catch((e) => {
-        console.error(e);
-      });
-    }
-
-    // don't mark cells dirty if they were just executed
     dirtyCells = dirtyCells.filter(
       (dirtyCell) =>
         !cells.some(
           (cell) =>
             cell.path === dirtyCell.path && cell.cellId === dirtyCell.cellId
         )
     );
-    this.markCellsDirty(dirtyCells);
+
+    const cellsToExecute = [
+      ...cells.map(({ path, cellId, language }) => ({
+        path,
+        cellId,
+        ext: extOfLanguage(language),
+      })),
+      ...(executeDirtyCells ? dirtyCells : []),
+    ];
+
+    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);
+        });
+      })
+    );
+
+    if (!executeDirtyCells) {
+      this.markCellsDirty(dirtyCells);
+    }
   }
 
   private removeCellsRPC(
@@ -388,7 +401,7 @@ class VitaleDevServer {
       language: string;
     }[]
   ) {
-    let dirtyCells: { path: string; cellId: string }[] = [];
+    let dirtyCells: { path: string; cellId: string; ext: string }[] = [];
 
     for (const { path, cellId, language } of cells) {
       const ext = extOfLanguage(language);
@@ -423,9 +436,9 @@ class VitaleDevServer {
           console.log("ping");
           return "pong";
         },
-        async executeCells(cells) {
+        async executeCells(cells, executeDirtyCells) {
           try {
-            return self.executeCellsRPC(cells);
+            return self.executeCellsRPC(cells, executeDirtyCells);
           } catch (e) {
             console.error(e);
           }

packages/server/src/types.ts

@@ -26,7 +26,8 @@ export type ServerFunctions = {
       cellId: string;
       language: string;
       code?: string;
-    }[]
+    }[],
+    executeDirtyCells: boolean
   ) => void;
 };
 

packages/vscode/src/controller.ts

@@ -21,7 +21,7 @@ function cellOutputToNotebookCellOutput(cellOutput: CellOutput) {
   );
 }
 
-function getRerunCellsWhenDirty() {
+export function getRerunCellsWhenDirty() {
   const config = vscode.workspace.getConfiguration("vitale");
   return config.get("rerunCellsWhenDirty", true);
 }
@@ -241,7 +241,9 @@ export class NotebookController {
     const notebook = await vscode.workspace.openNotebookDocument(
       vscode.Uri.parse(notebookUri)
     );
-    const cells = notebook.getCells().filter((cell) => cell.metadata.dirty);
+    const cells = notebook
+      .getCells()
+      .filter((cell) => cell.metadata.dirty || cell.metadata.docDirty);
     this.executeCells(cells);
   }
 
@@ -322,9 +324,6 @@ export class NotebookController {
     await Promise.all(
       notebookCells.map((cell) => this.setCellDirty(cell, true))
     );
-    if (getRerunCellsWhenDirty()) {
-      this.executeCells(notebookCells, false);
-    }
   }
 
   private async startCellExecution(path: string, id: string) {
@@ -378,28 +377,23 @@ export class NotebookController {
     this.executeCells(notebookCells);
   }
 
-  private async executeCells(
-    notebookCells: vscode.NotebookCell[],
-    sendDirtyDocs = true
-  ) {
+  private async executeCells(notebookCells: vscode.NotebookCell[]) {
     if (notebookCells.length === 0) {
       return;
     }
     const cells = notebookCells.map((cell) => ({
       path: cell.notebook.uri.fsPath,
       cellId: cell.metadata.id,
       language: cell.document.languageId,
-      code: sendDirtyDocs ? cell.document.getText() : undefined,
+      code: cell.document.getText(),
     }));
 
-    if (sendDirtyDocs) {
-      await Promise.all(
-        notebookCells.map((cell) => this.setCellDocDirty(cell, false))
-      );
-    }
+    await Promise.all(
+      notebookCells.map((cell) => this.setCellDocDirty(cell, false))
+    );
 
     const client = await this.getClient();
-    client.executeCells(cells);
+    client.executeCells(cells, getRerunCellsWhenDirty());
   }
 
   async removeCells(notebookCells: vscode.NotebookCell[]) {

packages/vscode/src/extension.ts

@@ -3,6 +3,7 @@ import { NotebookSerializer } from "./serializer";
 import { NotebookController } from "./controller";
 import { NotebookCellStatusBarItemProvider } from "./cellStatusBarItemProvider";
 import { makeHandleDidChangeNotebookDocument } from "./handleDidChangeNotebookDocument";
+import { makeHandleDidChangeNotebookEditorSelection } from "./handleDidChangeNotebookEditorSelection";
 import { handleDidChangeTextDocument } from "./handleDidChangeTextDocument";
 
 export function activate(context: vscode.ExtensionContext) {
@@ -32,6 +33,9 @@ export function activate(context: vscode.ExtensionContext) {
       makeHandleDidChangeNotebookDocument(controller)
     ),
     vscode.workspace.onDidChangeTextDocument(handleDidChangeTextDocument),
+    vscode.window.onDidChangeNotebookEditorSelection(
+      makeHandleDidChangeNotebookEditorSelection(controller)
+    ),
     controller,
     vscode.notebooks.registerNotebookCellStatusBarItemProvider(
       "vitale-notebook",

packages/vscode/src/handleDidChangeNotebookEditorSelection.ts

@@ -0,0 +1,18 @@
+import type { NotebookEditorSelectionChangeEvent } from "vscode";
+import { getRerunCellsWhenDirty } from "./controller";
+import type { NotebookController } from "./controller";
+
+export function makeHandleDidChangeNotebookEditorSelection(
+  controller: NotebookController
+) {
+  return (e: NotebookEditorSelectionChangeEvent) => {
+    // what we really want here is to run when an edited cell loses focus
+    // but it doesn't seem to be possible in the VS Code API
+    // 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());
+    }
+  };
+}