From 21b0a034d6917a8c2793c1ee2de39d22299a87cb Mon Sep 17 00:00:00 2001
From: Lendemor <thomas.brandeho@gmail.com>
Date: Wed, 20 Nov 2024 20:29:43 +0100
Subject: [PATCH] state typing

---
 reflex/state.py | 49 +++++++++++++++++++++++++++----------------------
 1 file changed, 27 insertions(+), 22 deletions(-)

diff --git a/reflex/state.py b/reflex/state.py
index 9ff6f0ea8..44c79d4a3 100644
--- a/reflex/state.py
+++ b/reflex/state.py
@@ -30,6 +30,7 @@ from typing import (
     Optional,
     Sequence,
     Set,
+    SupportsIndex,
     Tuple,
     Type,
     TypeVar,
@@ -1026,7 +1027,7 @@ class BaseState(Base, ABC, extra=pydantic.Extra.allow):
         setattr(cls, prop._var_field_name, prop)
 
     @classmethod
-    def _create_event_handler(cls, fn):
+    def _create_event_handler(cls, fn: Any):
         """Create an event handler for the given function.
 
         Args:
@@ -1143,14 +1144,14 @@ class BaseState(Base, ABC, extra=pydantic.Extra.allow):
 
         cls._check_overwritten_dynamic_args(list(args.keys()))
 
-        def argsingle_factory(param):
-            def inner_func(self) -> str:
+        def argsingle_factory(param: str):
+            def inner_func(self: BaseState) -> str:  # type: ignore
                 return self.router.page.params.get(param, "")
 
             return inner_func
 
-        def arglist_factory(param):
-            def inner_func(self) -> List[str]:
+        def arglist_factory(param: str):
+            def inner_func(self: BaseState) -> List[str]:
                 return self.router.page.params.get(param, [])
 
             return inner_func
@@ -2467,7 +2468,9 @@ class StateProxy(wrapt.ObjectProxy):
     """
 
     def __init__(
-        self, state_instance, parent_state_proxy: Optional["StateProxy"] = None
+        self,
+        state_instance: BaseState,
+        parent_state_proxy: Optional["StateProxy"] = None,
     ):
         """Create a proxy for a state instance.
 
@@ -3565,10 +3568,10 @@ class MutableProxy(wrapt.ObjectProxy):
 
     def _mark_dirty(
         self,
-        wrapped=None,
-        instance=None,
-        args=tuple(),
-        kwargs=None,
+        wrapped: Callable | None = None,
+        instance: BaseState | None = None,
+        args: tuple = tuple(),
+        kwargs: dict | None = None,
     ) -> Any:
         """Mark the state as dirty, then call a wrapped function.
 
@@ -3608,7 +3611,9 @@ class MutableProxy(wrapt.ObjectProxy):
             )
         return value
 
-    def _wrap_recursive_decorator(self, wrapped, instance, args, kwargs) -> Any:
+    def _wrap_recursive_decorator(
+        self, wrapped: Callable, instance: BaseState, args: list, kwargs: dict
+    ) -> Any:
         """Wrap a function that returns a possibly mutable value.
 
         Intended for use with `FunctionWrapper` from the `wrapt` library.
@@ -3667,7 +3672,7 @@ class MutableProxy(wrapt.ObjectProxy):
 
         return value
 
-    def __getitem__(self, key) -> Any:
+    def __getitem__(self, key: Any) -> Any:
         """Get the item on the proxied object and return a proxy if mutable.
 
         Args:
@@ -3690,7 +3695,7 @@ class MutableProxy(wrapt.ObjectProxy):
             # Recursively wrap mutable items retrieved through this proxy.
             yield self._wrap_recursive(value)
 
-    def __delattr__(self, name):
+    def __delattr__(self, name: str):
         """Delete the attribute on the proxied object and mark state dirty.
 
         Args:
@@ -3698,7 +3703,7 @@ class MutableProxy(wrapt.ObjectProxy):
         """
         self._mark_dirty(super().__delattr__, args=(name,))
 
-    def __delitem__(self, key):
+    def __delitem__(self, key: str):
         """Delete the item on the proxied object and mark state dirty.
 
         Args:
@@ -3706,7 +3711,7 @@ class MutableProxy(wrapt.ObjectProxy):
         """
         self._mark_dirty(super().__delitem__, args=(key,))
 
-    def __setitem__(self, key, value):
+    def __setitem__(self, key: str, value: Any):
         """Set the item on the proxied object and mark state dirty.
 
         Args:
@@ -3715,7 +3720,7 @@ class MutableProxy(wrapt.ObjectProxy):
         """
         self._mark_dirty(super().__setitem__, args=(key, value))
 
-    def __setattr__(self, name, value):
+    def __setattr__(self, name: str, value: Any):
         """Set the attribute on the proxied object and mark state dirty.
 
         If the attribute starts with "_self_", then the state is NOT marked
@@ -3739,7 +3744,7 @@ class MutableProxy(wrapt.ObjectProxy):
         """
         return copy.copy(self.__wrapped__)
 
-    def __deepcopy__(self, memo=None) -> Any:
+    def __deepcopy__(self, memo: dict[int, Any] | None = None) -> Any:
         """Return a deepcopy of the proxy.
 
         Args:
@@ -3750,7 +3755,7 @@ class MutableProxy(wrapt.ObjectProxy):
         """
         return copy.deepcopy(self.__wrapped__, memo=memo)
 
-    def __reduce_ex__(self, protocol_version):
+    def __reduce_ex__(self, protocol_version: SupportsIndex):
         """Get the state for redis serialization.
 
         This method is called by cloudpickle to serialize the object.
@@ -3811,10 +3816,10 @@ class ImmutableMutableProxy(MutableProxy):
 
     def _mark_dirty(
         self,
-        wrapped=None,
-        instance=None,
-        args=tuple(),
-        kwargs=None,
+        wrapped: Callable | None = None,
+        instance: BaseState | None = None,
+        args: tuple = tuple(),
+        kwargs: dict | None = None,
     ) -> Any:
         """Raise an exception when an attempt is made to modify the object.