From d773d704b401840d1d0c915bf4be99ee21fc8edc Mon Sep 17 00:00:00 2001
From: Khaleel Al-Adhami <khaleel.aladhami@gmail.com>
Date: Mon, 4 Nov 2024 10:59:32 -0800
Subject: [PATCH] fix typing issues

---
 reflex/vars/base.py             | 27 +++++++++++++--------------
 reflex/vars/object.py           |  6 ++++++
 tests/units/vars/test_object.py |  8 ++------
 3 files changed, 21 insertions(+), 20 deletions(-)

diff --git a/reflex/vars/base.py b/reflex/vars/base.py
index 303cabcc2..2599ec5fd 100644
--- a/reflex/vars/base.py
+++ b/reflex/vars/base.py
@@ -75,7 +75,6 @@ from reflex.utils.types import (
 if TYPE_CHECKING:
     from reflex.state import BaseState
 
-    from .function import FunctionVar
     from .number import (
         BooleanVar,
         NumberVar,
@@ -583,36 +582,33 @@ class Var(Generic[VAR_TYPE]):
         # Encode the _var_data into the formatted output for tracking purposes.
         return f"{constants.REFLEX_VAR_OPENING_TAG}{hashed_var}{constants.REFLEX_VAR_CLOSING_TAG}{self._js_expr}"
 
-    @overload
-    def to(self, output: Type[StringVar]) -> StringVar: ...
-
     @overload
     def to(self, output: Type[str]) -> StringVar: ...
 
     @overload
-    def to(self, output: Type[BooleanVar]) -> BooleanVar: ...
+    def to(self, output: Type[bool]) -> BooleanVar: ...
 
     @overload
-    def to(
-        self, output: Type[NumberVar], var_type: type[int] | type[float] = float
-    ) -> NumberVar: ...
+    def to(self, output: type[int] | type[float] = float) -> NumberVar: ...
 
     @overload
     def to(
         self,
-        output: Type[ArrayVar],
-        var_type: type[list] | type[tuple] | type[set] = list,
+        output: type[list] | type[tuple] | type[set],
     ) -> ArrayVar: ...
 
     @overload
     def to(
-        self, output: Type[ObjectVar], var_type: types.GenericType = dict
-    ) -> ObjectVar: ...
+        self, output: Type[ObjectVar], var_type: Type[VAR_INSIDE]
+    ) -> ObjectVar[VAR_INSIDE]: ...
 
     @overload
     def to(
-        self, output: Type[FunctionVar], var_type: Type[Callable] = Callable
-    ) -> FunctionVar: ...
+        self, output: Type[ObjectVar], var_type: None = None
+    ) -> ObjectVar[VAR_TYPE]: ...
+
+    @overload
+    def to(self, output: VAR_SUBCLASS, var_type: None = None) -> VAR_SUBCLASS: ...
 
     @overload
     def to(
@@ -1189,6 +1185,9 @@ class Var(Generic[VAR_TYPE]):
 
 OUTPUT = TypeVar("OUTPUT", bound=Var)
 
+VAR_SUBCLASS = TypeVar("VAR_SUBCLASS", bound=Var)
+VAR_INSIDE = TypeVar("VAR_INSIDE")
+
 
 class ToOperation:
     """A var operation that converts a var to another type."""
diff --git a/reflex/vars/object.py b/reflex/vars/object.py
index 4d5827527..e60ea09e3 100644
--- a/reflex/vars/object.py
+++ b/reflex/vars/object.py
@@ -239,6 +239,12 @@ class ObjectVar(Var[OBJECT_TYPE], python_types=dict):
         name: str,
     ) -> ObjectVar[dict[OTHER_KEY_TYPE, VALUE_TYPE]]: ...
 
+    @overload
+    def __getattr__(
+        self: ObjectVar,
+        name: str,
+    ) -> ObjectItemOperation: ...
+
     def __getattr__(self, name) -> Var:
         """Get an attribute of the var.
 
diff --git a/tests/units/vars/test_object.py b/tests/units/vars/test_object.py
index d762b208e..019c8ac4b 100644
--- a/tests/units/vars/test_object.py
+++ b/tests/units/vars/test_object.py
@@ -1,4 +1,4 @@
-from typing import assert_type, reveal_type
+from typing import assert_type
 
 import pytest
 
@@ -96,18 +96,14 @@ def test_state_to_operation(type_: GenericType) -> None:
 def test_typing() -> None:
     # Bare
     var = ObjectState.bare.to(ObjectVar)
-    reveal_type(var)
     _ = assert_type(var, ObjectVar[Bare])
 
     var = ObjectState.base.to(ObjectVar, Base)
-    reveal_type(var)
     _ = assert_type(var, ObjectVar[Base])
 
     # Base
     var = ObjectState.base.to(ObjectVar)
-    reveal_type(var)
     _ = assert_type(var, ObjectVar[Base])
 
     var = ObjectState.base.to(LiteralObjectVar, Base)
-    reveal_type(var)
-    _ = assert_type(var, LiteralObjectVar[Base])
+    _ = assert_type(var, ObjectVar[Base])