Use kernel id as keys of the map used to store widget manager (#2)

Author: trungleduc

python/jupyterlab_widgets/src/plugin.ts

@@ -5,12 +5,13 @@ import { ISettingRegistry } from '@jupyterlab/settingregistry';
 
 import * as nbformat from '@jupyterlab/nbformat';
 
-import { IConsoleTracker, CodeConsole } from '@jupyterlab/console';
-
-import { DocumentRegistry } from '@jupyterlab/docregistry';
+import {
+  IConsoleTracker,
+  CodeConsole,
+  ConsolePanel,
+} from '@jupyterlab/console';
 
 import {
-  INotebookModel,
   INotebookTracker,
   Notebook,
   NotebookPanel,
@@ -33,8 +34,6 @@ import { toArray, filter } from '@lumino/algorithm';
 
 import { DisposableDelegate } from '@lumino/disposable';
 
-import { AttachedProperty } from '@lumino/properties';
-
 import { WidgetRenderer } from './renderer';
 
 import {
@@ -55,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[] = [];
 
@@ -131,16 +131,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) {
@@ -167,43 +220,43 @@ export function registerWidgetManager(
   });
 }
 
-export function registerConsoleWidgetManager(
-  console: CodeConsole,
-  rendermime: IRenderMimeRegistry,
+export async function registerWidgetManager(
+  panel: NotebookPanel,
   renderers: IterableIterator<WidgetRenderer>
-): DisposableDelegate {
-  let wManager = Private.widgetManagerProperty.get(console);
-  if (!wManager) {
-    wManager = new KernelWidgetManager(
-      console.sessionContext.session?.kernel!,
-      rendermime
-    );
-    WIDGET_REGISTRY.forEach((data) => wManager!.register(data));
-    Private.widgetManagerProperty.set(console, wManager);
-  }
-
-  for (const r of renderers) {
-    r.manager = wManager;
-  }
-
-  // Replace the placeholder widget renderer with one bound to this widget
-  // manager.
-  rendermime.removeMimeType(WIDGET_VIEW_MIMETYPE);
-  rendermime.addFactory(
-    {
-      safe: false,
-      mimeTypes: [WIDGET_VIEW_MIMETYPE],
-      createRenderer: (options) => new WidgetRenderer(options, wManager),
-    },
-    0
+): 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
   );
+}
 
-  return new DisposableDelegate(() => {
-    if (rendermime) {
-      rendermime.removeMimeType(WIDGET_VIEW_MIMETYPE);
-    }
-    wManager!.dispose();
-  });
+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
+  );
 }
 
 /**
@@ -247,12 +300,17 @@ 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(
         (
@@ -300,48 +358,30 @@ function activateWidgetExtension(
   );
 
   if (tracker !== null) {
-    tracker.forEach((panel) => {
-      registerWidgetManager(
-        panel.context,
-        panel.content.rendermime,
-        chain(
-          notebookWidgetRenderers(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(
-          notebookWidgetRenderers(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 panel.sessionContext.ready;
-      registerConsoleWidgetManager(
-        panel.console,
-        panel.console.rendermime,
-        chain(consoleWidgetRenderers(panel.console))
-      );
+      await registerConsoleWidgetManager(panel, rendererIterator(panel));
     });
     consoleTracker.widgetAdded.connect(async (sender, panel) => {
-      await panel.sessionContext.ready;
-      registerConsoleWidgetManager(
-        panel.console,
-        panel.console.rendermime,
-        chain(consoleWidgetRenderers(panel.console))
-      );
+      await registerConsoleWidgetManager(panel, rendererIterator(panel));
     });
   }
   if (settingRegistry !== null) {
@@ -415,14 +455,23 @@ function activateWidgetExtension(
 
 namespace Private {
   /**
-   * A private attached property for a widget manager.
+   * A type alias for keys of `widgetManagerProperty` .
    */
-  export const widgetManagerProperty = new AttachedProperty<
-    DocumentRegistry.Context | CodeConsole,
-    WidgetManager | KernelWidgetManager | undefined
-  >({
-    name: 'widgetManager',
-    create: (owner: DocumentRegistry.Context | CodeConsole): 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
+  >();
 }