From bea8a1d10d528648408d592dabd9a45ae5528617 Mon Sep 17 00:00:00 2001
From: David Brochart <david.brochart@gmail.com>
Date: Wed, 1 Dec 2021 18:01:28 +0100
Subject: [PATCH] Save document through Yjs websocket

---
 .../services/contents/filemanager.py          | 23 ++++++++++++++++---
 1 file changed, 20 insertions(+), 3 deletions(-)

diff --git a/jupyter_server/services/contents/filemanager.py b/jupyter_server/services/contents/filemanager.py
index 54e2dca38d..c6f63b1733 100644
--- a/jupyter_server/services/contents/filemanager.py
+++ b/jupyter_server/services/contents/filemanager.py
@@ -2,6 +2,7 @@
 # Copyright (c) Jupyter Development Team.
 # Distributed under the terms of the Modified BSD License.
 import errno
+import json
 import mimetypes
 import os
 import shutil
@@ -40,6 +41,12 @@
     # windows + py2
     from jupyter_server.utils import samefile_simple as samefile
 
+try:
+    from jupyterlab.handlers.yjs_echo_ws import ROOMS
+    YJS_SUPPORT = True
+except ImportError:
+    YJS_SUPPORT = False
+
 _script_exporter = None
 
 
@@ -435,7 +442,7 @@ def save(self, model, path=""):
 
         if "type" not in model:
             raise web.HTTPError(400, u"No file type provided")
-        if "content" not in model and model["type"] != "directory":
+        if not YJS_SUPPORT and "content" not in model and model["type"] != "directory":
             raise web.HTTPError(400, u"No file content provided")
 
         os_path = self._get_os_path(path)
@@ -445,7 +452,12 @@ def save(self, model, path=""):
 
         try:
             if model["type"] == "notebook":
-                nb = nbformat.from_dict(model["content"])
+                if YJS_SUPPORT:
+                    nb_dict = json.loads(ROOMS[path].source)
+                    model["content"] = nb_dict
+                else:
+                    nb_dict = model["content"]
+                nb = nbformat.from_dict(nb_dict)
                 self.check_and_sign(nb, path)
                 self._save_notebook(os_path, nb)
                 # One checkpoint should always exist for notebooks.
@@ -453,7 +465,12 @@ def save(self, model, path=""):
                     self.create_checkpoint(path)
             elif model["type"] == "file":
                 # Missing format will be handled internally by _save_file.
-                self._save_file(os_path, model["content"], model.get("format"))
+                if YJS_SUPPORT:
+                    file_str = ROOMS[path].source
+                    model["content"] = file_str
+                else:
+                    file_str = model["content"]
+                self._save_file(os_path, file_str, model.get("format"))
             elif model["type"] == "directory":
                 self._save_directory(os_path, model, path)
             else: