From a3857c9010611b6daa0f6c36c1e49c74682ad111 Mon Sep 17 00:00:00 2001 From: Muhammad Huzaifa Date: Wed, 21 Oct 2026 01:48:07 -0500 Subject: [PATCH] Fix static typing for set and frozenset array vars --- .../news/+set-array-static-types.bugfix.md | 1 + .../reflex-base/src/reflex_base/vars/base.py | 14 ++++-- .../src/reflex_base/vars/sequence.py | 44 ++++++++++++----- tests/units/test_var.py | 48 +++++++++++++++++-- 4 files changed, 88 insertions(+), 19 deletions(-) create mode 100644 packages/reflex-base/news/+set-array-static-types.bugfix.md diff --git a/packages/reflex-base/news/+set-array-static-types.bugfix.md b/packages/reflex-base/news/+set-array-static-types.bugfix.md new file mode 100644 index 00000000000..8d9296819fe --- /dev/null +++ b/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. diff --git a/packages/reflex-base/src/reflex_base/vars/base.py b/packages/reflex-base/src/reflex_base/vars/base.py index 332a0c956e2..7a38bfd179f 100644 --- a/packages/reflex-base/src/reflex_base/vars/base.py +++ b/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 diff --git a/packages/reflex-base/src/reflex_base/vars/sequence.py b/packages/reflex-base/src/reflex_base/vars/sequence.py index 6798fa0cc9c..78bfc7e9b87 100644 --- a/packages/reflex-base/src/reflex_base/vars/sequence.py +++ b/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: diff --git a/tests/units/test_var.py b/tests/units/test_var.py index b7f81370f27..8c3ffb02dd0 100644 --- a/tests/units/test_var.py +++ b/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 = (