Fix static typing for set and frozenset array vars

packages/reflex-base/news/+set-array-static-types.bugfix.md

@@ -0,0 +1 @@
+Fix static typing for `set` and `frozenset` array vars, preserving their container and element types when creating literals and indexing.

packages/reflex-base/src/reflex_base/vars/base.py

@@ -92,7 +92,13 @@
     from .color import LiteralColorVar
     from .number import BooleanVar, LiteralBooleanVar, LiteralNumberVar, NumberVar
     from .object import LiteralObjectVar, ObjectVar
-    from .sequence import ArrayVar, LiteralArrayVar, LiteralStringVar, StringVar
+    from .sequence import (
+        ARRAY_VAR_TYPE,
+        ArrayVar,
+        LiteralArrayVar,
+        LiteralStringVar,
+        StringVar,
+    )
 
 
 VAR_TYPE = TypeVar("VAR_TYPE", covariant=True)
@@ -1041,9 +1047,9 @@ def create(
     @classmethod
     def create(
         cls,
-        value: SEQUENCE_TYPE,
+        value: ARRAY_VAR_TYPE,
         _var_data: VarData | None = None,
-    ) -> LiteralArrayVar[SEQUENCE_TYPE]: ...
+    ) -> LiteralArrayVar[ARRAY_VAR_TYPE]: ...
 
     @overload
     @classmethod
@@ -2094,7 +2100,7 @@ def var_operation(
 ) -> Callable[P, StringVar]: ...
 
 
-LIST_T = TypeVar("LIST_T", bound=Sequence)
+LIST_T = TypeVar("LIST_T", bound=Sequence | set | frozenset)
 
 
 @overload

packages/reflex-base/src/reflex_base/vars/sequence.py

@@ -50,8 +50,12 @@
     from .function import FunctionVar
     from .object import ObjectVar
 
-ARRAY_VAR_TYPE = TypeVar("ARRAY_VAR_TYPE", bound=Sequence, covariant=True)
-OTHER_ARRAY_VAR_TYPE = TypeVar("OTHER_ARRAY_VAR_TYPE", bound=Sequence, covariant=True)
+ARRAY_VAR_TYPE = TypeVar(
+    "ARRAY_VAR_TYPE", bound=Sequence | set | frozenset, covariant=True
+)
+OTHER_ARRAY_VAR_TYPE = TypeVar(
+    "OTHER_ARRAY_VAR_TYPE", bound=Sequence | set | frozenset, covariant=True
+)
 MAPPING_VAR_TYPE = TypeVar("MAPPING_VAR_TYPE", bound=Mapping, covariant=True)
 
 OTHER_TUPLE = TypeVar("OTHER_TUPLE")
@@ -166,43 +170,57 @@ def __getitem__(
 
     @overload
     def __getitem__(
-        self: ArrayVar[Sequence[bool]], i: int | NumberVar
+        self: ArrayVar[Sequence[bool] | set[bool] | frozenset[bool]], i: int | NumberVar
     ) -> BooleanVar: ...
 
     @overload
     def __getitem__(
         self: (
-            ArrayVar[Sequence[int]]
-            | ArrayVar[Sequence[float]]
-            | ArrayVar[Sequence[int | float]]
+            ArrayVar[Sequence[int] | set[int] | frozenset[int]]
+            | ArrayVar[Sequence[float] | set[float] | frozenset[float]]
+            | ArrayVar[
+                Sequence[int | float] | set[int | float] | frozenset[int | float]
+            ]
         ),
         i: int | NumberVar,
     ) -> NumberVar: ...
 
     @overload
-    def __getitem__(self: ArrayVar[Sequence[str]], i: int | NumberVar) -> StringVar: ...
+    def __getitem__(
+        self: ArrayVar[Sequence[str] | set[str] | frozenset[str]], i: int | NumberVar
+    ) -> StringVar: ...
 
     @overload
     def __getitem__(
-        self: ArrayVar[Sequence[OTHER_ARRAY_VAR_TYPE]],
+        self: ArrayVar[
+            Sequence[OTHER_ARRAY_VAR_TYPE]
+            | set[OTHER_ARRAY_VAR_TYPE]
+            | frozenset[OTHER_ARRAY_VAR_TYPE]
+        ],
         i: int | NumberVar,
     ) -> ArrayVar[OTHER_ARRAY_VAR_TYPE]: ...
 
     @overload
     def __getitem__(
-        self: ArrayVar[Sequence[MAPPING_VAR_TYPE]],
+        self: ArrayVar[
+            Sequence[MAPPING_VAR_TYPE]
+            | set[MAPPING_VAR_TYPE]
+            | frozenset[MAPPING_VAR_TYPE]
+        ],
         i: int | NumberVar,
     ) -> ObjectVar[MAPPING_VAR_TYPE]: ...
 
     @overload
     def __getitem__(
-        self: ArrayVar[Sequence[SQLA_TYPE]],
+        self: ArrayVar[Sequence[SQLA_TYPE] | set[SQLA_TYPE] | frozenset[SQLA_TYPE]],
         i: int | NumberVar,
     ) -> ObjectVar[SQLA_TYPE]: ...
 
     @overload
     def __getitem__(
-        self: ArrayVar[Sequence[DATACLASS_TYPE]],
+        self: ArrayVar[
+            Sequence[DATACLASS_TYPE] | set[DATACLASS_TYPE] | frozenset[DATACLASS_TYPE]
+        ],
         i: int | NumberVar,
     ) -> ObjectVar[DATACLASS_TYPE]: ...
 
@@ -613,7 +631,9 @@ class LiteralArrayVar(
 ):
     """Base class for immutable literal array vars."""
 
-    _var_value: Sequence[Var | Any] = dataclasses.field(default=())
+    _var_value: Sequence[Var | Any] | set[Var | Any] | frozenset[Var | Any] = (
+        dataclasses.field(default=())
+    )
 
     @cached_property_no_lock
     def _cached_var_name(self) -> str:

tests/units/test_var.py

@@ -44,7 +44,12 @@
     DestructuredArg,
     FunctionStringVar,
 )
-from reflex_base.vars.number import LiteralBooleanVar, LiteralNumberVar, NumberVar
+from reflex_base.vars.number import (
+    BooleanVar,
+    LiteralBooleanVar,
+    LiteralNumberVar,
+    NumberVar,
+)
 from reflex_base.vars.object import LiteralObjectVar, ObjectVar
 from reflex_base.vars.sequence import (
     ArrayVar,
@@ -424,8 +429,21 @@ def test_date_comparison_uses_compare_datetime():
 
 def test_frozenset_literal_types_its_elements():
     """A frozen set literal infers its element type like a set does."""
-    # The static array bound is Sequence, which no set satisfies; see the PR.
-    var = cast(ArrayVar, v(frozenset({"a"})))
+    var = LiteralVar.create(frozenset({"a"}))
+    typing.assert_type(var, LiteralArrayVar[frozenset[str]])
+    typing.assert_type(
+        LiteralArrayVar.create(frozenset({"a"})), LiteralArrayVar[frozenset[str]]
+    )
+    typing.assert_type(var[0], StringVar)
+    typing.assert_type(var.reverse(), ArrayVar[frozenset[str]])
+    typing.assert_type(var + var, ArrayVar[frozenset[str]])
+    typing.assert_type(var * 2, ArrayVar[frozenset[str]])
+    typing.assert_type(var[:1], ArrayVar[frozenset[str]])
+    typing.assert_type(LiteralVar.create(frozenset({1}))[0], NumberVar)
+    typing.assert_type(LiteralVar.create(frozenset({True}))[0], BooleanVar)
+    typing.assert_type(
+        LiteralVar.create(frozenset({frozenset({"a"})}))[0], ArrayVar[frozenset[str]]
+    )
 
     assert var._var_type == frozenset[str]
     assert isinstance(var[0], StringVar)
@@ -435,6 +453,30 @@ def test_frozenset_literal_types_its_elements():
     )
 
 
+def test_set_literal_types_its_elements():
+    """A set literal preserves its array and element types."""
+    var = LiteralVar.create({"a"})
+    typing.assert_type(var, LiteralArrayVar[set[str]])
+    typing.assert_type(LiteralArrayVar.create({"a"}), LiteralArrayVar[set[str]])
+    typing.assert_type(var[0], StringVar)
+    typing.assert_type(var.reverse(), ArrayVar[set[str]])
+    typing.assert_type(var + var, ArrayVar[set[str]])
+    typing.assert_type(var * 2, ArrayVar[set[str]])
+    typing.assert_type(var[:1], ArrayVar[set[str]])
+    typing.assert_type(LiteralVar.create({1})[0], NumberVar)
+    typing.assert_type(LiteralVar.create({True})[0], BooleanVar)
+    typing.assert_type(
+        LiteralVar.create({frozenset({"a"})})[0], ArrayVar[frozenset[str]]
+    )
+
+    assert var._var_type == set[str]
+    assert isinstance(var[0], StringVar)
+    assert re.fullmatch(
+        r'\["a"\]\.map\(\(\((\w+)\) => \1\.toUpperCase\(\)\)\)',
+        str(var.map(lambda item: item.upper())),
+    )
+
+
 def test_datetime_equality_with_other_type_uses_default_comparison():
     value = v(datetime(2024, 1, 1))
     expected_equality = (