Control when a Cell is executing #279

packages/react/src/components/cell/Cell.tsx

@@ -64,15 +64,18 @@ export const Cell = (props: ICellProps) => {
       }
 
       // Perform auto-start for code or markdown cells
+      cellStore.setIsExecuting(id, true);
       if (adapter.cell instanceof CodeCell) {
-        CodeCell.execute(
-          adapter.cell,
-          adapter.sessionContext
-        );
+        CodeCell
+          .execute(adapter.cell, adapter.sessionContext)
+          .then(() => {
+            cellStore.setIsExecuting(id, false);
+          });
       }
 
       if (adapter.cell instanceof MarkdownCell) {
         adapter.cell.rendered = true;
+        cellStore.setIsExecuting(id, false);
       }
     });
 
@@ -88,6 +91,7 @@ export const Cell = (props: ICellProps) => {
     if (id && defaultKernel && serverSettings) {
       defaultKernel.ready.then(() => {
         const adapter = new CellAdapter({
+          id,
           type,
           source,
           serverSettings,

packages/react/src/components/cell/CellAdapter.ts

@@ -52,19 +52,22 @@ import ClassicWidgetManager from '../../jupyter/ipywidgets/classic/manager';
 import Kernel from '../../jupyter/kernel/Kernel';
 import getMarked from '../notebook/marked/marked';
 import CellCommands from './CellCommands';
+import { cellStore } from './CellState';
 
 interface BoxOptions {
   showToolbar?: boolean;
 }
 export class CellAdapter {
+  private _id: string;
   private _cell: CodeCell | MarkdownCell | RawCell;
   private _kernel: Kernel;
   private _panel: BoxPanel;
   private _sessionContext: SessionContext;
   private _type: 'code' | 'markdown' | 'raw'
 
   constructor(options: CellAdapter.ICellAdapterOptions) {
-    const { type, source, serverSettings, kernel, boxOptions } = options;
+    const { id, type, source, serverSettings, kernel, boxOptions } = options;
+    this._id = id;
     this._kernel = kernel;
     this._type = type;
     this.setupCell(type, source, serverSettings, kernel, boxOptions);
@@ -303,7 +306,9 @@ export class CellAdapter {
     });
     handler.editor = editor;
 
-    CellCommands(commands, this._cell!, this._sessionContext, handler);
+    // Init general commands/complete for the cell
+    CellCommands(this._id, commands, this._cell!, this._sessionContext, handler);
+
     completer.hide();
     completer.addClass('jp-Completer-Cell');
     Widget.attach(completer, document.body);
@@ -312,10 +317,16 @@ export class CellAdapter {
     const runButton = new ToolbarButton({
       icon: runIcon,
       onClick: () => {
+        cellStore.getState().setIsExecuting(this._id, true);
         if (this._type === 'code') {
-          CodeCell.execute(this._cell as CodeCell, this._sessionContext);
+          CodeCell
+            .execute((this._cell as CodeCell), this._sessionContext)
+            .then(() => {
+              cellStore.getState().setIsExecuting(this._id, false);
+            });
         } else if (this._type === 'markdown') {
           (this._cell as MarkdownCell).rendered = true;
+          cellStore.getState().setIsExecuting(this._id, false);
         }
       },
       tooltip: 'Run',
@@ -376,16 +387,23 @@ export class CellAdapter {
   }
 
   execute = () => {
+    cellStore.getState().setIsExecuting(this._id, true);
     if (this._type === 'code') {
-      CodeCell.execute((this._cell as CodeCell), this._sessionContext);
+      CodeCell
+        .execute((this._cell as CodeCell), this._sessionContext)
+        .then(() => {
+          cellStore.getState().setIsExecuting(this._id, false);
+        });
     } else if (this._type === 'markdown') {
       (this._cell as MarkdownCell).rendered = true;
+      cellStore.getState().setIsExecuting(this._id, false);
     }
   };
 }
 
 export namespace CellAdapter {
   export type ICellAdapterOptions = {
+    id: string;
     type: 'code' | 'markdown' | 'raw';
     source: string;
     serverSettings: ServerConnection.ISettings;

packages/react/src/components/cell/CellCommands.ts

@@ -8,13 +8,15 @@ import { CommandRegistry } from '@lumino/commands';
 import { CompletionHandler } from '@jupyterlab/completer';
 import { CodeCell, MarkdownCell, RawCell } from '@jupyterlab/cells';
 import { SessionContext } from '@jupyterlab/apputils';
+import { cellStore } from './CellState';
 
 const cmdIds = {
   invoke: 'completer:invoke',
   select: 'completer:select',
 };
 
 export const CellCommands = (
+  cellId: string,
   commandRegistry: CommandRegistry,
   cell: CodeCell | MarkdownCell | RawCell,
   sessionContext: SessionContext,
@@ -30,10 +32,16 @@ export const CellCommands = (
   });
   commandRegistry.addCommand('run:cell', {
     execute: () => {
+      cellStore.getState().setIsExecuting(cellId, true);
       if (cell instanceof CodeCell) {
-        CodeCell.execute(cell, sessionContext)
+        CodeCell
+          .execute(cell, sessionContext)
+          .then(() => {
+            cellStore.getState().setIsExecuting(cellId, false);
+          })
       } else if (cell instanceof MarkdownCell) {
         (cell as MarkdownCell).rendered = true;
+        cellStore.getState().setIsExecuting(cellId, false);
       }
     },
   });

packages/react/src/components/cell/CellState.ts

@@ -8,29 +8,34 @@ import { createStore } from 'zustand/vanilla';
 import { useStore } from 'zustand';
 import CellAdapter from './CellAdapter';
 
+// Individual, for each cell
 export interface ICellState {
   source?: string;
   outputsCount?: number;
   adapter?: CellAdapter;
-  isKernelSessionAvailable?: boolean; // Individual, for cell
+  isKernelSessionAvailable?: boolean;
+  isExecuting?: boolean;
 }
 
+// For all cells
 export interface ICellsState {
   cells: Map<string, ICellState>;
-  areAllKernelSessionsReady: boolean; // Control the state for all cells
+  areAllKernelSessionsReady: boolean; // Control if all cells has sessions available
+  isAnyCellExecuting: boolean; // Control if at least one cell is executing
 }
 
 export type CellState = ICellsState & {
   setCells: (cells: Map<string, ICellState>) => void;
   setSource: (id: string, source: string) => void;
   setOutputsCount: (id: string, outputsCount: number) => void;
-  setIsKernelSessionAvailable: (id: string, kernelAvailable: boolean) => void;
   setAdapter: (id: string, adapter?: CellAdapter) => void;
-  getAdapter: (id: string) => CellAdapter | undefined;
+  setIsKernelSessionAvailable: (id: string, kernelAvailable: boolean) => void;
   getSource: (id: string) => string | undefined;
   getOutputsCount: (id: string) => number | undefined;
+  getAdapter: (id: string) => CellAdapter | undefined;
   getIsKernelSessionAvailable: (id: string) => boolean | undefined;
   execute: (id?: string) => void;
+  setIsExecuting: (id: string, isExecuting: boolean) => void; // Necessary for the cases when execute directly from jupyterlab/cells
 };
 
 /**
@@ -45,13 +50,27 @@ const areAllKernelSessionsAvailable = (cells: Map<string, ICellState>): boolean
   return true;
 };
 
+/**
+ * Check if any cell is currently executing
+ */
+export const isAnyCellRunning = (cells: Map<string, ICellState>): boolean => {
+  for (const cell of cells.values()) {
+    if (cell.isExecuting) {
+      return true;
+    }
+  }
+  return false;
+};
+
 export const cellStore = createStore<CellState>((set, get) => ({
   cells: new Map<string, ICellState>(),
   source: '',
   outputsCount: 0,
-  isKernelSessionAvailable: false,
-  areAllKernelSessionsReady: false,
   adapter: undefined,
+  isKernelSessionAvailable: false,
+  isExecuting: false,
+  areAllKernelSessionsReady: false, // prop refers to all cells
+  isAnyCellExecuting: false, // prop refers to all cells
   setCells: (cells: Map<string, ICellState>) => set((cell: CellState) => ({ cells })),
 
   setSource: (id: string, source: string) => {
@@ -111,10 +130,29 @@ export const cellStore = createStore<CellState>((set, get) => ({
     const cells = get().cells;
     const cell = cells.get(id);
     if (cell) {
-      cell.adapter?.execute() 
+      cell.adapter?.execute();
+    } else {
+      get().cells.forEach((cell) => {
+        cell.adapter?.execute();
+      });
+    }
+    set((state: CellState) => ({ cells }));
+  },
+  setIsExecuting: (id: string, isExecuting: boolean) => {
+    /**
+     * Set is executing
+     * This is necessary to update state in cases when we execute directly from jupyterlab/cells
+     */
+    const cells = get().cells;
+    const cell = cells.get(id);
+    if (cell) {
+      cell.isExecuting = isExecuting;
     } else {
-      get().cells.forEach((cell) => cell.adapter?.execute())
+      cells.set(id, { isExecuting });
     }
+
+    const isAnyCellExecuting = isAnyCellRunning(cells);
+    set((state: CellState) => ({ cells, isAnyCellExecuting }));
   },
 }));