From 0b771db10a2b4f806c8c7898172716d938b8c36c Mon Sep 17 00:00:00 2001
From: wassaf shahzad <wassafshahzad@gmail.com>
Date: Wed, 21 Feb 2024 00:02:20 +0100
Subject: [PATCH] fixed bug in var type for iterable types (#2617)

* fixed bug in var type for iterable types

* added test cases

* formatting issue

* fixed black formatting  issues
---
 reflex/vars.py    |  4 +++-
 tests/test_var.py | 19 +++++++++++++++++++
 2 files changed, 22 insertions(+), 1 deletion(-)

diff --git a/reflex/vars.py b/reflex/vars.py
index 2bc0fc572..93f8363b9 100644
--- a/reflex/vars.py
+++ b/reflex/vars.py
@@ -623,7 +623,9 @@ class Var:
 
             # Get the type of the indexed var.
             if types.is_generic_alias(self._var_type):
-                type_ = types.get_args(self._var_type)[0]
+                index = i if not isinstance(i, Var) else 0
+                type_ = types.get_args(self._var_type)
+                type_ = type_[index % len(type_)]
             elif types._issubclass(self._var_type, str):
                 type_ = str
 
diff --git a/tests/test_var.py b/tests/test_var.py
index 3d7d23d89..a4797532c 100644
--- a/tests/test_var.py
+++ b/tests/test_var.py
@@ -459,6 +459,25 @@ def test_var_indexing_lists(var):
     assert str(var[-1]) == f"{{{var._var_name}.at(-1)}}"
 
 
+@pytest.mark.parametrize(
+    "var, type_",
+    [
+        (BaseVar(_var_name="list", _var_type=List[int]), [int, int]),
+        (BaseVar(_var_name="tuple", _var_type=Tuple[int, str]), [int, str]),
+    ],
+)
+def test_var_indexing_types(var, type_):
+    """Test that indexing returns valid types.
+
+    Args:
+        var   : The list, typle base var.
+        type_ : The type on indexed object.
+
+    """
+    assert var[2]._var_type == type_[0]
+    assert var[3]._var_type == type_[1]
+
+
 def test_var_indexing_str():
     """Test that we can index into str vars."""
     str_var = BaseVar(_var_name="str", _var_type=str)