Repository: reflex-dev/reflex

Issue:
[REF-2682] Foreach over dict gets incorrect arg type
**Describe the bug**
The arg var resulting from iterating over a Var of type dict should be a tuple of the key and value types. But it's actually just the same type as the Var itself.

**To Reproduce**

```python
import reflex as rx


class MyState(rx.State):

    labels: dict[str, str] = {
        "A": "Alpha",
        "a": "alpha",
        "b": "beta",
    }
    some_values: dict[str, str] = {
        "A": "a",
    }
    some_values_dict: dict[str, list[str]] = {
        "A": ["a", "b"],
    }


def make_nice_label(entry: list[str, str]) -> rx.Component:
    key = entry[0]
    value = entry[1]
    print(repr(entry), repr(key), repr(value))
    labels = MyState.labels
    return rx.vstack(rx.heading(labels[key]), rx.text(labels[value]))


def make_nice_label_dict(entry: list[str, str]) -> rx.Component:
    key = entry[0]
    values = entry[1]
    print(repr(entry), repr(key), repr(values))
    labels = MyState.labels
    return rx.vstack(
        rx.heading(labels[key]),
        rx.foreach(
            values,
            lambda value: rx.text(labels[value]),
        ),
    )


@rx.page(route="/")
def index() -> rx.Component:
    return rx.foreach(
        MyState.some_values,
        make_nice_label,
    ), rx.foreach(
        MyState.some_values_dict,
        make_nice_label_dict,
    )


app = rx.App()
```

**Expected behavior**
Should display

```
Alpha
alpha

Alpha
alpha
beta
```

Should print

```console
BaseVar(_var_name='entry', _var_type=typing.Tuple[str, str], _var_is_local=False, _var_is_string=False, _var_full_name_needs_state_prefix=False, _var_data=None) BaseVar(_var_name='entry.at(0)', _var_type=<class 'str'>, _var_is_local=False, _var_is_string=False, _var_full_name_needs_state_prefix=False, _var_data=None) BaseVar(_var_name='entry.at(1)', _var_type=<class 'str'>, _var_is_local=False, _var_is_string=False, _var_full_name_needs_state_prefix=False, _var_data=None)
BaseVar(_var_name='entry', _var_type=typing.Tuple[str, list[str]], _var_is_local=False, _var_is_string=False, _var_full_name_needs_state_prefix=False, _var_data=None) BaseVar(_var_name='entry.at(0)', _var_type=<class 'str'>, _var_is_local=False, _var_is_string=False, _var_full_name_needs_state_prefix=False, _var_data=None) BaseVar(_var_name='entry.at(1)', _var_type=list[str], _var_is_local=False, _var_is_string=False, _var_full_name_needs_state_prefix=False, _var_data=None)
```

Instead, it crashes:

```console
BaseVar(_var_name='entry', _var_type=dict[str, str], _var_is_local=False, _var_is_string=False, _var_full_name_needs_state_prefix=False, _var_data=None) BaseVar(_var_name='entry[0]', _var_type=<class 'str'>, _var_is_local=False, _var_is_string=False, _var_full_name_needs_state_prefix=False, _var_data=None) BaseVar(_var_name='entry[1]', _var_type=<class 'str'>, _var_is_local=False, _var_is_string=False, _var_full_name_needs_state_prefix=False, _var_data=None)
BaseVar(_var_name='entry', _var_type=dict[str, list[str]], _var_is_local=False, _var_is_string=False, _var_full_name_needs_state_prefix=False, _var_data=None) BaseVar(_var_name='entry[0]', _var_type=list[str], _var_is_local=False, _var_is_string=False, _var_full_name_needs_state_prefix=False, _var_data=None) BaseVar(_var_name='entry[1]', _var_type=list[str], _var_is_local=False, _var_is_string=False, _var_full_name_needs_state_prefix=False, _var_data=None)
Traceback (most recent call last):
  File "/Users/masen/code/reflex-dev/VENV-dev311/bin/reflex", line 8, in <module>
    sys.exit(cli())
             ^^^^^
  File "/Users/masen/code/reflex-dev/VENV-dev311/lib/python3.11/site-packages/typer/main.py", line 326, in __call__
    raise e
  File "/Users/masen/code/reflex-dev/VENV-dev311/lib/python3.11/site-packages/typer/main.py", line 309, in __call__
    return get_command(self)(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/masen/code/reflex-dev/VENV-dev311/lib/python3.11/site-packages/click/core.py", line 1157, in __call__
    return self.main(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/masen/code/reflex-dev/VENV-dev311/lib/python3.11/site-packages/typer/core.py", line 723, in main
    return _main(
           ^^^^^^
  File "/Users/masen/code/reflex-dev/VENV-dev311/lib/python3.11/site-packages/typer/core.py", line 193, in _main
    rv = self.invoke(ctx)
         ^^^^^^^^^^^^^^^^
  File "/Users/masen/code/reflex-dev/VENV-dev311/lib/python3.11/site-packages/click/core.py", line 1688, in invoke
    return _process_result(sub_ctx.command.invoke(sub_ctx))
                           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/masen/code/reflex-dev/VENV-dev311/lib/python3.11/site-packages/click/core.py", line 1434, in invoke
    return ctx.invoke(self.callback, **ctx.params)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/masen/code/reflex-dev/VENV-dev311/lib/python3.11/site-packages/click/core.py", line 783, in invoke
    return __callback(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/masen/code/reflex-dev/VENV-dev311/lib/python3.11/site-packages/typer/main.py", line 688, in wrapper
    return callback(**use_params)
           ^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/masen/code/reflex-dev/reflex/reflex/reflex.py", line 256, in run
    _run(env, frontend, backend, frontend_port, backend_port, backend_host, loglevel)
  File "/Users/masen/code/reflex-dev/reflex/reflex/reflex.py", line 182, in _run
    prerequisites.get_compiled_app()
  File "/Users/masen/code/reflex-dev/reflex/reflex/utils/prerequisites.py", line 242, in get_compiled_app
    app._apply_decorated_pages()
  File "/Users/masen/code/reflex-dev/reflex/reflex/app.py", line 735, in _apply_decorated_pages
    self.add_page(render, **kwargs)
  File "/Users/masen/code/reflex-dev/reflex/reflex/app.py", line 455, in add_page
    component = self._generate_component(component)
                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/masen/code/reflex-dev/reflex/reflex/app.py", line 406, in _generate_component
    raise e
  File "/Users/masen/code/reflex-dev/reflex/reflex/app.py", line 395, in _generate_component
    return component if isinstance(component, Component) else component()
                                                              ^^^^^^^^^^^
  File "/Users/masen/code/reflex-dev/repro-complex-indexing/repro_complex_indexing/repro_complex_indexing.py", line 46, in index
    ), rx.foreach(
       ^^^^^^^^^^^
  File "/Users/masen/code/reflex-dev/reflex/reflex/components/core/foreach.py", line 64, in create
    component._render(props=dict(index_var_name="i")).render_component()
  File "/Users/masen/code/reflex-dev/reflex/reflex/components/tags/iter_tag.py", line 117, in render_component
    component = self.render_fn(arg)
                ^^^^^^^^^^^^^^^^^^^
  File "/Users/masen/code/reflex-dev/repro-complex-indexing/repro_complex_indexing/repro_complex_indexing.py", line 33, in make_nice_label_dict
    rx.heading(labels[key]),
               ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~^^^^^
  File "/Users/masen/code/reflex-dev/reflex/reflex/vars.py", line 661, in __getitem__
    raise TypeError(
TypeError: Index must be one of the following types: int, str, int or str Var
```

**Specifics (please complete the following information):**

* Python Version: 3.11.8
* Reflex Version: 0.4.9
* OS: macOS 14
* Browser (Optional):

**Additional context**
In `IterTag.get_iterable_var_type`, we could use `Tuple[self.iterable._var_type.__args__[0], self.iterable._var_type.__args__[1]]` if the \_var_type is `dict`

<sub>From [SyncLinear.com](https://synclinear.com) | [REF-2682](https://linear.app/reflex-dev/issue/REF-2682/foreach-over-dict-gets-incorrect-arg-type)</sub>


Discussion:

Patch:
diff --git a/reflex/components/tags/iter_tag.py b/reflex/components/tags/iter_tag.py
index a0797193cc2..7ca05047057 100644
--- a/reflex/components/tags/iter_tag.py
+++ b/reflex/components/tags/iter_tag.py
@@ -2,7 +2,7 @@
 from __future__ import annotations
 
 import inspect
-from typing import TYPE_CHECKING, Any, Callable, List, Type
+from typing import TYPE_CHECKING, Any, Callable, List, Tuple, Type, Union, get_args
 
 from reflex.components.tags.tag import Tag
 from reflex.vars import BaseVar, Var
@@ -33,11 +33,14 @@ def get_iterable_var_type(self) -> Type:
             The type of the iterable var.
         """
         try:
-            return (
-                self.iterable._var_type
-                if self.iterable._var_type.mro()[0] == dict
-                else self.iterable._var_type.__args__[0]
-            )
+            if self.iterable._var_type.mro()[0] == dict:
+                # Arg is a tuple of (key, value).
+                return Tuple[get_args(self.iterable._var_type)]  # type: ignore
+            elif self.iterable._var_type.mro()[0] == tuple:
+                # Arg is a union of any possible values in the tuple.
+                return Union[get_args(self.iterable._var_type)]  # type: ignore
+            else:
+                return get_args(self.iterable._var_type)[0]
         except Exception:
             return Any
 


Test patch:
diff --git a/tests/components/core/test_foreach.py b/tests/components/core/test_foreach.py
index bfd97265ed7..db0488e6840 100644
--- a/tests/components/core/test_foreach.py
+++ b/tests/components/core/test_foreach.py
@@ -1,10 +1,11 @@
-from typing import Dict, List, Set, Tuple
+from typing import Dict, List, Set, Tuple, Union
 
 import pytest
 
 from reflex.components import box, foreach, text, theme
 from reflex.components.core import Foreach
 from reflex.state import BaseState
+from reflex.vars import Var
 
 try:
     # When pydantic v2 is installed
@@ -39,29 +40,36 @@ class ForEachState(BaseState):
     )
     colors_set: Set[str] = {"red", "green"}
     bad_annotation_list: list = [["red", "orange"], ["yellow", "blue"]]
+    color_index_tuple: Tuple[int, str] = (0, "red")
 
 
 def display_color(color):
+    assert color._var_type == str
     return box(text(color))
 
 
 def display_color_name(color):
+    assert color._var_type == Dict[str, str]
     return box(text(color["name"]))
 
 
 def display_shade(color):
+    assert color._var_type == Dict[str, List[str]]
     return box(text(color["shades"][0]))
 
 
 def display_primary_colors(color):
+    assert color._var_type == Tuple[str, str]
     return box(text(color[0]), text(color[1]))
 
 
 def display_color_with_shades(color):
+    assert color._var_type == Tuple[str, List[str]]
     return box(text(color[0]), text(color[1][0]))
 
 
 def display_nested_color_with_shades(color):
+    assert color._var_type == Tuple[str, Dict[str, List[Dict[str, str]]]]
     return box(text(color[0]), text(color[1]["red"][0]["shade"]))
 
 
@@ -70,21 +78,31 @@ def show_shade(item):
 
 
 def display_nested_color_with_shades_v2(color):
+    assert color._var_type == Tuple[str, Dict[str, List[Dict[str, str]]]]
     return box(text(foreach(color[1], show_shade)))
 
 
 def display_color_tuple(color):
+    assert color._var_type == str
     return box(text(color, "tuple"))
 
 
 def display_colors_set(color):
+    assert color._var_type == str
     return box(text(color, "set"))
 
 
-def display_nested_list_element(element: str, index: int):
+def display_nested_list_element(element: Var[str], index: Var[int]):
+    assert element._var_type == List[str]
+    assert index._var_type == int
     return box(text(element[index]))
 
 
+def display_color_index_tuple(color):
+    assert color._var_type == Union[int, str]
+    return box(text(color, "index_tuple"))
+
+
 seen_index_vars = set()
 
 
@@ -171,6 +189,14 @@ def display_nested_list_element(element: str, index: int):
                 "iterable_type": "list",
             },
         ),
+        (
+            ForEachState.color_index_tuple,
+            display_color_index_tuple,
+            {
+                "iterable_state": "for_each_state.color_index_tuple",
+                "iterable_type": "tuple",
+            },
+        ),
     ],
 )
 def test_foreach_render(state_var, render_fn, render_dict):