Add undo_manager to Y documents

jupyter_ydoc/ybasedoc.py

@@ -2,9 +2,9 @@
 # Distributed under the terms of the Modified BSD License.
 
 from abc import ABC, abstractmethod
-from typing import Any, Callable, Dict, Optional
+from typing import Any, Callable, Optional
 
-from pycrdt import Doc, Map, Subscription
+from pycrdt import Doc, Map, Subscription, UndoManager
 
 
 class YBaseDoc(ABC):
@@ -15,6 +15,11 @@ class YBaseDoc(ABC):
     subscribe to changes in the document.
     """
 
+    _ydoc: Doc
+    _ystate: Map
+    _subscriptions: dict[Any, Subscription]
+    _undo_manager: UndoManager
+
     def __init__(self, ydoc: Optional[Doc] = None):
         """
         Constructs a YBaseDoc.
@@ -27,7 +32,8 @@ def __init__(self, ydoc: Optional[Doc] = None):
         else:
             self._ydoc = ydoc
         self._ystate = self._ydoc.get("state", type=Map)
-        self._subscriptions: Dict[Any, Subscription] = {}
+        self._subscriptions = {}
+        self._undo_manager = UndoManager(doc=self._ydoc, capture_timeout_millis=0)
 
     @property
     @abstractmethod
@@ -40,6 +46,15 @@ def version(self) -> str:
         """
 
     @property
+    def undo_manager(self) -> UndoManager:
+        """
+        A :class:`pycrdt.UndoManager` for the document.
+
+        :return: The document's undo manager.
+        :rtype: :class:`pycrdt.UndoManager`
+        """
+        return self._undo_manager
+
     def ystate(self) -> Map:
         """
         A :class:`pycrdt.Map` containing the state of the document.

jupyter_ydoc/yblob.py

@@ -33,6 +33,7 @@ def __init__(self, ydoc: Optional[Doc] = None):
         """
         super().__init__(ydoc)
         self._ysource = self._ydoc.get("source", type=Map)
+        self.undo_manager.expand_scope(self._ysource)
 
     @property
     def version(self) -> str:

jupyter_ydoc/ynotebook.py

@@ -56,6 +56,7 @@ def __init__(self, ydoc: Optional[Doc] = None):
         super().__init__(ydoc)
         self._ymeta = self._ydoc.get("meta", type=Map)
         self._ycells = self._ydoc.get("cells", type=Array)
+        self.undo_manager.expand_scope(self._ycells)
 
     @property
     def version(self) -> str:

jupyter_ydoc/yunicode.py

@@ -32,6 +32,7 @@ def __init__(self, ydoc: Optional[Doc] = None):
         """
         super().__init__(ydoc)
         self._ysource = self._ydoc.get("source", type=Text)
+        self.undo_manager.expand_scope(self._ysource)
 
     @property
     def version(self) -> str:

pyproject.toml

@@ -13,7 +13,7 @@ requires-python = ">=3.8"
 keywords = ["jupyter", "pycrdt", "yjs"]
 dependencies = [
     "importlib_metadata >=3.6; python_version<'3.10'",
-    "pycrdt >=0.8.16,<0.9.0",
+    "pycrdt >=0.9.0,<0.10.0",
 ]
 
 [[project.authors]]
@@ -31,7 +31,7 @@ test = [
     "pytest",
     "pytest-asyncio",
     "websockets >=10.0",
-    "pycrdt-websocket >=0.12.6,<0.13.0",
+    "pycrdt-websocket >=0.14.1,<0.15.0",
 ]
 docs = [
     "sphinx",

tests/test_ydocs.py

@@ -1,7 +1,7 @@
 # Copyright (c) Jupyter Development Team.
 # Distributed under the terms of the Modified BSD License.
 
-from jupyter_ydoc import YBlob
+from jupyter_ydoc import YBlob, YNotebook
 
 
 def test_yblob():
@@ -22,3 +22,34 @@ def callback(topic, event):
     assert topic == "source"
     assert event.keys["bytes"]["oldValue"] == b"012"
     assert event.keys["bytes"]["newValue"] == b"345"
+
+
+def test_ynotebook_undo_manager():
+    ynotebook = YNotebook()
+    cell0 = {
+        "cell_type": "code",
+        "source": "Hello",
+    }
+    ynotebook.append_cell(cell0)
+    source = ynotebook.ycells[0]["source"]
+    source += ", World!\n"
+    cell1 = {
+        "cell_type": "code",
+        "source": "print(1 + 1)\n",
+    }
+    ynotebook.append_cell(cell1)
+    assert len(ynotebook.ycells) == 2
+    assert str(ynotebook.ycells[0]["source"]) == "Hello, World!\n"
+    assert str(ynotebook.ycells[1]["source"]) == "print(1 + 1)\n"
+    assert ynotebook.undo_manager.can_undo()
+    ynotebook.undo_manager.undo()
+    assert len(ynotebook.ycells) == 1
+    assert str(ynotebook.ycells[0]["source"]) == "Hello, World!\n"
+    assert ynotebook.undo_manager.can_undo()
+    ynotebook.undo_manager.undo()
+    assert len(ynotebook.ycells) == 1
+    assert str(ynotebook.ycells[0]["source"]) == "Hello"
+    assert ynotebook.undo_manager.can_undo()
+    ynotebook.undo_manager.undo()
+    assert len(ynotebook.ycells) == 0
+    assert not ynotebook.undo_manager.can_undo()