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
This commit is contained in:
wassaf shahzad 2024-02-21 00:02:20 +01:00 committed by GitHub
parent 37f66207fc
commit 0b771db10a
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 22 additions and 1 deletions

View File

@ -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

View File

@ -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)