From 0a8aaea599f95bb6c4730c743f6c4b636134ca6a Mon Sep 17 00:00:00 2001
From: Masen Furer <m_github@0x26.net>
Date: Thu, 25 Apr 2024 09:10:55 -0700
Subject: [PATCH] [REF-2682] Foreach over dict uses Tuple arg value (#3160)

* test_foreach: assert on arg _var_type

* [REF-2682] Foreach over dict uses Tuple arg value

When iterating over a Var with _var_type dict, the resulting arg value
_var_type should be Tuple[key, value] so it can be correctly used with other
var operations.

Fix #3157

* Correct _var_type for iteration over Tuple of multiple types

The arg value when iterating over a tuple could be any of the possible values
mentioned in the annotation.

When only one type is used, the Union collapses to the base type, at least in py3.11

* Add comments
---
 reflex/components/tags/iter_tag.py    | 15 ++++++++------
 tests/components/core/test_foreach.py | 30 +++++++++++++++++++++++++--
 2 files changed, 37 insertions(+), 8 deletions(-)

diff --git a/reflex/components/tags/iter_tag.py b/reflex/components/tags/iter_tag.py
index a0797193c..7ca050470 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 @@ class IterTag(Tag):
             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
 
diff --git a/tests/components/core/test_foreach.py b/tests/components/core/test_foreach.py
index bfd97265e..db0488e68 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 @@ seen_index_vars = set()
                 "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):