reflex_copy/tests/units/vars/test_base.py

from typing import List, Mapping, Union

import pytest

from reflex.state import State
from reflex.vars.base import computed_var, figure_out_type


class CustomDict(dict[str, str]):
    """A custom dict with generic arguments."""

    pass


class ChildCustomDict(CustomDict):
    """A child of CustomDict."""

    pass


class GenericDict(dict):
    """A generic dict with no generic arguments."""

    pass


class ChildGenericDict(GenericDict):
    """A child of GenericDict."""

    pass


@pytest.mark.parametrize(
    ("value", "expected"),
    [
        (1, int),
        (1.0, float),
        ("a", str),
        ([1, 2, 3], List[int]),
        ([1, 2.0, "a"], List[Union[int, float, str]]),
        ({"a": 1, "b": 2}, Mapping[str, int]),
        ({"a": 1, 2: "b"}, Mapping[Union[int, str], Union[str, int]]),
        (CustomDict(), CustomDict),
        (ChildCustomDict(), ChildCustomDict),
        (GenericDict({1: 1}), Mapping[int, int]),
        (ChildGenericDict({1: 1}), Mapping[int, int]),
    ],
)
def test_figure_out_type(value, expected):
    assert figure_out_type(value) == expected


def test_computed_var_replace() -> None:
    class StateTest(State):
        @computed_var(cache=True)
        def cv(self) -> int:
            return 1

    cv = StateTest.cv
    assert cv._var_type is int

    replaced = cv._replace(_var_type=float)
    assert replaced._var_type is float