From 38409998b5db79a6fb7bc78dba78325f2116b305 Mon Sep 17 00:00:00 2001
From: Jeremy Tuloup <jeremy.tuloup@gmail.com>
Date: Mon, 22 Aug 2022 16:40:37 +0200
Subject: [PATCH] Provide a custom OutputModel for now

---
 packages/voila/package.json   |  2 +
 packages/voila/src/manager.ts |  7 ++-
 packages/voila/src/output.ts  | 98 +++++++++++++++++++++++++++++++++++
 3 files changed, 106 insertions(+), 1 deletion(-)
 create mode 100644 packages/voila/src/output.ts

diff --git a/packages/voila/package.json b/packages/voila/package.json
index d1cb468f0..db1d7d804 100644
--- a/packages/voila/package.json
+++ b/packages/voila/package.json
@@ -10,6 +10,7 @@
   "dependencies": {
     "@jupyter-widgets/base": "^6.0.0",
     "@jupyter-widgets/controls": "^5.0.0",
+    "@jupyter-widgets/output": "^6.0.0",
     "@jupyter-widgets/jupyterlab-manager": "^5.0.2",
     "@jupyterlab/json-extension": "^3.0.0",
     "@jupyterlab/markdownviewer-extension": "^3.0.0",
@@ -21,6 +22,7 @@
     "@jupyterlab/docregistry": "^3.0.0",
     "@jupyterlab/logconsole": "^3.0.0",
     "@jupyterlab/mainmenu": "^3.0.0",
+    "@jupyterlab/nbformat": "^3.0.0",
     "@jupyterlab/notebook": "^3.0.0",
     "@jupyterlab/outputarea": "^3.0.0",
     "@jupyterlab/rendermime": "^3.0.0",
diff --git a/packages/voila/src/manager.ts b/packages/voila/src/manager.ts
index e427c63e4..364a67b59 100644
--- a/packages/voila/src/manager.ts
+++ b/packages/voila/src/manager.ts
@@ -25,6 +25,8 @@ import { Widget } from '@lumino/widgets';
 
 import { Kernel } from '@jupyterlab/services';
 
+import { OutputModel } from './output';
+
 const WIDGET_MIMETYPE = 'application/vnd.jupyter.widget-view+json';
 
 /**
@@ -112,7 +114,10 @@ export class WidgetManager extends KernelWidgetManager {
     this.register({
       name: '@jupyter-widgets/output',
       version: output.OUTPUT_WIDGET_VERSION,
-      exports: output as any
+      exports: {
+        ...(output as any),
+        OutputModel
+      }
     });
   }
 }
diff --git a/packages/voila/src/output.ts b/packages/voila/src/output.ts
new file mode 100644
index 000000000..58aeb94d6
--- /dev/null
+++ b/packages/voila/src/output.ts
@@ -0,0 +1,98 @@
+import { LabWidgetManager } from '@jupyter-widgets/jupyterlab-manager';
+
+import * as outputBase from '@jupyter-widgets/output';
+
+import * as nbformat from '@jupyterlab/nbformat';
+
+import { OutputAreaModel } from '@jupyterlab/outputarea';
+
+import { KernelMessage } from '@jupyterlab/services';
+
+/**
+ * Adapt the upstream output model to Voila using a KernelWidgetManager:
+ * https://github.com/jupyter-widgets/ipywidgets/blob/58adde2cfbe2f78a8ec6756d38a7c637f5e599f8/python/jupyterlab_widgets/src/output.ts#L22
+ *
+ * TODO: remove when https://github.com/jupyter-widgets/ipywidgets/pull/3561 (or similar) is merged and released in ipywidgets.
+ */
+export class OutputModel extends outputBase.OutputModel {
+  defaults(): Backbone.ObjectHash {
+    return { ...super.defaults(), msg_id: '', outputs: [] };
+  }
+
+  initialize(attributes: any, options: any): void {
+    super.initialize(attributes, options);
+    // The output area model is trusted since widgets are only rendered in trusted contexts.
+    this._outputs = new OutputAreaModel({ trusted: true });
+    this._msgHook = (msg): boolean => {
+      this.add(msg);
+      return false;
+    };
+
+    this.listenTo(this, 'change:msg_id', this.reset_msg_id);
+    this.listenTo(this, 'change:outputs', this.setOutputs);
+    this.setOutputs();
+  }
+
+  /**
+   * Reset the message id.
+   */
+  reset_msg_id(): void {
+    const kernel = this.widget_manager.kernel;
+    const msgId = this.get('msg_id');
+    const oldMsgId = this.previous('msg_id');
+
+    // Clear any old handler.
+    if (oldMsgId && kernel) {
+      kernel.removeMessageHook(oldMsgId, this._msgHook);
+    }
+
+    // Register any new handler.
+    if (msgId && kernel) {
+      kernel.registerMessageHook(msgId, this._msgHook);
+    }
+  }
+
+  add(msg: KernelMessage.IIOPubMessage): void {
+    const msgType = msg.header.msg_type;
+    switch (msgType) {
+      case 'execute_result':
+      case 'display_data':
+      case 'stream':
+      case 'error': {
+        const model = msg.content as nbformat.IOutput;
+        model.output_type = msgType as nbformat.OutputType;
+        this._outputs.add(model);
+        break;
+      }
+      case 'clear_output':
+        this.clear_output((msg as KernelMessage.IClearOutputMsg).content.wait);
+        break;
+      default:
+        break;
+    }
+    this.set('outputs', this._outputs.toJSON(), { newMessage: true });
+    this.save_changes();
+  }
+
+  clear_output(wait = false): void {
+    this._outputs.clear(wait);
+  }
+
+  get outputs(): OutputAreaModel {
+    return this._outputs;
+  }
+
+  setOutputs(model?: any, value?: any, options?: any): void {
+    if (!(options && options.newMessage)) {
+      // fromJSON does not clear the existing output
+      this.clear_output();
+      // fromJSON does not copy the message, so we make a deep copy
+      this._outputs.fromJSON(JSON.parse(JSON.stringify(this.get('outputs'))));
+    }
+  }
+
+  widget_manager!: LabWidgetManager;
+
+  private _msgHook!: (msg: KernelMessage.IIOPubMessage) => boolean;
+  private _outputs!: OutputAreaModel;
+}