From c8df48728b5a51dfb6a7c39397341508899cd2ca Mon Sep 17 00:00:00 2001
From: Khaleel Al-Adhami <khaleel.aladhami@gmail.com>
Date: Tue, 22 Oct 2024 11:49:23 -0700
Subject: [PATCH] add additional typing for calling events

---
 reflex/event.py | 88 ++++++++++++++++++++++++++-----------------------
 1 file changed, 47 insertions(+), 41 deletions(-)

diff --git a/reflex/event.py b/reflex/event.py
index cfe40ef9a..fdab9650a 100644
--- a/reflex/event.py
+++ b/reflex/event.py
@@ -16,6 +16,7 @@ from typing import (
     Generic,
     List,
     Optional,
+    Self,
     Tuple,
     Type,
     TypeVar,
@@ -1445,13 +1446,8 @@ class LiteralEventChainVar(ArgsFunctionOperation, LiteralVar, EventChainVar):
         )
 
 
-G = ParamSpec("G")
-
-IndividualEventType = Union[EventSpec, EventHandler, Callable[G, Any], Var[Any]]
-
-EventType = Union[IndividualEventType[G], List[IndividualEventType[G]]]
-
 P = ParamSpec("P")
+Q = ParamSpec("Q")
 T = TypeVar("T")
 V = TypeVar("V")
 V2 = TypeVar("V2")
@@ -1474,49 +1470,47 @@ if sys.version_info >= (3, 10):
             self.func = func
 
         @overload
-        def __get__(
-            self: EventCallback[[V], T], instance: None, owner
-        ) -> Callable[[Union[Var[V], V]], EventSpec]: ...
+        def __call__(
+            self: EventCallback[Concatenate[V, Q], T], value: V | Var[V]
+        ) -> EventCallback[Q, T]: ...
 
         @overload
-        def __get__(
-            self: EventCallback[[V, V2], T], instance: None, owner
-        ) -> Callable[[Union[Var[V], V], Union[Var[V2], V2]], EventSpec]: ...
+        def __call__(
+            self: EventCallback[Concatenate[V, V2, Q], T],
+            value: V | Var[V],
+            value2: V2 | Var[V2],
+        ) -> EventCallback[Q, T]: ...
 
         @overload
-        def __get__(
-            self: EventCallback[[V, V2, V3], T], instance: None, owner
-        ) -> Callable[
-            [Union[Var[V], V], Union[Var[V2], V2], Union[Var[V3], V3]],
-            EventSpec,
-        ]: ...
+        def __call__(
+            self: EventCallback[Concatenate[V, V2, V3, Q], T],
+            value: V | Var[V],
+            value2: V2 | Var[V2],
+            value3: V3 | Var[V3],
+        ) -> EventCallback[Q, T]: ...
 
         @overload
-        def __get__(
-            self: EventCallback[[V, V2, V3, V4], T], instance: None, owner
-        ) -> Callable[
-            [
-                Union[Var[V], V],
-                Union[Var[V2], V2],
-                Union[Var[V3], V3],
-                Union[Var[V4], V4],
-            ],
-            EventSpec,
-        ]: ...
+        def __call__(
+            self: EventCallback[Concatenate[V, V2, V3, V4, Q], T],
+            value: V | Var[V],
+            value2: V2 | Var[V2],
+            value3: V3 | Var[V3],
+            value4: V4 | Var[V4],
+        ) -> EventCallback[Q, T]: ...
+
+        def __call__(self, *values) -> EventCallback:  # type: ignore
+            """Call the function with the values.
+
+            Args:
+                *values: The values to call the function with.
+
+            Returns:
+                The function with the values.
+            """
+            return self.func(*values)  # type: ignore
 
         @overload
-        def __get__(
-            self: EventCallback[[V, V2, V3, V4, V5], T], instance: None, owner
-        ) -> Callable[
-            [
-                Union[Var[V], V],
-                Union[Var[V2], V2],
-                Union[Var[V3], V3],
-                Union[Var[V4], V4],
-                Union[Var[V5], V5],
-            ],
-            EventSpec,
-        ]: ...
+        def __get__(self, instance: None, owner) -> Self: ...
 
         @overload
         def __get__(self, instance, owner) -> Callable[P, T]: ...
@@ -1548,6 +1542,9 @@ if sys.version_info >= (3, 10):
         return func  # type: ignore
 else:
 
+    class EventCallback(Generic[P, T]):
+        """A descriptor that wraps a function to be used as an event."""
+
     def event_handler(func: Callable[P, T]) -> Callable[P, T]:
         """Wrap a function to be used as an event.
 
@@ -1560,6 +1557,15 @@ else:
         return func
 
 
+G = ParamSpec("G")
+
+IndividualEventType = Union[
+    EventSpec, EventHandler, Callable[G, Any], EventCallback[G, Any], Var[Any]
+]
+
+EventType = Union[IndividualEventType[G], List[IndividualEventType[G]]]
+
+
 class EventNamespace(types.SimpleNamespace):
     """A namespace for event related classes."""