reflex 0.5.10a3__py3-none-any.whl → 0.6.0a1__py3-none-any.whl

reflex/state.py

@@ -11,6 +11,7 @@ import os
 import uuid
 from abc import ABC, abstractmethod
 from collections import defaultdict
+from pathlib import Path
 from types import FunctionType, MethodType
 from typing import (
     TYPE_CHECKING,
@@ -23,6 +24,7 @@ from typing import (
     Optional,
     Sequence,
     Set,
+    Tuple,
     Type,
     Union,
     cast,
@@ -32,6 +34,13 @@ import dill
 from sqlalchemy.orm import DeclarativeBase
 
 from reflex.config import get_config
+from reflex.ivars.base import (
+    DynamicRouteVar,
+    ImmutableComputedVar,
+    ImmutableVar,
+    immutable_computed_var,
+    is_computed_var,
+)
 
 try:
     import pydantic.v1 as pydantic
@@ -51,19 +60,23 @@ from reflex.event import (
     EventSpec,
     fix_events,
 )
-from reflex.utils import console, format, prerequisites, types
-from reflex.utils.exceptions import ImmutableStateError, LockExpiredError
+from reflex.utils import console, format, path_ops, prerequisites, types
+from reflex.utils.exceptions import (
+    DynamicRouteArgShadowsStateVar,
+    ImmutableStateError,
+    LockExpiredError,
+)
 from reflex.utils.exec import is_testing_env
 from reflex.utils.serializers import SerializedType, serialize, serializer
 from reflex.utils.types import override
-from reflex.vars import BaseVar, ComputedVar, Var, computed_var
+from reflex.vars import VarData
 
 if TYPE_CHECKING:
     from reflex.components.component import Component
 
 
 Delta = Dict[str, Any]
-var = computed_var
+var = immutable_computed_var
 
 
 # If the state is this large, it's considered a performance issue.
@@ -296,16 +309,16 @@ class BaseState(Base, ABC, extra=pydantic.Extra.allow):
     """The state of the app."""
 
     # A map from the var name to the var.
-    vars: ClassVar[Dict[str, Var]] = {}
+    vars: ClassVar[Dict[str, ImmutableVar]] = {}
 
     # The base vars of the class.
-    base_vars: ClassVar[Dict[str, BaseVar]] = {}
+    base_vars: ClassVar[Dict[str, ImmutableVar]] = {}
 
     # The computed vars of the class.
-    computed_vars: ClassVar[Dict[str, ComputedVar]] = {}
+    computed_vars: ClassVar[Dict[str, ImmutableComputedVar]] = {}
 
     # Vars inherited by the parent state.
-    inherited_vars: ClassVar[Dict[str, Var]] = {}
+    inherited_vars: ClassVar[Dict[str, ImmutableVar]] = {}
 
     # Backend base vars that are never sent to the client.
     backend_vars: ClassVar[Dict[str, Any]] = {}
@@ -415,7 +428,7 @@ class BaseState(Base, ABC, extra=pydantic.Extra.allow):
         return f"{self.__class__.__name__}({self.dict()})"
 
     @classmethod
-    def _get_computed_vars(cls) -> list[ComputedVar]:
+    def _get_computed_vars(cls) -> list[ImmutableComputedVar]:
         """Helper function to get all computed vars of a instance.
 
         Returns:
@@ -424,8 +437,8 @@ class BaseState(Base, ABC, extra=pydantic.Extra.allow):
         return [
             v
             for mixin in cls._mixins() + [cls]
-            for v in mixin.__dict__.values()
-            if isinstance(v, ComputedVar)
+            for name, v in mixin.__dict__.items()
+            if is_computed_var(v) and name not in cls.inherited_vars
         ]
 
     @classmethod
@@ -469,7 +482,7 @@ class BaseState(Base, ABC, extra=pydantic.Extra.allow):
         cls._check_overridden_methods()
 
         # Computed vars should not shadow builtin state props.
-        cls._check_overriden_basevars()
+        cls._check_overridden_basevars()
 
         # Reset subclass tracking for this class.
         cls.class_subclasses = set()
@@ -487,26 +500,17 @@ class BaseState(Base, ABC, extra=pydantic.Extra.allow):
             if cls.get_name() in set(
                 c.get_name() for c in parent_state.class_subclasses
             ):
-                if is_testing_env():
-                    # Clear existing subclass with same name when app is reloaded via
-                    # utils.prerequisites.get_app(reload=True)
-                    parent_state.class_subclasses = set(
-                        c
-                        for c in parent_state.class_subclasses
-                        if c.get_name() != cls.get_name()
-                    )
-                else:
-                    # During normal operation, subclasses cannot have the same name, even if they are
-                    # defined in different modules.
-                    raise StateValueError(
-                        f"The substate class '{cls.get_name()}' has been defined multiple times. "
-                        "Shadowing substate classes is not allowed."
-                    )
+                # This should not happen, since we have added module prefix to state names in #3214
+                raise StateValueError(
+                    f"The substate class '{cls.get_name()}' has been defined multiple times. "
+                    "Shadowing substate classes is not allowed."
+                )
             # Track this new subclass in the parent state's subclasses set.
             parent_state.class_subclasses.add(cls)
 
         # Get computed vars.
         computed_vars = cls._get_computed_vars()
+        cls._check_overridden_computed_vars()
 
         new_backend_vars = {
             name: value
@@ -521,13 +525,18 @@ class BaseState(Base, ABC, extra=pydantic.Extra.allow):
 
         # Set the base and computed vars.
         cls.base_vars = {
-            f.name: BaseVar(_var_name=f.name, _var_type=f.outer_type_)._var_set_state(
-                cls
-            )
+            f.name: ImmutableVar(
+                _var_name=format.format_state_name(cls.get_full_name()) + "." + f.name,
+                _var_type=f.outer_type_,
+                _var_data=VarData.from_state(cls),
+            ).guess_type()
             for f in cls.get_fields().values()
             if f.name not in cls.get_skip_vars()
         }
-        cls.computed_vars = {v._var_name: v._var_set_state(cls) for v in computed_vars}
+        cls.computed_vars = {
+            v._var_name: v._replace(merge_var_data=VarData.from_state(cls))
+            for v in computed_vars
+        }
         cls.vars = {
             **cls.inherited_vars,
             **cls.base_vars,
@@ -548,12 +557,12 @@ class BaseState(Base, ABC, extra=pydantic.Extra.allow):
 
         for mixin in cls._mixins():
             for name, value in mixin.__dict__.items():
-                if isinstance(value, ComputedVar):
+                if name in cls.inherited_vars:
+                    continue
+                if is_computed_var(value):
                     fget = cls._copy_fn(value.fget)
-                    newcv = value._replace(fget=fget)
+                    newcv = value._replace(fget=fget, _var_data=VarData.from_state(cls))
                     # cleanup refs to mixin cls in var_data
-                    newcv._var_data = None
-                    newcv._var_set_state(cls)
                     setattr(cls, name, newcv)
                     cls.computed_vars[newcv._var_name] = newcv
                     cls.vars[newcv._var_name] = newcv
@@ -579,6 +588,9 @@ class BaseState(Base, ABC, extra=pydantic.Extra.allow):
             cls.event_handlers[name] = handler
             setattr(cls, name, handler)
 
+        # Initialize per-class var dependency tracking.
+        cls._computed_var_dependencies = defaultdict(set)
+        cls._substate_var_dependencies = defaultdict(set)
         cls._init_var_dependency_dicts()
 
     @staticmethod
@@ -648,10 +660,6 @@ class BaseState(Base, ABC, extra=pydantic.Extra.allow):
         Additional updates tracking dicts for vars and substates that always
         need to be recomputed.
         """
-        # Initialize per-class var dependency tracking.
-        cls._computed_var_dependencies = defaultdict(set)
-        cls._substate_var_dependencies = defaultdict(set)
-
         inherited_vars = set(cls.inherited_vars).union(
             set(cls.inherited_backend_vars),
         )
@@ -716,7 +724,7 @@ class BaseState(Base, ABC, extra=pydantic.Extra.allow):
             )
 
     @classmethod
-    def _check_overriden_basevars(cls):
+    def _check_overridden_basevars(cls):
         """Check for shadow base vars and raise error if any.
 
         Raises:
@@ -728,6 +736,22 @@ class BaseState(Base, ABC, extra=pydantic.Extra.allow):
                     f"The computed var name `{computed_var_._var_name}` shadows a base var in {cls.__module__}.{cls.__name__}; use a different name instead"
                 )
 
+    @classmethod
+    def _check_overridden_computed_vars(cls) -> None:
+        """Check for shadow computed vars and raise error if any.
+
+        Raises:
+            NameError: When a computed var shadows another.
+        """
+        for name, cv in cls.__dict__.items():
+            if not is_computed_var(cv):
+                continue
+            name = cv._var_name
+            if name in cls.inherited_vars or name in cls.inherited_backend_vars:
+                raise NameError(
+                    f"The computed var name `{cv._var_name}` shadows a var in {cls.__module__}.{cls.__name__}; use a different name instead"
+                )
+
     @classmethod
     def get_skip_vars(cls) -> set[str]:
         """Get the vars to skip when serializing.
@@ -847,7 +871,7 @@ class BaseState(Base, ABC, extra=pydantic.Extra.allow):
         return getattr(substate, name)
 
     @classmethod
-    def _init_var(cls, prop: BaseVar):
+    def _init_var(cls, prop: ImmutableVar):
         """Initialize a variable.
 
         Args:
@@ -890,8 +914,11 @@ class BaseState(Base, ABC, extra=pydantic.Extra.allow):
             )
 
         # create the variable based on name and type
-        var = BaseVar(_var_name=name, _var_type=type_)
-        var._var_set_state(cls)
+        var = ImmutableVar(
+            _var_name=format.format_state_name(cls.get_full_name()) + "." + name,
+            _var_type=type_,
+            _var_data=VarData.from_state(cls),
+        ).guess_type()
 
         # add the pydantic field dynamically (must be done before _init_var)
         cls.add_field(var, default_value)
@@ -910,13 +937,18 @@ class BaseState(Base, ABC, extra=pydantic.Extra.allow):
         cls._init_var_dependency_dicts()
 
     @classmethod
-    def _set_var(cls, prop: BaseVar):
+    def _set_var(cls, prop: ImmutableVar):
         """Set the var as a class member.
 
         Args:
             prop: The var instance to set.
         """
-        setattr(cls, prop._var_name, prop)
+        acutal_var_name = (
+            prop._var_name
+            if "." not in prop._var_name
+            else prop._var_name.split(".")[-1]
+        )
+        setattr(cls, acutal_var_name, prop)
 
     @classmethod
     def _create_event_handler(cls, fn):
@@ -936,7 +968,7 @@ class BaseState(Base, ABC, extra=pydantic.Extra.allow):
         cls.setvar = cls.event_handlers["setvar"] = EventHandlerSetVar(state_cls=cls)
 
     @classmethod
-    def _create_setter(cls, prop: BaseVar):
+    def _create_setter(cls, prop: ImmutableVar):
         """Create a setter for the var.
 
         Args:
@@ -949,14 +981,17 @@ class BaseState(Base, ABC, extra=pydantic.Extra.allow):
             setattr(cls, setter_name, event_handler)
 
     @classmethod
-    def _set_default_value(cls, prop: BaseVar):
+    def _set_default_value(cls, prop: ImmutableVar):
         """Set the default value for the var.
 
         Args:
             prop: The var to set the default value for.
         """
         # Get the pydantic field for the var.
-        field = cls.get_fields()[prop._var_name]
+        if "." in prop._var_name:
+            field = cls.get_fields()[prop._var_name.split(".")[-1]]
+        else:
+            field = cls.get_fields()[prop._var_name]
         if field.required:
             default_value = prop.get_default_value()
             if default_value is not None:
@@ -968,7 +1003,7 @@ class BaseState(Base, ABC, extra=pydantic.Extra.allow):
             and not types.is_optional(prop._var_type)
         ):
             # Ensure frontend uses null coalescing when accessing.
-            prop._var_type = Optional[prop._var_type]
+            object.__setattr__(prop, "_var_type", Optional[prop._var_type])
 
     @staticmethod
     def _get_base_functions() -> dict[str, FunctionType]:
@@ -990,20 +1025,22 @@ class BaseState(Base, ABC, extra=pydantic.Extra.allow):
         Args:
             args: a dict of args
         """
+        if not args:
+            return
+
+        cls._check_overwritten_dynamic_args(list(args.keys()))
 
         def argsingle_factory(param):
-            @ComputedVar
             def inner_func(self) -> str:
                 return self.router.page.params.get(param, "")
 
-            return inner_func
+            return DynamicRouteVar(fget=inner_func, cache=True)
 
         def arglist_factory(param):
-            @ComputedVar
-            def inner_func(self) -> List:
+            def inner_func(self) -> List[str]:
                 return self.router.page.params.get(param, [])
 
-            return inner_func
+            return DynamicRouteVar(fget=inner_func, cache=True)
 
         for param, value in args.items():
             if value == constants.RouteArgType.SINGLE:
@@ -1012,13 +1049,37 @@ class BaseState(Base, ABC, extra=pydantic.Extra.allow):
                 func = arglist_factory(param)
             else:
                 continue
-            # to allow passing as a prop
-            func._var_name = param
-            cls.vars[param] = cls.computed_vars[param] = func._var_set_state(cls)  # type: ignore
+            # to allow passing as a prop, evade python frozen rules (bad practice)
+            object.__setattr__(func, "_var_name", param)
+            # cls.vars[param] = cls.computed_vars[param] = func._var_set_state(cls)  # type: ignore
+            cls.vars[param] = cls.computed_vars[param] = func._replace(
+                _var_data=VarData.from_state(cls)
+            )
             setattr(cls, param, func)
 
-            # Reinitialize dependency tracking dicts.
-            cls._init_var_dependency_dicts()
+        # Reinitialize dependency tracking dicts.
+        cls._init_var_dependency_dicts()
+
+    @classmethod
+    def _check_overwritten_dynamic_args(cls, args: list[str]):
+        """Check if dynamic args are shadowing existing vars. Recursively checks all child states.
+
+        Args:
+            args: a dict of args
+
+        Raises:
+            DynamicRouteArgShadowsStateVar: If a dynamic arg is shadowing an existing var.
+        """
+        for arg in args:
+            if (
+                arg in cls.computed_vars
+                and not isinstance(cls.computed_vars[arg], DynamicRouteVar)
+            ) or arg in cls.base_vars:
+                raise DynamicRouteArgShadowsStateVar(
+                    f"Dynamic route arg '{arg}' is shadowing an existing var in {cls.__module__}.{cls.__name__}"
+                )
+        for substate in cls.get_substates():
+            substate._check_overwritten_dynamic_args(args)
 
     def __getattribute__(self, name: str) -> Any:
         """Get the state var.
@@ -1769,12 +1830,12 @@ class BaseState(Base, ABC, extra=pydantic.Extra.allow):
             prop_name: self.get_value(getattr(self, prop_name))
             for prop_name in self.base_vars
         }
-        if initial:
+        if initial and include_computed:
             computed_vars = {
                 # Include initial computed vars.
                 prop_name: (
                     cv._initial_value
-                    if isinstance(cv, ComputedVar)
+                    if is_computed_var(cv)
                     and not isinstance(cv._initial_value, types.Unset)
                     else self.get_value(getattr(self, prop_name))
                 )
@@ -2272,11 +2333,12 @@ class StateProxy(wrapt.ObjectProxy):
         Returns:
             The state update.
         """
+        original_mutable = self._self_mutable
         self._self_mutable = True
         try:
             return self.__wrapped__._as_state_update(*args, **kwargs)
         finally:
-            self._self_mutable = False
+            self._self_mutable = original_mutable
 
 
 class StateUpdate(Base):
@@ -2318,7 +2380,7 @@ class StateManager(Base, ABC):
                 token_expiration=config.redis_token_expiration,
                 lock_expiration=config.redis_lock_expiration,
             )
-        return StateManagerMemory(state=state)
+        return StateManagerDisk(state=state)
 
     @abstractmethod
     async def get_state(self, token: str) -> BaseState:
@@ -2425,6 +2487,266 @@ class StateManagerMemory(StateManager):
             await self.set_state(token, state)
 
 
+def _default_token_expiration() -> int:
+    """Get the default token expiration time.
+
+    Returns:
+        The default token expiration time.
+    """
+    return get_config().redis_token_expiration
+
+
+def _serialize_type(type_: Any) -> str:
+    """Serialize a type.
+
+    Args:
+        type_: The type to serialize.
+
+    Returns:
+        The serialized type.
+    """
+    if not inspect.isclass(type_):
+        return f"{type_}"
+    return f"{type_.__module__}.{type_.__qualname__}"
+
+
+def state_to_schema(
+    state: BaseState,
+) -> List[
+    Tuple[
+        str,
+        str,
+        Any,
+        Union[bool, None],
+    ]
+]:
+    """Convert a state to a schema.
+
+    Args:
+        state: The state to convert to a schema.
+
+    Returns:
+        The schema.
+    """
+    return list(
+        sorted(
+            (
+                field_name,
+                model_field.name,
+                _serialize_type(model_field.type_),
+                (
+                    model_field.required
+                    if isinstance(model_field.required, bool)
+                    else None
+                ),
+            )
+            for field_name, model_field in state.__fields__.items()
+        )
+    )
+
+
+def reset_disk_state_manager():
+    """Reset the disk state manager."""
+    states_directory = prerequisites.get_web_dir() / constants.Dirs.STATES
+    if states_directory.exists():
+        for path in states_directory.iterdir():
+            path.unlink()
+
+
+class StateManagerDisk(StateManager):
+    """A state manager that stores states in memory."""
+
+    # The mapping of client ids to states.
+    states: Dict[str, BaseState] = {}
+
+    # The mutex ensures the dict of mutexes is updated exclusively
+    _state_manager_lock = asyncio.Lock()
+
+    # The dict of mutexes for each client
+    _states_locks: Dict[str, asyncio.Lock] = pydantic.PrivateAttr({})
+
+    # The token expiration time (s).
+    token_expiration: int = pydantic.Field(default_factory=_default_token_expiration)
+
+    class Config:
+        """The Pydantic config."""
+
+        fields = {
+            "_states_locks": {"exclude": True},
+        }
+        keep_untouched = (functools.cached_property,)
+
+    def __init__(self, state: Type[BaseState]):
+        """Create a new state manager.
+
+        Args:
+            state: The state class to use.
+        """
+        super().__init__(state=state)
+
+        path_ops.mkdir(self.states_directory)
+
+        self._purge_expired_states()
+
+    @functools.cached_property
+    def states_directory(self) -> Path:
+        """Get the states directory.
+
+        Returns:
+            The states directory.
+        """
+        return prerequisites.get_web_dir() / constants.Dirs.STATES
+
+    def _purge_expired_states(self):
+        """Purge expired states from the disk."""
+        import time
+
+        for path in path_ops.ls(self.states_directory):
+            # check path is a pickle file
+            if path.suffix != ".pkl":
+                continue
+
+            # load last edited field from file
+            last_edited = path.stat().st_mtime
+
+            # check if the file is older than the token expiration time
+            if time.time() - last_edited > self.token_expiration:
+                # remove the file
+                path.unlink()
+
+    def token_path(self, token: str) -> Path:
+        """Get the path for a token.
+
+        Args:
+            token: The token to get the path for.
+
+        Returns:
+            The path for the token.
+        """
+        return (self.states_directory / f"{token}.pkl").absolute()
+
+    async def load_state(self, token: str, root_state: BaseState) -> BaseState:
+        """Load a state object based on the provided token.
+
+        Args:
+            token: The token used to identify the state object.
+            root_state: The root state object.
+
+        Returns:
+            The loaded state object.
+        """
+        if token in self.states:
+            return self.states[token]
+
+        client_token, substate_address = _split_substate_key(token)
+
+        token_path = self.token_path(token)
+
+        if token_path.exists():
+            try:
+                with token_path.open(mode="rb") as file:
+                    (substate_schema, substate) = dill.load(file)
+                if substate_schema == state_to_schema(substate):
+                    await self.populate_substates(client_token, substate, root_state)
+                    return substate
+            except Exception:
+                pass
+
+        return root_state.get_substate(substate_address.split(".")[1:])
+
+    async def populate_substates(
+        self, client_token: str, state: BaseState, root_state: BaseState
+    ):
+        """Populate the substates of a state object.
+
+        Args:
+            client_token: The client token.
+            state: The state object to populate.
+            root_state: The root state object.
+        """
+        for substate in state.get_substates():
+            substate_token = _substate_key(client_token, substate)
+
+            substate = await self.load_state(substate_token, root_state)
+
+            state.substates[substate.get_name()] = substate
+            substate.parent_state = state
+
+    @override
+    async def get_state(
+        self,
+        token: str,
+    ) -> BaseState:
+        """Get the state for a token.
+
+        Args:
+            token: The token to get the state for.
+
+        Returns:
+            The state for the token.
+        """
+        client_token, substate_address = _split_substate_key(token)
+
+        root_state_token = _substate_key(client_token, substate_address.split(".")[0])
+
+        return await self.load_state(
+            root_state_token, self.state(_reflex_internal_init=True)
+        )
+
+    async def set_state_for_substate(self, client_token: str, substate: BaseState):
+        """Set the state for a substate.
+
+        Args:
+            client_token: The client token.
+            substate: The substate to set.
+        """
+        substate_token = _substate_key(client_token, substate)
+
+        self.states[substate_token] = substate
+
+        state_dilled = dill.dumps((state_to_schema(substate), substate))
+        if not self.states_directory.exists():
+            self.states_directory.mkdir(parents=True, exist_ok=True)
+        self.token_path(substate_token).write_bytes(state_dilled)
+
+        for substate_substate in substate.substates.values():
+            await self.set_state_for_substate(client_token, substate_substate)
+
+    @override
+    async def set_state(self, token: str, state: BaseState):
+        """Set the state for a token.
+
+        Args:
+            token: The token to set the state for.
+            state: The state to set.
+        """
+        client_token, substate = _split_substate_key(token)
+        await self.set_state_for_substate(client_token, state)
+
+    @override
+    @contextlib.asynccontextmanager
+    async def modify_state(self, token: str) -> AsyncIterator[BaseState]:
+        """Modify the state for a token while holding exclusive lock.
+
+        Args:
+            token: The token to modify the state for.
+
+        Yields:
+            The state for the token.
+        """
+        # Memory state manager ignores the substate suffix and always returns the top-level state.
+        client_token, substate = _split_substate_key(token)
+        if client_token not in self._states_locks:
+            async with self._state_manager_lock:
+                if client_token not in self._states_locks:
+                    self._states_locks[client_token] = asyncio.Lock()
+
+        async with self._states_locks[client_token]:
+            state = await self.get_state(token)
+            yield state
+            await self.set_state(token, state)
+
+
 # Workaround https://github.com/cloudpipe/cloudpickle/issues/408 for dynamic pydantic classes
 if not isinstance(State.validate.__func__, FunctionType):
     cython_function_or_method = type(State.validate.__func__)
@@ -2453,15 +2775,6 @@ def _default_lock_expiration() -> int:
     return get_config().redis_lock_expiration
 
 
-def _default_token_expiration() -> int:
-    """Get the default token expiration time.
-
-    Returns:
-        The default token expiration time.
-    """
-    return get_config().redis_token_expiration
-
-
 class StateManagerRedis(StateManager):
     """A state manager that stores states in redis."""
 
@@ -3340,5 +3653,7 @@ def reload_state_module(
         if subclass.__module__ == module and module is not None:
             state.class_subclasses.remove(subclass)
             state._always_dirty_substates.discard(subclass.get_name())
-    state._init_var_dependency_dicts()
+            state._computed_var_dependencies = defaultdict(set)
+            state._substate_var_dependencies = defaultdict(set)
+            state._init_var_dependency_dicts()
     state.get_class_substate.cache_clear()