Tinkering Towards a JupyterLab Cell Status Indicator Extension

Having managed to build one extension (a JupyterLab version of empinken), I thought I’d have ago an another, a JuptyerLab version of the classic notebook innovationoutside/nb_cell_execution_status extension which colours a code cell run indication according to whether the cell is running (or queued for execution), has run successfully, or ran with an error.

You can try it, via JuptyerLite (howto), here: innovationoutside/nb_cell_execution_status

The code in its minimal form is really simple, although that isn’t to say it didn’t take me a disproportionate amount of time to find the methods I guessed might help implement the feature and then try to get the actual extension compiled and working. I still really, really, really struggle with the whole JupyterLab thing. The more I use it, the more I don’t want to, if that were possible…!;-)

As a starting point, I used the JupyterLab Typescript extension cookiecutter (e.g. as per here). I needed to import one additional package, import { NotebookActions } from '@jupyterlab/notebook';, which also need adding in the dependencies of the package.json file as "@jupyterlab/notebook": "^3.3.4".

I would have had a nice git diff from a single check-in after the “base” cookiecutter check-in, showing just what was required. But the reality is that there are dozens of commits in the repo showing the numerous craptastic and ignorant end-user-dev attempts I kept making to try to get things to work. Have I mentioned before how much I dislike eveything about JupyterLab, and its development process?!;-)

Essentially, all the extension requires that you subscribe to signals that fire when a code cell has been queued for execution, or when it completes execution.

Specifically, when one or more code cells are selected and run, a NotebookActions.executionScheduled signal is issued. We can subscribe to this signal, remove any classes previously added to the cell DOM in the notebook HTML that relate to the status of previous successful or unsuccessful cell execution (cell.inputArea.promptNode.classList.remove()), and add a new class (cell.inputArea.promptNode.classList.add()) to indicate a scheduled status.

When the cell has completed execution, a NotebookActions.executed signal raised. This also carries with it information regarding whether the cell executed successfully or not; we can remove the scheduled-or-executing class from the cell DOM in the notebook HTML, and then use this cell executed status to set a class identifying the executed status.

Note, there are also various notebook API methods, but I’m not really sure how to tap into those appropriately….

import {
  JupyterFrontEnd,
  JupyterFrontEndPlugin
} from '@jupyterlab/application';

//https://jupyterlab.readthedocs.io/en/3.1.x/api/classes/notebook.notebookactions-1.html
import { NotebookActions } from '@jupyterlab/notebook';
//import { IKernelConnection } from '@jupyterlab/services/Kernel IKernelConnection';
import { ISettingRegistry } from '@jupyterlab/settingregistry';

/**
 * Initialization data for the jupyterlab_cell_status extension.
 */
const plugin: JupyterFrontEndPlugin<void> = {
  activate,
  id: 'jupyterlab_cell_status_extension:plugin',
  autoStart: true,
  optional: [ISettingRegistry],//, IKernelConnection]
};

function activate (
  app: JupyterFrontEnd, settingRegistry: ISettingRegistry | null): void {
    console.log("jupyterlab_cell_status_extension:plugin activating...");
    /*
    // We can use settings to set the status colours
    if (settingRegistry) {
      settingRegistry
        .load(plugin.id)
        .then(settings => {
          console.log("jupyterlab_cell_status_extension:plugin: loading settings...");
          const root = document.documentElement;
          const updateSettings = (): void => {
            const queue_color = settings.get('status_queue').composite as string;
            const success_color = settings.get('status_success').composite as string;
            const error_color = settings.get('status_error').composite as string;
            
            root.style.setProperty('--jp-cell-status-queue', queue_color);
            root.style.setProperty('--jp-cell-status-success', success_color);
            root.style.setProperty('--jp-cell-status-error', error_color);
            
          };
          updateSettings();
          console.log("jupyterlab_cell_status_extension:plugin: loaded settings...");
          // We can auto update the color
          settings.changed.connect(updateSettings);
        })
        .catch(reason => {
          console.error('Failed to load settings for jupyterlab_cell_status_extension.', reason);
        });
    }
    */

    /*
    // This was a start at sketching whether I could
    // reset things if the kernel was restarted.
    // Didn't get very far though?
    IKernelConnection.connectionStatusChanged.connect((kernel, conn_stat) => {

    console.log("KERNEL ****"+conn_stat)
    });
    */
    NotebookActions.executed.connect((_, args) => {
      // The following construction seems to say 
      // something akin to: const cell = args["cell"]
      const { cell } = args;
      const { success } = args;
      // If we have a code cell, update the status
      if (cell.model.type == 'code') {
        if (success)
          cell.inputArea.promptNode.classList.add("executed-success");
        else
          cell.inputArea.promptNode.classList.add("executed-error");
          cell.inputArea.promptNode.classList.remove("scheduled");
      }
    });

    NotebookActions.executionScheduled.connect((_, args) => {
      const { cell } = args;
      // If we have a code cell
      // set the status class to "scheduled"
      // and remove the other classes
      if (cell.model.type == 'code') {
        cell.inputArea.promptNode.classList.remove("executed-success");
        cell.inputArea.promptNode.classList.remove("executed-error");
        cell.inputArea.promptNode.classList.add("scheduled");
      }
    });  

    console.log("jupyterlab_cell_status_extension:plugin activated...");
  }

export default plugin;

To build a Python wheel for an installable version of the extension, I use the build command:

jlpm clean && jlpm run build && python3 -m build

When cell is running, or queued, by default the cell status area now highlights a blie. Successful cell execution is denote green, and unsuccessful cell execution is red.

Another disproportionate amount of time was spent trying to figure out why the settings weren’t being loaded. I tracked this down to a mismatch between the extension id defined in the index.ts file and the name in the schema/pugin.json file. It still didn’t work, but a repeated “comment out all the settings code, rebuild, commit to Github, republish the JupyerLite distro” process, uncommenting a line at a time revealed: no errors, and it had started working. So I have no idea what went wrong.

If the itereated process sounds faintly stupid, and it took maybe 5 minutes each iteration (two to build the extension, two to republish the JupyterLite distribution). That’s partly because I was suspicious that my local build process wasn’t looking on the correct path. (I know that JupyterLite builds don’t respect paths on my local machine but work fine when run on Github via a Github Action.)

Anyway, the settings seem to work now to allow you to set the status indicator colours, though after a change you need to close a notebook then open it again to see the change reflected.

If you close a notebook and open it again, whether or not you stop the kernel, then the status colouring is lost. I did start to wonder whether I should persist the status as notebook metadata (eg how to set tags can be found here) to allow a notebook opened against a still running kernel to show the cell status (classing the cell based on cell metadata), and then managing persistent tags and classes based on cell execution status and notebook kernel signals (eg removing all status indications if the kernel is stopped. I’m not sure how to access the notebook/cells from inside a IKernelConnection.connectionStatusChanged connected handler though? (My thinking is the cell status indicators should give you a reasonable indication of whether a particular cell has contributed to the kernel state. I wonder if there’s any code in nbsafety-project/nbsafety that might help me think this through a bit more?)

I guess the next step is to see if I can also add the audio indicators that I’d added to the classic notebook extension too…

Author: Tony Hirst

I'm a Senior Lecturer at The Open University, with an interest in #opendata policy and practice, as well as general web tinkering...

%d bloggers like this: