Widgets support in JupyterLab console

Co-authored-by: Duc Trung LE <leductrungxf@gmail.com>

Author: 2 people

.gitignore

@@ -49,3 +49,5 @@ ui-tests/playwright-report
 **/lite/.cache
 **/*.doit.*
 **/docs/typedoc/
+
+.yarn

python/jupyterlab_widgets/package.json

@@ -51,6 +51,8 @@
     "@jupyter-widgets/controls": "^5.0.8",
     "@jupyter-widgets/output": "^6.0.7",
     "@jupyterlab/application": "^3.0.0 || ^4.0.0",
+    "@jupyterlab/apputils": "^3.0.0 || ^4.0.0",
+    "@jupyterlab/console": "^3.0.0 || ^4.0.0",
     "@jupyterlab/docregistry": "^3.0.0 || ^4.0.0",
     "@jupyterlab/logconsole": "^3.0.0 || ^4.0.0",
     "@jupyterlab/mainmenu": "^3.0.0 || ^4.0.0",
@@ -65,7 +67,6 @@
     "@lumino/algorithm": "^1.11.1 || ^2.0.0",
     "@lumino/coreutils": "^1.11.1 || ^2.1",
     "@lumino/disposable": "^1.10.1 || ^2.1",
-    "@lumino/properties": "^1.8.1 || ^2.1",
     "@lumino/signaling": "^1.10.1 || ^2.1",
     "@lumino/widgets": "^1.30.0 || ^2.1",
     "@types/backbone": "1.4.14",

python/jupyterlab_widgets/src/plugin.ts

@@ -2,12 +2,16 @@
 // Distributed under the terms of the Modified BSD License.
 
 import { ISettingRegistry } from '@jupyterlab/settingregistry';
+
 import * as nbformat from '@jupyterlab/nbformat';
 
-import { DocumentRegistry } from '@jupyterlab/docregistry';
+import {
+  IConsoleTracker,
+  CodeConsole,
+  ConsolePanel,
+} from '@jupyterlab/console';
 
 import {
-  INotebookModel,
   INotebookTracker,
   Notebook,
   NotebookPanel,
@@ -30,11 +34,13 @@ import { filter } from '@lumino/algorithm';
 
 import { DisposableDelegate } from '@lumino/disposable';
 
-import { AttachedProperty } from '@lumino/properties';
-
 import { WidgetRenderer } from './renderer';
 
-import { WidgetManager, WIDGET_VIEW_MIMETYPE } from './manager';
+import {
+  WidgetManager,
+  WIDGET_VIEW_MIMETYPE,
+  KernelWidgetManager,
+} from './manager';
 
 import { OutputModel, OutputView, OUTPUT_WIDGET_VERSION } from './output';
 
@@ -48,6 +54,7 @@ import '@jupyter-widgets/base/css/index.css';
 import '@jupyter-widgets/controls/css/widgets-base.css';
 import { KernelMessage } from '@jupyterlab/services';
 import { ITranslator, nullTranslator } from '@jupyterlab/translation';
+import { ISessionContext } from '@jupyterlab/apputils';
 
 const WIDGET_REGISTRY: base.IWidgetRegistryData[] = [];
 
@@ -59,7 +66,7 @@ const SETTINGS: WidgetManager.Settings = { saveState: false };
 /**
  * Iterate through all widget renderers in a notebook.
  */
-function* widgetRenderers(
+function* notebookWidgetRenderers(
   nb: Notebook
 ): Generator<WidgetRenderer, void, unknown> {
   for (const cell of nb.widgets) {
@@ -77,6 +84,25 @@ function* widgetRenderers(
   }
 }
 
+/**
+ * Iterate through all widget renderers in a console.
+ */
+function* consoleWidgetRenderers(
+  console: CodeConsole
+): Generator<WidgetRenderer, void, unknown> {
+  for (const cell of toArray(console.cells)) {
+    if (cell.model.type === 'code') {
+      for (const codecell of (cell as unknown as CodeCell).outputArea.widgets) {
+        for (const output of toArray(codecell.children())) {
+          if (output instanceof WidgetRenderer) {
+            yield output;
+          }
+        }
+      }
+    }
+  }
+}
+
 /**
  * Iterate through all matching linked output views
  */
@@ -109,16 +135,69 @@ function* chain<T>(
   }
 }
 
-export function registerWidgetManager(
-  context: DocumentRegistry.IContext<INotebookModel>,
+/**
+ * Get the kernel id of current notebook or console panel, this value
+ * is used as key for `Private.widgetManagerProperty` to store the widget
+ * manager of current notebook or console panel.
+ *
+ * @param {ISessionContext} sessionContext The session context of notebook or
+ * console panel.
+ */
+async function getWidgetManagerOwner(
+  sessionContext: ISessionContext
+): Promise<Private.IWidgetManagerOwner> {
+  await sessionContext.ready;
+  return sessionContext.session!.kernel!.id;
+}
+
+/**
+ * Common handler for registering both notebook and console
+ * `WidgetManager`
+ *
+ * @param {(Notebook | CodeConsole)} content Context of panel.
+ * @param {ISessionContext} sessionContext Session context of panel.
+ * @param {IRenderMimeRegistry} rendermime Rendermime of panel.
+ * @param {IterableIterator<WidgetRenderer>} renderers Iterator of
+ * `WidgetRenderer` inside panel
+ * @param {(() => WidgetManager | KernelWidgetManager)} widgetManagerFactory
+ * function to create widget manager.
+ */
+async function registerWidgetHandler(
+  content: Notebook | CodeConsole,
+  sessionContext: ISessionContext,
   rendermime: IRenderMimeRegistry,
-  renderers: IterableIterator<WidgetRenderer>
-): DisposableDelegate {
-  let wManager = Private.widgetManagerProperty.get(context);
+  renderers: IterableIterator<WidgetRenderer>,
+  widgetManagerFactory: () => WidgetManager | KernelWidgetManager
+): Promise<DisposableDelegate> {
+  const wManagerOwner = await getWidgetManagerOwner(sessionContext);
+  let wManager = Private.widgetManagerProperty.get(wManagerOwner);
+  let currentOwner: string;
+
   if (!wManager) {
-    wManager = new WidgetManager(context, rendermime, SETTINGS);
+    wManager = widgetManagerFactory();
     WIDGET_REGISTRY.forEach((data) => wManager!.register(data));
-    Private.widgetManagerProperty.set(context, wManager);
+    Private.widgetManagerProperty.set(wManagerOwner, wManager);
+    currentOwner = wManagerOwner;
+    content.disposed.connect((_) => {
+      const currentwManager = Private.widgetManagerProperty.get(currentOwner);
+      if (currentwManager) {
+        Private.widgetManagerProperty.delete(currentOwner);
+      }
+    });
+
+    sessionContext.kernelChanged.connect((_, args) => {
+      const { newValue } = args;
+      if (newValue) {
+        const newKernelId = newValue.id;
+        const oldwManager = Private.widgetManagerProperty.get(currentOwner);
+
+        if (oldwManager) {
+          Private.widgetManagerProperty.delete(currentOwner);
+          Private.widgetManagerProperty.set(newKernelId, oldwManager);
+        }
+        currentOwner = newKernelId;
+      }
+    });
   }
 
   for (const r of renderers) {
@@ -145,6 +224,45 @@ export function registerWidgetManager(
   });
 }
 
+export async function registerWidgetManager(
+  panel: NotebookPanel,
+  renderers: IterableIterator<WidgetRenderer>
+): Promise<DisposableDelegate> {
+  const content = panel.content;
+  const context = panel.context;
+  const sessionContext = context.sessionContext;
+  const rendermime = content.rendermime;
+  const widgetManagerFactory = () =>
+    new WidgetManager(context, rendermime, SETTINGS);
+
+  return registerWidgetHandler(
+    content,
+    sessionContext,
+    rendermime,
+    renderers,
+    widgetManagerFactory
+  );
+}
+
+export async function registerConsoleWidgetManager(
+  panel: ConsolePanel,
+  renderers: IterableIterator<WidgetRenderer>
+): Promise<DisposableDelegate> {
+  const content = panel.console;
+  const sessionContext = content.sessionContext;
+  const rendermime = content.rendermime;
+  const widgetManagerFactory = () =>
+    new KernelWidgetManager(sessionContext.session!.kernel!, rendermime);
+
+  return registerWidgetHandler(
+    content,
+    sessionContext,
+    rendermime,
+    renderers,
+    widgetManagerFactory
+  );
+}
+
 /**
  * The widget manager provider.
  */
@@ -154,6 +272,7 @@ export const managerPlugin: JupyterFrontEndPlugin<base.IJupyterWidgetRegistry> =
     requires: [IRenderMimeRegistry],
     optional: [
       INotebookTracker,
+      IConsoleTracker,
       ISettingRegistry,
       IMainMenu,
       ILoggerRegistry,
@@ -175,6 +294,7 @@ function activateWidgetExtension(
   app: JupyterFrontEnd,
   rendermime: IRenderMimeRegistry,
   tracker: INotebookTracker | null,
+  consoleTracker: IConsoleTracker | null,
   settingRegistry: ISettingRegistry | null,
   menu: IMainMenu | null,
   loggerRegistry: ILoggerRegistry | null,
@@ -183,15 +303,23 @@ function activateWidgetExtension(
   const { commands } = app;
   const trans = (translator ?? nullTranslator).load('jupyterlab_widgets');
 
-  const bindUnhandledIOPubMessageSignal = (nb: NotebookPanel): void => {
+  const bindUnhandledIOPubMessageSignal = async (
+    nb: NotebookPanel
+  ): Promise<void> => {
     if (!loggerRegistry) {
       return;
     }
+    const wManagerOwner = await getWidgetManagerOwner(
+      nb.context.sessionContext
+    );
+    const wManager = Private.widgetManagerProperty.get(wManagerOwner);
 
-    const wManager = Private.widgetManagerProperty.get(nb.context);
     if (wManager) {
       wManager.onUnhandledIOPubMessage.connect(
-        (sender: WidgetManager, msg: KernelMessage.IIOPubMessage) => {
+        (
+          sender: WidgetManager | KernelWidgetManager,
+          msg: KernelMessage.IIOPubMessage
+        ) => {
           const logger = loggerRegistry.getLogger(nb.context.path);
           let level: LogLevel = 'warning';
           if (
@@ -233,32 +361,32 @@ function activateWidgetExtension(
   );
 
   if (tracker !== null) {
-    tracker.forEach((panel) => {
-      registerWidgetManager(
-        panel.context,
-        panel.content.rendermime,
-        chain(
-          widgetRenderers(panel.content),
-          outputViews(app, panel.context.path)
-        )
+    const rendererIterator = (panel: NotebookPanel) =>
+      chain(
+        notebookWidgetRenderers(panel.content),
+        outputViews(app, panel.context.path)
       );
-
+    tracker.forEach(async (panel) => {
+      await registerWidgetManager(panel, rendererIterator(panel));
       bindUnhandledIOPubMessageSignal(panel);
     });
-    tracker.widgetAdded.connect((sender, panel) => {
-      registerWidgetManager(
-        panel.context,
-        panel.content.rendermime,
-        chain(
-          widgetRenderers(panel.content),
-          outputViews(app, panel.context.path)
-        )
-      );
-
+    tracker.widgetAdded.connect(async (sender, panel) => {
+      await registerWidgetManager(panel, rendererIterator(panel));
       bindUnhandledIOPubMessageSignal(panel);
     });
   }
 
+  if (consoleTracker !== null) {
+    const rendererIterator = (panel: ConsolePanel) =>
+      chain(consoleWidgetRenderers(panel.console));
+
+    consoleTracker.forEach(async (panel) => {
+      await registerConsoleWidgetManager(panel, rendererIterator(panel));
+    });
+    consoleTracker.widgetAdded.connect(async (sender, panel) => {
+      await registerConsoleWidgetManager(panel, rendererIterator(panel));
+    });
+  }
   if (settingRegistry !== null) {
     // Add a command for automatically saving (jupyter-)widget state.
     commands.addCommand('@jupyter-widgets/jupyterlab-manager:saveWidgetState', {
@@ -378,13 +506,23 @@ export default [
 ];
 namespace Private {
   /**
-   * A private attached property for a widget manager.
+   * A type alias for keys of `widgetManagerProperty` .
    */
-  export const widgetManagerProperty = new AttachedProperty<
-    DocumentRegistry.Context,
-    WidgetManager | undefined
-  >({
-    name: 'widgetManager',
-    create: (owner: DocumentRegistry.Context): undefined => undefined,
-  });
+  export type IWidgetManagerOwner = string;
+
+  /**
+   * A type alias for values of `widgetManagerProperty` .
+   */
+  export type IWidgetManagerValue =
+    | WidgetManager
+    | KernelWidgetManager
+    | undefined;
+
+  /**
+   * A private map for a widget manager.
+   */
+  export const widgetManagerProperty = new Map<
+    IWidgetManagerOwner,
+    IWidgetManagerValue
+  >();
 }