Set correct type when indexing into Var[str] (#2469)

* Index into strings

* Write tests
This commit is contained in:
Nikhil Rao 2024-01-29 12:57:02 +07:00 committed by GitHub
parent b2c749fc9f
commit 01c2a1ed7d
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 20 additions and 5 deletions

View File

@ -572,11 +572,10 @@ class Var:
)
# Get the type of the indexed var.
type_ = (
types.get_args(self._var_type)[0]
if types.is_generic_alias(self._var_type)
else Any
)
if types.is_generic_alias(self._var_type):
type_ = types.get_args(self._var_type)[0]
elif types._issubclass(self._var_type, str):
type_ = str
# Use `at` to support negative indices.
return self._replace(

View File

@ -406,6 +406,22 @@ def test_var_indexing_lists(var):
assert str(var[-1]) == f"{{{var._var_name}.at(-1)}}"
def test_var_indexing_str():
"""Test that we can index into str vars."""
str_var = BaseVar(_var_name="str", _var_type=str)
# Test that indexing gives a type of Var[str].
assert isinstance(str_var[0], Var)
assert str_var[0]._var_type == str
# Test basic indexing.
assert str(str_var[0]) == "{str.at(0)}"
assert str(str_var[1]) == "{str.at(1)}"
# Test negative indexing.
assert str(str_var[-1]) == "{str.at(-1)}"
@pytest.mark.parametrize(
"var, index",
[