Use pycrdt's typed containers

.github/workflows/test.yml

@@ -49,6 +49,7 @@ jobs:
         run: |
           micromamba install pip nodejs=18
           pip install ".[test]"
+          pip install "pycrdt >=0.12.1"
       - name: Build JavaScript assets
         working-directory: javascript
         run: |

jupyter_ydoc/__init__.py

@@ -5,9 +5,13 @@
 
 from ._version import __version__ as __version__
 from .yblob import YBlob as YBlob
+from .yblob import YBlobDoc as YBlobDoc
 from .yfile import YFile as YFile
+from .yfile import YFileDoc as YFileDoc
 from .ynotebook import YNotebook as YNotebook
+from .ynotebook import YNotebookDoc as YNotebookDoc
 from .yunicode import YUnicode as YUnicode
+from .yunicode import YUnicodeDoc as YUnicodeDoc
 
 # See compatibility note on `group` keyword in
 # https://docs.python.org/3/library/importlib.metadata.html#entry-points

jupyter_ydoc/ybasedoc.py

@@ -1,10 +1,22 @@
 # Copyright (c) Jupyter Development Team.
 # Distributed under the terms of the Modified BSD License.
 
+from __future__ import annotations
+
 from abc import ABC, abstractmethod
-from typing import Any, Callable, Dict, Optional
+from typing import Any, Callable
+
+from pycrdt import Awareness, Doc, Subscription, TypedDoc, TypedMap, UndoManager
+
+
+class YState(TypedMap):
+    dirty: bool
+    hash: str
+    path: str
 
-from pycrdt import Awareness, Doc, Map, Subscription, UndoManager
+
+class YDoc(TypedDoc):
+    state: YState
 
 
 class YBaseDoc(ABC):
@@ -15,12 +27,12 @@ class YBaseDoc(ABC):
     subscribe to changes in the document.
     """
 
-    _ydoc: Doc
-    _ystate: Map
-    _subscriptions: Dict[Any, Subscription]
+    _ydoc: YDoc
+    _ystate: YState
+    _subscriptions: dict[Any, Subscription]
     _undo_manager: UndoManager
 
-    def __init__(self, ydoc: Optional[Doc] = None, awareness: Optional[Awareness] = None):
+    def __init__(self, ydoc: YDoc | Doc | None = None, awareness: Awareness | None = None):
         """
         Constructs a YBaseDoc.
 
@@ -30,15 +42,15 @@ def __init__(self, ydoc: Optional[Doc] = None, awareness: Optional[Awareness] =
                           between clients.
         :type awareness: :class:`pycrdt.Awareness`, optional.
         """
-        if ydoc is None:
-            self._ydoc = Doc()
-        else:
+        if isinstance(ydoc, YDoc):
             self._ydoc = ydoc
+        else:
+            self._ydoc = YDoc(ydoc)
         self.awareness = awareness
 
-        self._ystate = self._ydoc.get("state", type=Map)
+        self._ydoc.state = self._ystate = YState()
         self._subscriptions = {}
-        self._undo_manager = UndoManager(doc=self._ydoc, capture_timeout_millis=0)
+        self._undo_manager = UndoManager(doc=self._ydoc._, capture_timeout_millis=0)
 
     @property
     @abstractmethod
@@ -60,22 +72,22 @@ def undo_manager(self) -> UndoManager:
         """
         return self._undo_manager
 
-    def ystate(self) -> Map:
+    def ystate(self) -> YState:
         """
-        A :class:`pycrdt.Map` containing the state of the document.
+        A :class:`YState` containing the state of the document.
 
         :return: The document's state.
-        :rtype: :class:`pycrdt.Map`
+        :rtype: :class:`YState`
         """
         return self._ystate
 
     @property
-    def ydoc(self) -> Doc:
+    def ydoc(self) -> YDoc:
         """
-        The underlying :class:`pycrdt.Doc` that contains the data.
+        The :class:`YDoc` that contains the data.
 
         :return: The document's ydoc.
-        :rtype: :class:`pycrdt.Doc`
+        :rtype: :class:`YDoc`
         """
         return self._ydoc
 
@@ -100,14 +112,17 @@ def source(self, value: Any):
         return self.set(value)
 
     @property
-    def dirty(self) -> Optional[bool]:
+    def dirty(self) -> bool | None:
         """
         Returns whether the document is dirty.
 
         :return: Whether the document is dirty.
-        :rtype: Optional[bool]
+        :rtype: bool | None
         """
-        return self._ystate.get("dirty")
+        try:
+            return self._ystate.dirty
+        except KeyError:
+            return None
 
     @dirty.setter
     def dirty(self, value: bool) -> None:
@@ -117,17 +132,20 @@ def dirty(self, value: bool) -> None:
         :param value: Whether the document is clean or dirty.
         :type value: bool
         """
-        self._ystate["dirty"] = value
+        self._ystate.dirty = value
 
     @property
-    def hash(self) -> Optional[str]:
+    def hash(self) -> str | None:
         """
         Returns the document hash as computed by contents manager.
 
         :return: The document hash.
-        :rtype: Optional[str]
+        :rtype: str | None
         """
-        return self._ystate.get("hash")
+        try:
+            return self._ystate.hash
+        except KeyError:
+            return None
 
     @hash.setter
     def hash(self, value: str) -> None:
@@ -137,17 +155,20 @@ def hash(self, value: str) -> None:
         :param value: The document hash.
         :type value: str
         """
-        self._ystate["hash"] = value
+        self._ystate.hash = value
 
     @property
-    def path(self) -> Optional[str]:
+    def path(self) -> str | None:
         """
         Returns document's path.
 
         :return: Document's path.
-        :rtype: Optional[str]
+        :rtype: str | None
         """
-        return self._ystate.get("path")
+        try:
+            return self._ystate.path
+        except KeyError:
+            return None
 
     @path.setter
     def path(self, value: str) -> None:
@@ -157,7 +178,7 @@ def path(self, value: str) -> None:
         :param value: Document's path.
         :type value: str
         """
-        self._ystate["path"] = value
+        self._ystate.path = value
 
     @abstractmethod
     def get(self) -> Any:

jupyter_ydoc/yblob.py

@@ -1,12 +1,22 @@
 # Copyright (c) Jupyter Development Team.
 # Distributed under the terms of the Modified BSD License.
 
+from __future__ import annotations
+
 from functools import partial
-from typing import Any, Callable, Optional
+from typing import Any, Callable
+
+from pycrdt import Awareness, Doc, TypedMap
+
+from .ybasedoc import YBaseDoc, YDoc
 
-from pycrdt import Awareness, Doc, Map
 
-from .ybasedoc import YBaseDoc
+class YBlobSource(TypedMap):
+    bytes: bytes
+
+
+class YBlobDoc(YDoc):
+    source: YBlobSource
 
 
 class YBlob(YBaseDoc):
@@ -24,7 +34,10 @@ class YBlob(YBaseDoc):
         }
     """
 
-    def __init__(self, ydoc: Optional[Doc] = None, awareness: Optional[Awareness] = None):
+    _ydoc: YBlobDoc
+    _ysource: YBlobSource
+
+    def __init__(self, ydoc: Doc | None = None, awareness: Awareness | None = None):
         """
         Constructs a YBlob.
 
@@ -34,9 +47,9 @@ def __init__(self, ydoc: Optional[Doc] = None, awareness: Optional[Awareness] =
                           between clients.
         :type awareness: :class:`pycrdt.Awareness`, optional.
         """
-        super().__init__(ydoc, awareness)
-        self._ysource = self._ydoc.get("source", type=Map)
-        self.undo_manager.expand_scope(self._ysource)
+        super().__init__(YBlobDoc(ydoc), awareness)
+        self._ydoc.source = self._ysource = YBlobSource()
+        self.undo_manager.expand_scope(self._ysource._)
 
     @property
     def version(self) -> str:
@@ -55,7 +68,10 @@ def get(self) -> bytes:
         :return: Document's content.
         :rtype: bytes
         """
-        return self._ysource.get("bytes", b"")
+        try:
+            return self._ysource.bytes
+        except KeyError:
+            return b""
 
     def set(self, value: bytes) -> None:
         """
@@ -64,7 +80,7 @@ def set(self, value: bytes) -> None:
         :param value: The content of the document.
         :type value: bytes
         """
-        self._ysource["bytes"] = value
+        self._ysource.bytes = value
 
     def observe(self, callback: Callable[[str, Any], None]) -> None:
         """
@@ -74,5 +90,5 @@ def observe(self, callback: Callable[[str, Any], None]) -> None:
         :type callback: Callable[[str, Any], None]
         """
         self.unobserve()
-        self._subscriptions[self._ystate] = self._ystate.observe(partial(callback, "state"))
-        self._subscriptions[self._ysource] = self._ysource.observe(partial(callback, "source"))
+        self._subscriptions[self._ystate] = self._ystate._.observe(partial(callback, "state"))
+        self._subscriptions[self._ysource] = self._ysource._.observe(partial(callback, "source"))

jupyter_ydoc/yfile.py

@@ -1,8 +1,14 @@
 # Copyright (c) Jupyter Development Team.
 # Distributed under the terms of the Modified BSD License.
 
-from .yunicode import YUnicode
+from .yunicode import YUnicode, YUnicodeDoc
 
+# For backwards-compatibility:
 
-class YFile(YUnicode):  # for backwards-compatibility
+
+class YFile(YUnicode):
+    pass
+
+
+class YFileDoc(YUnicodeDoc):
     pass