adds type hints to public APIs

observ/dep.py

@@ -7,7 +7,7 @@
 
 
 class Dep:
-    __slots__ = ["_subs", "__weakref__"]
+    __slots__ = ("_subs", "__weakref__")
     stack: List["Watcher"] = []  # noqa: F821
 
     def __init__(self) -> None:

observ/dict_proxy.py

@@ -51,7 +51,7 @@
 }
 
 
-class DictProxyBase(Proxy):
+class DictProxyBase(Proxy[dict]):
     def _orphaned_keydeps(self):
         return set(proxy_db.attrs(self)["keydep"].keys()) - set(self.target.keys())
 

observ/list_proxy.py

@@ -45,7 +45,7 @@
 }
 
 
-class ListProxyBase(Proxy):
+class ListProxyBase(Proxy[list]):
     pass
 
 

observ/proxy.py

@@ -1,9 +1,14 @@
+from __future__ import annotations
+
 from functools import partial
+from typing import cast, Generic, Literal, TypedDict, TypeVar
 
 from .proxy_db import proxy_db
 
+T = TypeVar("T")
+
 
-class Proxy:
+class Proxy(Generic[T]):
     """
     Proxy for an object/target.
 
@@ -16,9 +21,9 @@ class Proxy:
     """
 
     __hash__ = None
-    __slots__ = ["target", "readonly", "shallow", "__weakref__"]
+    __slots__ = ("target", "readonly", "shallow", "__weakref__")
 
-    def __init__(self, target, readonly=False, shallow=False):
+    def __init__(self, target: T, readonly=False, shallow=False):
         self.target = target
         self.readonly = readonly
         self.shallow = shallow
@@ -33,7 +38,7 @@ def __del__(self):
 TYPE_LOOKUP = {}
 
 
-def proxy(target, readonly=False, shallow=False):
+def proxy(target: T, readonly=False, shallow=False) -> T:
     """
     Returns a Proxy for the given object. If a proxy for the given
     configuration already exists, it will return that instead of
@@ -67,14 +72,26 @@ def proxy(target, readonly=False, shallow=False):
             return proxy_type(target, readonly=readonly, shallow=shallow)
 
     if isinstance(target, tuple):
-        return tuple(proxy(x, readonly=readonly, shallow=shallow) for x in target)
+        return cast(
+            T, tuple(proxy(x, readonly=readonly, shallow=shallow) for x in target)
+        )
 
     # We can't proxy a plain value
-    return target
+    return cast(T, target)
+
+
+try:
+    # for Python >= 3.11
+    class Ref(TypedDict, Generic[T]):
+        value: T
 
+    def ref(target: T) -> Ref[T]:
+        return proxy(Ref(value=target))
 
-def ref(target):
-    return proxy({"value": target})
+except TypeError:
+    # before python 3.11 a TypedDict cannot inherit from a non-TypedDict class
+    def ref(target: T) -> dict[Literal["value"], T]:
+        return proxy({"value": target})
 
 
 reactive = proxy
@@ -83,7 +100,7 @@ def ref(target):
 shallow_readonly = partial(proxy, shallow=True, readonly=True)
 
 
-def to_raw(target):
+def to_raw(target: Proxy[T] | T) -> T:
     """
     Returns a raw object from which any trace of proxy has been replaced
     with its wrapped target value.
@@ -92,15 +109,15 @@ def to_raw(target):
         return to_raw(target.target)
 
     if isinstance(target, list):
-        return [to_raw(t) for t in target]
+        return cast(T, [to_raw(t) for t in target])
 
     if isinstance(target, dict):
-        return {key: to_raw(value) for key, value in target.items()}
+        return cast(T, {key: to_raw(value) for key, value in target.items()})
 
     if isinstance(target, tuple):
-        return tuple(to_raw(t) for t in target)
+        return cast(T, tuple(to_raw(t) for t in target))
 
     if isinstance(target, set):
-        return {to_raw(t) for t in target}
+        return cast(T, {to_raw(t) for t in target})
 
     return target

observ/proxy_db.py

@@ -15,7 +15,7 @@ class ProxyDb:
     removed from the collection.
     """
 
-    __slots__ = ["db"]
+    __slots__ = ("db",)
 
     def __init__(self):
         self.db = {}

observ/scheduler.py

@@ -9,7 +9,7 @@
 
 
 class Scheduler:
-    __slots__ = [
+    __slots__ = (
         "_queue",
         "_queue_indices",
         "flushing",
@@ -19,7 +19,7 @@ class Scheduler:
         "waiting",
         "request_flush",
         "detect_cycles",
-    ]
+    )
 
     def __init__(self):
         self._queue = []

observ/set_proxy.py

@@ -54,7 +54,7 @@
 }
 
 
-class SetProxyBase(Proxy):
+class SetProxyBase(Proxy[set]):
     pass
 
 

observ/store.py

@@ -1,5 +1,5 @@
 from functools import partial, wraps
-from typing import Callable, Collection, TypeVar
+from typing import Callable, Generic, TypeVar
 
 import patchdiff
 
@@ -51,12 +51,15 @@ def decorator_computed(fn: T) -> T:
     return decorator_computed(_fn)
 
 
-class Store:
+S = TypeVar("S")
+
+
+class Store(Generic[S]):
     """
     Store that tracks mutations to state in order to enable undo/redo functionality
     """
 
-    def __init__(self, state: Collection, strict=True):
+    def __init__(self, state: S, strict=True):
         """
         Creates a store with the given state as the initial state.
         When `strict` is False, calling mutations that do not result

observ/watcher.py

@@ -9,27 +9,28 @@
 from functools import partial, wraps
 import inspect
 from itertools import count
-from typing import Any, Callable, Optional, TypeVar
+from typing import Any, Callable, Generic, TypeVar, Union
 from weakref import ref, WeakSet
 
 from .dep import Dep
 from .dict_proxy import DictProxyBase
 from .list_proxy import ListProxyBase
-from .proxy import Proxy
 from .scheduler import scheduler
 from .set_proxy import SetProxyBase
 
 
-T = TypeVar("T", bound=Callable[[], Any])
+T = TypeVar("T")
+Watchable = Union[Callable[[], T], T]
+WatchCallback = Union[Callable[[], Any], Callable[[T], Any], Callable[[T, T], Any]]
 
 
 def watch(
-    fn: Callable[[], Any] | Proxy | list[Proxy],
-    callback: Optional[Callable] = None,
+    fn: Watchable[T],
+    callback: WatchCallback[T] | None = None,
     sync: bool = False,
     deep: bool | None = None,
     immediate: bool = False,
-):
+) -> Watcher[T]:
     watcher = Watcher(fn, sync=sync, lazy=False, deep=deep, callback=callback)
     if immediate:
         watcher.dirty = True
@@ -42,8 +43,8 @@ def watch(
 watch_effect = partial(watch, immediate=False, deep=True, callback=None)
 
 
-def computed(_fn=None, *, deep=True):
-    def decorator_computed(fn: T) -> T:
+def computed(_fn: Callable[[], T] | None = None, *, deep=True) -> Callable[[], T]:
+    def decorator_computed(fn: Callable[[], T]) -> Callable[[], T]:
         """
         Create a watcher for an expression.
         Note: make sure fn doesn't need any arguments to run
@@ -113,8 +114,8 @@ class WrongNumberOfArgumentsError(TypeError):
     pass
 
 
-class Watcher:
-    __slots__ = [
+class Watcher(Generic[T]):
+    __slots__ = (
         "id",
         "fn",
         "_deps",
@@ -128,15 +129,15 @@ class Watcher:
         "value",
         "_number_of_callback_args",
         "__weakref__",
-    ]
+    )
 
     def __init__(
         self,
-        fn: Callable[[], Any] | Proxy | list[Proxy],
+        fn: Watchable[T],
         sync: bool = False,
         lazy: bool = True,
         deep: bool | None = None,
-        callback: Callable = None,
+        callback: WatchCallback[T] | None = None,
     ) -> None:
         """
         sync: Ignore the scheduler

tests/test_collections.py

@@ -55,6 +55,10 @@
     "readonly",
     "shallow",
     "__weakref__",
+    # exclude attributes added by typing.Generic
+    "_is_protocol",
+    "__orig_bases__",
+    "__parameters__",
 }