Handle KernelWidgetManager in the JupyterLab OutputModel

Author: jtpio

python/jupyterlab_widgets/src/output.ts

@@ -7,7 +7,7 @@ import { JupyterLuminoPanelWidget } from '@jupyter-widgets/base';
 
 import { Panel } from '@lumino/widgets';
 
-import { WidgetManager } from './manager';
+import { KernelWidgetManager, WidgetManager } from './manager';
 
 import { OutputAreaModel, OutputArea } from '@jupyterlab/outputarea';
 
@@ -33,11 +33,14 @@ export class OutputModel extends outputBase.OutputModel {
       return false;
     };
 
-    this.widget_manager.context.sessionContext.kernelChanged.connect(
-      (sender, args) => {
-        this._handleKernelChanged(args);
-      }
-    );
+    // if the context is available, react on kernel changes
+    if (this.widget_manager instanceof WidgetManager) {
+      this.widget_manager.context.sessionContext.kernelChanged.connect(
+        (sender, args) => {
+          this._handleKernelChanged(args);
+        }
+      );
+    }
     this.listenTo(this, 'change:msg_id', this.reset_msg_id);
     this.listenTo(this, 'change:outputs', this.setOutputs);
     this.setOutputs();
@@ -60,7 +63,12 @@ export class OutputModel extends outputBase.OutputModel {
    * Reset the message id.
    */
   reset_msg_id(): void {
-    const kernel = this.widget_manager.context.sessionContext?.session?.kernel;
+    let kernel;
+    if (this.widget_manager instanceof WidgetManager) {
+      kernel = this.widget_manager.context.sessionContext?.session?.kernel;
+    } else if (this.widget_manager instanceof KernelWidgetManager) {
+      kernel = this.widget_manager.kernel;
+    }
     const msgId = this.get('msg_id');
     const oldMsgId = this.previous('msg_id');
 
@@ -114,7 +122,7 @@ export class OutputModel extends outputBase.OutputModel {
     }
   }
 
-  widget_manager: WidgetManager;
+  widget_manager: WidgetManager | KernelWidgetManager;
 
   private _msgHook: (msg: KernelMessage.IIOPubMessage) => boolean;
   private _outputs: OutputAreaModel;