Add undo_manager to Y documents (#248) (#251)

This is a cherry-pick of #248.

.github/workflows/test.yml

@@ -2,9 +2,9 @@ name: Tests
 
 on:
   push:
-    branches: [main]
+    branches: [2.x]
   pull_request:
-    branches: [main]
+    branches: [2.x]
 
 jobs:
   pre-commit:

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
+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.
@@ -26,8 +31,9 @@ def __init__(self, ydoc: Optional[Doc] = None):
             self._ydoc = Doc()
         else:
             self._ydoc = ydoc
-        self._ydoc["state"] = self._ystate = Map()
-        self._subscriptions: Dict[Any, str] = {}
+        self._ystate = self._ydoc.get("state", type=Map)
+        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

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

jupyter_ydoc/ynotebook.py

@@ -54,8 +54,9 @@ def __init__(self, ydoc: Optional[Doc] = None):
         :type ydoc: :class:`pycrdt.Doc`, optional.
         """
         super().__init__(ydoc)
-        self._ydoc["meta"] = self._ymeta = Map()
-        self._ydoc["cells"] = self._ycells = Array()
+        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

@@ -31,7 +31,8 @@ def __init__(self, ydoc: Optional[Doc] = None):
         :type ydoc: :class:`pycrdt.Doc`, optional.
         """
         super().__init__(ydoc)
-        self._ydoc["source"] = self._ysource = Text()
+        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.7"
 keywords = ["jupyter", "ypy"]
 dependencies = [
     "importlib_metadata >=3.6; python_version<'3.10'",
-    "pycrdt >=0.8.1,<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

@@ -0,0 +1,35 @@
+# Copyright (c) Jupyter Development Team.
+# Distributed under the terms of the Modified BSD License.
+
+from jupyter_ydoc import YNotebook
+
+
+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()