Support Doc

python/pycrdt/_doc.py

@@ -1,17 +1,17 @@
 from __future__ import annotations
 
-from typing import Any, Callable, Iterable, Type, TypeVar, cast
+from typing import Any, Callable, Generic, Iterable, Type, TypeVar, cast
 
 from ._base import BaseDoc, BaseType, base_types, forbid_read_transaction
 from ._pycrdt import Doc as _Doc
 from ._pycrdt import SubdocsEvent, Subscription, TransactionEvent
 from ._pycrdt import Transaction as _Transaction
 from ._transaction import NewTransaction, ReadTransaction, Transaction
 
-T_BaseType = TypeVar("T_BaseType", bound=BaseType)
+T = TypeVar("T", bound=BaseType)
 
 
-class Doc(BaseDoc):
+class Doc(BaseDoc, Generic[T]):
     """
     A shared document.
 
@@ -23,7 +23,7 @@ class Doc(BaseDoc):
 
     def __init__(
         self,
-        init: dict[str, BaseType] = {},
+        init: dict[str, T] = {},
         *,
         client_id: int | None = None,
         doc: _Doc | None = None,
@@ -165,7 +165,7 @@ def apply_update(self, update: bytes) -> None:
             assert txn._txn is not None
             self._doc.apply_update(txn._txn, update)
 
-    def __setitem__(self, key: str, value: BaseType) -> None:
+    def __setitem__(self, key: str, value: T) -> None:
         """
         Sets a document root type:
         ```py
@@ -185,7 +185,7 @@ def __setitem__(self, key: str, value: BaseType) -> None:
         prelim = value._integrate(self, integrated)
         value._init(prelim)
 
-    def __getitem__(self, key: str) -> BaseType:
+    def __getitem__(self, key: str) -> T:
         """
         Gets the document root type corresponding to the given key:
         ```py
@@ -207,7 +207,7 @@ def __iter__(self) -> Iterable[str]:
         """
         return iter(self.keys())
 
-    def get(self, key: str, *, type: type[T_BaseType]) -> T_BaseType:
+    def get(self, key: str, *, type: type[T]) -> T:
         """
         Gets the document root type corresponding to the given key.
         If it already exists, it will be cast to the given type (if different),
@@ -230,29 +230,29 @@ def keys(self) -> Iterable[str]:
         """
         return self._roots.keys()
 
-    def values(self) -> Iterable[BaseType]:
+    def values(self) -> Iterable[T]:
         """
         Returns:
             An iterable over the document root types.
         """
         return self._roots.values()
 
-    def items(self) -> Iterable[tuple[str, BaseType]]:
+    def items(self) -> Iterable[tuple[str, T]]:
         """
         Returns:
             An iterable over the key-value pairs of document root types.
         """
         return self._roots.items()
 
     @property
-    def _roots(self) -> dict[str, BaseType]:
+    def _roots(self) -> dict[str, T]:
         with self.transaction() as txn:
             assert txn._txn is not None
             return {
                 key: (
                     None
                     if val is None
-                    else cast(Type[BaseType], base_types[type(val)])(_integrated=val, _doc=self)
+                    else cast(Type[T], base_types[type(val)])(_integrated=val, _doc=self)
                 )
                 for key, val in self._doc.roots(txn._txn).items()
             }

tests/test_types.py

@@ -1,20 +1,20 @@
 from typing import TypedDict
 
 import pytest
-from pycrdt import Array, Doc, Map
+from pycrdt import Array, Doc, Map, Text
 
 
 @pytest.mark.mypy_testing
 def mypy_test_array():
-    doc = Doc()
+    doc: Doc[Array] = Doc()
     array0: Array[int] = doc.get("array0", type=Array)
     array0.append(0)
     array0.append("foo")  # E: Argument 1 to "append" of "Array" has incompatible type "str"; expected "int"  [arg-type]
 
 
 @pytest.mark.mypy_testing
 def mypy_test_uniform_map():
-    doc = Doc()
+    doc: Doc[Map] = Doc()
     map0: Map[bool] = doc.get("map0", type=Map)
     map0["foo"] = True
     map0["foo"] = "bar"  # E: Incompatible types in assignment (expression has type "str", target has type "bool")
@@ -24,7 +24,7 @@ def mypy_test_uniform_map():
 
 @pytest.mark.mypy_testing
 def mypy_test_typed_map():
-    doc = Doc()
+    doc: Doc[Map] = Doc()
 
     MyMap = TypedDict(
         "MyMap",
@@ -46,3 +46,40 @@ def mypy_test_typed_map():
     v1: str = map0["toggle"]  # E: Incompatible types in assignment (expression has type "bool", variable has type "str")
     v2: bool = map0["toggle"]
     map0["key0"]  # E: TypedDict "MyMap@29" has no key "key0"
+
+
+@pytest.mark.mypy_testing
+def mypy_test_uniform_doc():
+    doc: Doc[Text] = Doc()
+    doc.get("text0", type=Text)
+    doc.get("array0", type=Array)  # E: Argument "type" to "get" of "Doc" has incompatible type "type[pycrdt._array.Array[Any]]"; expected "type[Text]"
+    doc.get("Map0", type=Map)  # E:  Argument "type" to "get" of "Doc" has incompatible type "type[pycrdt._map.Map[Any]]"; expected "type[Text]"
+
+
+@pytest.mark.mypy_testing
+def mypy_test_typed_doc():
+    MyMap = TypedDict(
+        "MyMap",
+        {
+            "name": str,
+            "toggle": bool,
+            "nested": Array[bool],
+        },
+    )
+
+    MyDoc = TypedDict(
+        "MyDoc",
+        {
+            "text0": Text,
+            "array0": Array[int],
+            "map0": MyMap,
+        }
+    )
+    doc: MyDoc = Doc()  # type: ignore[assignment]
+    map0: MyMap = Map()  # type: ignore[assignment]
+    doc["map0"] = map0
+    doc["map0"] = Array()  # E: Value of "map0" has incompatible type "Array[Never]"; expected "MyMap@61"
+    doc["text0"] = Text()
+    array0: Array[bool] = Array()
+    doc["array0"] = array0  # E: Value of "array0" has incompatible type "Array[bool]"; expected "Array[int]"
+    doc["array0"] = Array()