reflex_copy/reflex/experimental/client_state.py

"""Handle client side state with `useState`."""

import dataclasses
import sys
from typing import Any, Callable, Optional, Type

from reflex import constants
from reflex.event import EventChain, EventHandler, EventSpec, call_script
from reflex.utils.imports import ImportVar
from reflex.vars import Var, VarData


def _client_state_ref(var_name: str) -> str:
    """Get the ref path for a ClientStateVar.

    Args:
        var_name: The name of the variable.

    Returns:
        An accessor for ClientStateVar ref as a string.
    """
    return f"refs['_client_state_{var_name}']"


@dataclasses.dataclass(
    eq=False,
    **{"slots": True} if sys.version_info >= (3, 10) else {},
)
class ClientStateVar(Var):
    """A Var that exists on the client via useState."""

    # The name of the var.
    _var_name: str = dataclasses.field()

    # Track the names of the getters and setters
    _setter_name: str = dataclasses.field()
    _getter_name: str = dataclasses.field()

    # The type of the var.
    _var_type: Type = dataclasses.field(default=Any)

    # Whether this is a local javascript variable.
    _var_is_local: bool = dataclasses.field(default=False)

    # Whether the var is a string literal.
    _var_is_string: bool = dataclasses.field(default=False)

    # _var_full_name should be prefixed with _var_state
    _var_full_name_needs_state_prefix: bool = dataclasses.field(default=False)

    # Extra metadata associated with the Var
    _var_data: Optional[VarData] = dataclasses.field(default=None)

    def __hash__(self) -> int:
        """Define a hash function for a var.

        Returns:
            The hash of the var.
        """
        return hash(
            (self._var_name, str(self._var_type), self._getter_name, self._setter_name)
        )

    @classmethod
    def create(cls, var_name, default=None) -> "ClientStateVar":
        """Create a local_state Var that can be accessed and updated on the client.

        The `ClientStateVar` should be included in the highest parent component
        that contains the components which will access and manipulate the client
        state. It has no visual rendering, including it ensures that the
        `useState` hook is called in the correct scope.

        To render the var in a component, use the `value` property.

        To update the var in a component, use the `set` property.

        To access the var in an event handler, use the `retrieve` method with
        `callback` set to the event handler which should receive the value.

        To update the var in an event handler, use the `push` method with the
        value to update.

        Args:
            var_name: The name of the variable.
            default: The default value of the variable.

        Returns:
            ClientStateVar
        """
        if default is None:
            default_var = Var.create_safe("", _var_is_local=False, _var_is_string=False)
        elif not isinstance(default, Var):
            default_var = Var.create_safe(default)
        else:
            default_var = default
        setter_name = f"set{var_name.capitalize()}"
        return cls(
            _var_name="",
            _setter_name=setter_name,
            _getter_name=var_name,
            _var_is_local=False,
            _var_is_string=False,
            _var_type=default_var._var_type,
            _var_data=VarData.merge(
                default_var._var_data,
                VarData(  # type: ignore
                    hooks={
                        f"const [{var_name}, {setter_name}] = useState({default_var._var_name_unwrapped})": None,
                        f"{_client_state_ref(var_name)} = {var_name}": None,
                        f"{_client_state_ref(setter_name)} = {setter_name}": None,
                    },
                    imports={
                        "react": {ImportVar(tag="useState", install=False)},
                        f"/{constants.Dirs.STATE_PATH}": [ImportVar(tag="refs")],
                    },
                ),
            ),
        )

    @property
    def value(self) -> Var:
        """Get a placeholder for the Var.

        This property can only be rendered on the frontend.

        To access the value in a backend event handler, see `retrieve`.

        Returns:
            an accessor for the client state variable.
        """
        return (
            Var.create_safe(
                _client_state_ref(self._getter_name),
                _var_is_local=False,
                _var_is_string=False,
            )
            .to(self._var_type)
            ._replace(
                merge_var_data=VarData(  # type: ignore
                    imports={
                        f"/{constants.Dirs.STATE_PATH}": [ImportVar(tag="refs")],
                    }
                )
            )
        )

    @property
    def set(self) -> Var:
        """Set the value of the client state variable.

        This property can only be attached to a frontend event trigger.

        To set a value from a backend event handler, see `push`.

        Returns:
            A special EventChain Var which will set the value when triggered.
        """
        return (
            Var.create_safe(
                _client_state_ref(self._setter_name),
                _var_is_local=False,
                _var_is_string=False,
            )
            .to(EventChain)
            ._replace(
                merge_var_data=VarData(  # type: ignore
                    imports={
                        f"/{constants.Dirs.STATE_PATH}": [ImportVar(tag="refs")],
                    }
                )
            )
        )

    def retrieve(self, callback: EventHandler | Callable | None = None) -> EventSpec:
        """Pass the value of the client state variable to a backend EventHandler.

        The event handler must `yield` or `return` the EventSpec to trigger the event.

        Args:
            callback: The callback to pass the value to.

        Returns:
            An EventSpec which will retrieve the value when triggered.
        """
        return call_script(_client_state_ref(self._getter_name), callback=callback)

    def push(self, value: Any) -> EventSpec:
        """Push a value to the client state variable from the backend.

        The event handler must `yield` or `return` the EventSpec to trigger the event.

        Args:
            value: The value to update.

        Returns:
            An EventSpec which will push the value when triggered.
        """
        return call_script(f"{_client_state_ref(self._setter_name)}({value})")