Save document through Yjs websocket

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,13 @@
     # windows + py2
     from jupyter_server.utils import samefile_simple as samefile
 
+try:
+    from jupyterlab.handlers.yjs_echo_ws import ROOMS
+
+    YJS_SUPPORT = True
+except Exception:
+    YJS_SUPPORT = False
+
 _script_exporter = None
 
 
@@ -437,6 +445,13 @@ def save(self, model, path=""):
 
         if "type" not in model:
             raise web.HTTPError(400, u"No file type provided")
+
+        if YJS_SUPPORT and path in ROOMS:
+            if model["type"] == "notebook":
+                model["content"] = json.loads(ROOMS[path].source)
+            elif model["type"] == "file":
+                model["content"] = ROOMS[path].source
+
         if "content" not in model and model["type"] != "directory":
             raise web.HTTPError(400, u"No file content provided")
 
@@ -763,6 +778,15 @@ async def save(self, model, path=""):
 
         if "type" not in model:
             raise web.HTTPError(400, u"No file type provided")
+
+        if YJS_SUPPORT and path in ROOMS:
+            await ROOMS[path].ready_to_save.wait()
+            if model["type"] == "notebook":
+                model["content"] = json.loads(ROOMS[path].source)
+            elif model["type"] == "file":
+                model["content"] = ROOMS[path].source
+            ROOMS[path].ready_to_save.clear()
+
         if "content" not in model and model["type"] != "directory":
             raise web.HTTPError(400, u"No file content provided")
 

jupyter_server/services/contents/largefilemanager.py

@@ -1,5 +1,6 @@
 import base64
 import io
+import json
 import os
 
 from anyio.to_thread import run_sync
@@ -8,6 +9,13 @@
 from jupyter_server.services.contents.filemanager import AsyncFileContentsManager
 from jupyter_server.services.contents.filemanager import FileContentsManager
 
+try:
+    from jupyterlab.handlers.yjs_echo_ws import ROOMS
+
+    YJS_SUPPORT = True
+except Exception:
+    YJS_SUPPORT = False
+
 
 class LargeFileManager(FileContentsManager):
     """Handle large file upload."""
@@ -29,6 +37,13 @@ def save(self, model, path=""):
                         model["type"]
                     ),
                 )
+
+            if YJS_SUPPORT and path in ROOMS:
+                if model["type"] == "notebook":
+                    model["content"] = json.loads(ROOMS[path].source)
+                elif model["type"] == "file":
+                    model["content"] = ROOMS[path].source
+
             if "content" not in model and model["type"] != "directory":
                 raise web.HTTPError(400, u"No file content provided")
 
@@ -104,6 +119,15 @@ async def save(self, model, path=""):
                         model["type"]
                     ),
                 )
+
+            if YJS_SUPPORT and path in ROOMS:
+                await ROOMS[path].ready_to_save.wait()
+                if model["type"] == "notebook":
+                    model["content"] = json.loads(ROOMS[path].source)
+                elif model["type"] == "file":
+                    model["content"] = ROOMS[path].source
+                ROOMS[path].ready_to_save.clear()
+
             if "content" not in model and model["type"] != "directory":
                 raise web.HTTPError(400, u"No file content provided")