[REF-3328] Implement __getitem__ for ArrayVar

reflex/experimental/vars/object.py

@@ -4,6 +4,7 @@
 
 import dataclasses
 import sys
+import typing
 from functools import cached_property
 from typing import Any, Dict, Tuple, Type, Union
 
@@ -89,6 +90,23 @@ def __getattr__(self, name) -> ObjectItemOperation:
         return ObjectItemOperation(self, name)
 
 
+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)]
+
+
 @dataclasses.dataclass(
     eq=False,
     frozen=True,
@@ -116,7 +134,14 @@ def __init__(
         """
         super(LiteralObjectVar, self).__init__(
             _var_name="",
-            _var_type=type(_var_value) if _var_type is None else _var_type,
+            _var_type=(
+                Dict[
+                    unionize(*map(type, _var_value.keys())),
+                    unionize(*map(type, _var_value.values())),
+                ]
+                if _var_type is None
+                else _var_type
+            ),
             _var_data=ImmutableVarData.merge(_var_data),
         )
         object.__setattr__(
@@ -132,20 +157,17 @@ def _key_type(self) -> Type:
         Returns:
             The type of the keys of the object.
         """
-        print(self._var_type)
-        return (
-            self._var_type.__args__[0] if hasattr(self._var_type, "__args__") else Any
-        )
+        args_list = typing.get_args(self._var_type)
+        return args_list[0] if args_list else Any
 
     def _value_type(self) -> Type:
         """Get the type of the values of the object.
 
         Returns:
             The type of the values of the object.
         """
-        return (
-            self._var_type.__args__[1] if hasattr(self._var_type, "__args__") else Any
-        )
+        args_list = typing.get_args(self._var_type)
+        return args_list[1] if args_list else Any
 
     def __getattr__(self, name):
         """Get an attribute of the var.

reflex/experimental/vars/sequence.py

@@ -7,6 +7,7 @@
 import json
 import re
 import sys
+import typing
 from functools import cached_property
 from typing import Any, List, Set, Tuple, Union, overload
 
@@ -1238,11 +1239,10 @@ def __init__(
             i: The index.
             _var_data: Additional hooks and imports associated with the Var.
         """
+        element_type = typing.get_args(a._var_type)
         super(ArrayItemOperation, self).__init__(
             _var_name="",
-            _var_type=(
-                a._var_type.__args__[0] if hasattr(a._var_type, "__args__") else Any
-            ),
+            _var_type=element_type[0] if element_type else Any,
             _var_data=ImmutableVarData.merge(_var_data),
         )
         object.__setattr__(self, "a", a if isinstance(a, Var) else LiteralArrayVar(a))
@@ -1319,7 +1319,7 @@ def __init__(
         """
         super(RangeOperation, self).__init__(
             _var_name="",
-            _var_type=list,
+            _var_type=List[int],
             _var_data=ImmutableVarData.merge(_var_data),
         )
         object.__setattr__(