Add ComputedVar overloads for BASE_TYPE, SQLA_TYPE, and DATACLASS_TYPE (

#4777) Allow typing to find __getattr__ for rx.var that returns an object-like model.
reflex-dev · Feb 8, 2025 · 8b2c729 · 8b2c729
1 parent 3de0415
commit 8b2c729
Showing 1 changed file with 42 additions and 0 deletions.
diff --git a/reflex/vars/base.py b/reflex/vars/base.py
@@ -2254,6 +2254,27 @@ def __get__(
         owner: Type,
     ) -> ArrayVar[tuple[LIST_INSIDE, ...]]: ...
 
+    @overload
+    def __get__(
+        self: ComputedVar[BASE_TYPE],
+        instance: None,
+        owner: Type,
+    ) -> ObjectVar[BASE_TYPE]: ...
+
+    @overload
+    def __get__(
+        self: ComputedVar[SQLA_TYPE],
+        instance: None,
+        owner: Type,
+    ) -> ObjectVar[SQLA_TYPE]: ...
+
+    if TYPE_CHECKING:
+
+        @overload
+        def __get__(
+            self: ComputedVar[DATACLASS_TYPE], instance: None, owner: Any
+        ) -> ObjectVar[DATACLASS_TYPE]: ...
+
     @overload
     def __get__(self, instance: None, owner: Type) -> ComputedVar[RETURN_TYPE]: ...
 
@@ -2500,6 +2521,27 @@ def __get__(
         owner: Type,
     ) -> ArrayVar[tuple[LIST_INSIDE, ...]]: ...
 
+    @overload
+    def __get__(
+        self: AsyncComputedVar[BASE_TYPE],
+        instance: None,
+        owner: Type,
+    ) -> ObjectVar[BASE_TYPE]: ...
+
+    @overload
+    def __get__(
+        self: AsyncComputedVar[SQLA_TYPE],
+        instance: None,
+        owner: Type,
+    ) -> ObjectVar[SQLA_TYPE]: ...
+
+    if TYPE_CHECKING:
+
+        @overload
+        def __get__(
+            self: AsyncComputedVar[DATACLASS_TYPE], instance: None, owner: Any
+        ) -> ObjectVar[DATACLASS_TYPE]: ...
+
     @overload
     def __get__(self, instance: None, owner: Type) -> AsyncComputedVar[RETURN_TYPE]: ...