From 8b2c7291d3ed605fc1cd82192555ac76ac8b76d3 Mon Sep 17 00:00:00 2001
From: Masen Furer <m_github@0x26.net>
Date: Fri, 7 Feb 2025 17:38:42 -0800
Subject: [PATCH] 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/vars/base.py | 42 ++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 42 insertions(+)

diff --git a/reflex/vars/base.py b/reflex/vars/base.py
index 0d8af8f3c..a24db4010 100644
--- a/reflex/vars/base.py
+++ b/reflex/vars/base.py
@@ -2254,6 +2254,27 @@ class ComputedVar(Var[RETURN_TYPE]):
         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 @@ class AsyncComputedVar(ComputedVar[RETURN_TYPE]):
         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]: ...