From 74e0922a4e7a5444428c5099e802db19302d448b Mon Sep 17 00:00:00 2001
From: David Brochart <david.brochart@gmail.com>
Date: Mon, 27 Nov 2023 19:39:53 +0100
Subject: [PATCH] Change notebook code cell stream output schema

---
 javascript/src/ycell.ts   | 27 +++++++++++++++++++++++++--
 jupyter_ydoc/ynotebook.py |  7 ++++++-
 2 files changed, 31 insertions(+), 3 deletions(-)

diff --git a/javascript/src/ycell.ts b/javascript/src/ycell.ts
index ccafbf8..f0c0664 100644
--- a/javascript/src/ycell.ts
+++ b/javascript/src/ycell.ts
@@ -761,13 +761,35 @@ export class YCodeCell
     return JSONExt.deepCopy(this._youtputs.toArray());
   }
 
+  createOutputs(outputs: Array<nbformat.IOutput>): Array<any> {
+    const new_outputs: Array<any> = [];
+    for (const output of outputs) {
+      if (output.output_type === 'stream') {
+        const { text, ...output_without_text } = output;
+        const _new_output: { [id: string]: any } = output_without_text;
+        const new_text = new Y.Array();
+        new_text.push(text as string[]);
+        _new_output['text'] = new_text;
+        const new_output = new Y.Map();
+        for (const [key, value] of Object.entries(_new_output)) {
+          new_output.set(key, value);
+        }
+        new_outputs.push(new_output);
+      } else {
+        new_outputs.push(output);
+      }
+    }
+    return new_outputs;
+  }
+
   /**
    * Replace all outputs.
    */
   setOutputs(outputs: Array<nbformat.IOutput>): void {
     this.transact(() => {
       this._youtputs.delete(0, this._youtputs.length);
-      this._youtputs.insert(0, outputs);
+      const new_outputs = this.createOutputs(outputs);
+      this._youtputs.insert(0, new_outputs);
     }, false);
   }
 
@@ -789,7 +811,8 @@ export class YCodeCell
       end < this._youtputs.length ? end - start : this._youtputs.length - start;
     this.transact(() => {
       this._youtputs.delete(start, fin);
-      this._youtputs.insert(start, outputs);
+      const new_outputs = this.createOutputs(outputs);
+      this._youtputs.insert(start, new_outputs);
     }, false);
   }
 
diff --git a/jupyter_ydoc/ynotebook.py b/jupyter_ydoc/ynotebook.py
index 47ae466..5418a66 100644
--- a/jupyter_ydoc/ynotebook.py
+++ b/jupyter_ydoc/ynotebook.py
@@ -170,7 +170,12 @@ def create_ycell(self, value: Dict[str, Any]) -> Y.YMap:
             if "attachments" in cell and not cell["attachments"]:
                 del cell["attachments"]
         elif cell_type == "code":
-            cell["outputs"] = Y.YArray(cell.get("outputs", []))
+            outputs = cell.get("outputs", [])
+            for idx, output in enumerate(outputs):
+                if output.get("output_type") == "stream":
+                    output["text"] = Y.YArray(output.get("text", []))
+                    outputs[idx] = Y.YMap(output)
+            cell["outputs"] = Y.YArray(outputs)
 
         return Y.YMap(cell)