add type hinting to existing types

reflex/experimental/vars/base.py

@@ -10,9 +10,15 @@
     TYPE_CHECKING,
     Any,
     Callable,
+    Dict,
+    Generic,
+    List,
     Optional,
+    Set,
+    Tuple,
     Type,
     TypeVar,
+    Union,
     overload,
 )
 
@@ -42,13 +48,15 @@
     from .object import ObjectVar, ToObjectOperation
     from .sequence import ArrayVar, StringVar, ToArrayOperation, ToStringOperation
 
+VAR_TYPE = TypeVar("VAR_TYPE")
+
 
 @dataclasses.dataclass(
     eq=False,
     frozen=True,
     **{"slots": True} if sys.version_info >= (3, 10) else {},
 )
-class ImmutableVar(Var):
+class ImmutableVar(Var, Generic[VAR_TYPE]):
     """Base class for immutable vars."""
 
     # The name of the var.
@@ -405,6 +413,8 @@ def guess_type(self) -> ImmutableVar:
             return self.to(ArrayVar, var_type)
         if issubclass(fixed_type, str):
             return self.to(StringVar)
+        if issubclass(fixed_type, Base):
+            return self.to(ObjectVar, var_type)
         return self
 
 
@@ -531,3 +541,43 @@ def wrapper(*args: P.args, **kwargs: P.kwargs) -> T:
         return wrapper
 
     return decorator
+
+
+def unionize(*args: Type) -> Type:
+    """Unionize the types.
+
+    Args:
+        args: The types to unionize.
+
+    Returns:
+        The unionized types.
+    """
+    if not args:
+        return Any
+    first, *rest = args
+    if not rest:
+        return first
+    return Union[first, unionize(*rest)]
+
+
+def figure_out_type(value: Any) -> Type:
+    """Figure out the type of the value.
+
+    Args:
+        value: The value to figure out the type of.
+
+    Returns:
+        The type of the value.
+    """
+    if isinstance(value, list):
+        return List[unionize(*(figure_out_type(v) for v in value))]
+    if isinstance(value, set):
+        return Set[unionize(*(figure_out_type(v) for v in value))]
+    if isinstance(value, tuple):
+        return Tuple[unionize(*(figure_out_type(v) for v in value)), ...]
+    if isinstance(value, dict):
+        return Dict[
+            unionize(*(figure_out_type(k) for k in value)),
+            unionize(*(figure_out_type(v) for v in value.values())),
+        ]
+    return type(value)

reflex/experimental/vars/function.py

@@ -11,7 +11,7 @@
 from reflex.vars import ImmutableVarData, Var, VarData
 
 
-class FunctionVar(ImmutableVar):
+class FunctionVar(ImmutableVar[Callable]):
     """Base class for immutable function vars."""
 
     def __call__(self, *args: Var | Any) -> ArgsFunctionOperation:

reflex/experimental/vars/number.py

@@ -15,7 +15,7 @@
 from reflex.vars import ImmutableVarData, Var, VarData
 
 
-class NumberVar(ImmutableVar):
+class NumberVar(ImmutableVar[Union[int, float]]):
     """Base class for immutable number vars."""
 
     def __add__(self, other: number_types | boolean_types) -> NumberAddOperation:
@@ -693,7 +693,7 @@ def _cached_var_name(self) -> str:
         return f"Math.trunc({str(value)})"
 
 
-class BooleanVar(ImmutableVar):
+class BooleanVar(ImmutableVar[bool]):
     """Base class for immutable boolean vars."""
 
     def __and__(self, other: bool) -> BooleanAndOperation: