"""A module to hold state proxy classes.""" from __future__ import annotations import asyncio import copy import dataclasses import functools import inspect import json from collections.abc import Callable, Sequence from types import MethodType from typing import TYPE_CHECKING, Any, SupportsIndex import pydantic import wrapt from pydantic import BaseModel as BaseModelV2 from pydantic.v1 import BaseModel as BaseModelV1 from sqlalchemy.orm import DeclarativeBase from reflex.base import Base from reflex.utils import prerequisites from reflex.utils.exceptions import ImmutableStateError from reflex.utils.serializers import serializer from reflex.vars.base import Var if TYPE_CHECKING: from reflex.state import BaseState, StateUpdate class StateProxy(wrapt.ObjectProxy): """Proxy of a state instance to control mutability of vars for a background task. Since a background task runs against a state instance without holding the state_manager lock for the token, the reference may become stale if the same state is modified by another event handler. The proxy object ensures that writes to the state are blocked unless explicitly entering a context which refreshes the state from state_manager and holds the lock for the token until exiting the context. After exiting the context, a StateUpdate may be emitted to the frontend to notify the client of the state change. A background task will be passed the `StateProxy` as `self`, so mutability can be safely performed inside an `async with self` block. class State(rx.State): counter: int = 0 @rx.event(background=True) async def bg_increment(self): await asyncio.sleep(1) async with self: self.counter += 1 """ def __init__( self, state_instance: BaseState, parent_state_proxy: StateProxy | None = None, ): """Create a proxy for a state instance. If `get_state` is used on a StateProxy, the resulting state will be linked to the given state via parent_state_proxy. The first state in the chain is the state that initiated the background task. Args: state_instance: The state instance to proxy. parent_state_proxy: The parent state proxy, for linked mutability and context tracking. """ super().__init__(state_instance) # compile is not relevant to backend logic self._self_app = prerequisites.get_and_validate_app().app self._self_substate_path = tuple(state_instance.get_full_name().split(".")) self._self_actx = None self._self_mutable = False self._self_actx_lock = asyncio.Lock() self._self_actx_lock_holder = None self._self_parent_state_proxy = parent_state_proxy def _is_mutable(self) -> bool: """Check if the state is mutable. Returns: Whether the state is mutable. """ if self._self_parent_state_proxy is not None: return self._self_parent_state_proxy._is_mutable() or self._self_mutable return self._self_mutable async def __aenter__(self) -> StateProxy: """Enter the async context manager protocol. Sets mutability to True and enters the `App.modify_state` async context, which refreshes the state from state_manager and holds the lock for the given state token until exiting the context. Background tasks should avoid blocking calls while inside the context. Returns: This StateProxy instance in mutable mode. Raises: ImmutableStateError: If the state is already mutable. """ if self._self_parent_state_proxy is not None: from reflex.state import State parent_state = ( await self._self_parent_state_proxy.__aenter__() ).__wrapped__ super().__setattr__( "__wrapped__", await parent_state.get_state( State.get_class_substate(self._self_substate_path) ), ) return self current_task = asyncio.current_task() if ( self._self_actx_lock.locked() and current_task == self._self_actx_lock_holder ): raise ImmutableStateError( "The state is already mutable. Do not nest `async with self` blocks." ) from reflex.state import _substate_key await self._self_actx_lock.acquire() self._self_actx_lock_holder = current_task self._self_actx = self._self_app.modify_state( token=_substate_key( self.__wrapped__.router.session.client_token, self._self_substate_path, ) ) mutable_state = await self._self_actx.__aenter__() super().__setattr__( "__wrapped__", mutable_state.get_substate(self._self_substate_path) ) self._self_mutable = True return self async def __aexit__(self, *exc_info: Any) -> None: """Exit the async context manager protocol. Sets proxy mutability to False and persists any state changes. Args: exc_info: The exception info tuple. """ if self._self_parent_state_proxy is not None: await self._self_parent_state_proxy.__aexit__(*exc_info) return if self._self_actx is None: return self._self_mutable = False try: await self._self_actx.__aexit__(*exc_info) finally: self._self_actx_lock_holder = None self._self_actx_lock.release() self._self_actx = None def __enter__(self): """Enter the regular context manager protocol. This is not supported for background tasks, and exists only to raise a more useful exception when the StateProxy is used incorrectly. Raises: TypeError: always, because only async contextmanager protocol is supported. """ raise TypeError("Background task must use `async with self` to modify state.") def __exit__(self, *exc_info: Any) -> None: """Exit the regular context manager protocol. Args: exc_info: The exception info tuple. """ pass def __getattr__(self, name: str) -> Any: """Get the attribute from the underlying state instance. Args: name: The name of the attribute. Returns: The value of the attribute. Raises: ImmutableStateError: If the state is not in mutable mode. """ if name in ["substates", "parent_state"] and not self._is_mutable(): raise ImmutableStateError( "Background task StateProxy is immutable outside of a context " "manager. Use `async with self` to modify state." ) value = super().__getattr__(name) if not name.startswith("_self_") and isinstance(value, MutableProxy): # ensure mutations to these containers are blocked unless proxy is _mutable return ImmutableMutableProxy( wrapped=value.__wrapped__, state=self, field_name=value._self_field_name, ) if isinstance(value, functools.partial) and value.args[0] is self.__wrapped__: # Rebind event handler to the proxy instance value = functools.partial( value.func, self, *value.args[1:], **value.keywords, ) if isinstance(value, MethodType) and value.__self__ is self.__wrapped__: # Rebind methods to the proxy instance value = type(value)(value.__func__, self) return value def __setattr__(self, name: str, value: Any) -> None: """Set the attribute on the underlying state instance. If the attribute is internal, set it on the proxy instance instead. Args: name: The name of the attribute. value: The value of the attribute. Raises: ImmutableStateError: If the state is not in mutable mode. """ if ( name.startswith("_self_") # wrapper attribute or self._is_mutable() # lock held # non-persisted state attribute or name in self.__wrapped__.get_skip_vars() ): super().__setattr__(name, value) return raise ImmutableStateError( "Background task StateProxy is immutable outside of a context " "manager. Use `async with self` to modify state." ) def get_substate(self, path: Sequence[str]) -> BaseState: """Only allow substate access with lock held. Args: path: The path to the substate. Returns: The substate. Raises: ImmutableStateError: If the state is not in mutable mode. """ if not self._is_mutable(): raise ImmutableStateError( "Background task StateProxy is immutable outside of a context " "manager. Use `async with self` to modify state." ) return self.__wrapped__.get_substate(path) async def get_state(self, state_cls: type[BaseState]) -> BaseState: """Get an instance of the state associated with this token. Args: state_cls: The class of the state. Returns: The state. Raises: ImmutableStateError: If the state is not in mutable mode. """ if not self._is_mutable(): raise ImmutableStateError( "Background task StateProxy is immutable outside of a context " "manager. Use `async with self` to modify state." ) return type(self)( await self.__wrapped__.get_state(state_cls), parent_state_proxy=self ) async def _as_state_update(self, *args, **kwargs) -> StateUpdate: """Temporarily allow mutability to access parent_state. Args: *args: The args to pass to the underlying state instance. **kwargs: The kwargs to pass to the underlying state instance. Returns: The state update. """ original_mutable = self._self_mutable self._self_mutable = True try: return await self.__wrapped__._as_state_update(*args, **kwargs) finally: self._self_mutable = original_mutable class ReadOnlyStateProxy(StateProxy): """A read-only proxy for a state.""" def __setattr__(self, name: str, value: Any) -> None: """Prevent setting attributes on the state for read-only proxy. Args: name: The attribute name. value: The attribute value. Raises: NotImplementedError: Always raised when trying to set an attribute on proxied state. """ if name.startswith("_self_"): # Special case attributes of the proxy itself, not applied to the wrapped object. super().__setattr__(name, value) return raise NotImplementedError("This is a read-only state proxy.") def mark_dirty(self): """Mark the state as dirty. Raises: NotImplementedError: Always raised when trying to mark the proxied state as dirty. """ raise NotImplementedError("This is a read-only state proxy.") class MutableProxy(wrapt.ObjectProxy): """A proxy for a mutable object that tracks changes.""" # Hint for finding the base class of the proxy. __base_proxy__ = "MutableProxy" # Methods on wrapped objects which should mark the state as dirty. __mark_dirty_attrs__ = { "add", "append", "clear", "difference_update", "discard", "extend", "insert", "intersection_update", "pop", "popitem", "remove", "reverse", "setdefault", "sort", "symmetric_difference_update", "update", } # Methods on wrapped objects might return mutable objects that should be tracked. __wrap_mutable_attrs__ = { "get", "setdefault", } # These internal attributes on rx.Base should NOT be wrapped in a MutableProxy. __never_wrap_base_attrs__ = set(Base.__dict__) - {"set"} | set( pydantic.BaseModel.__dict__ ) # These types will be wrapped in MutableProxy __mutable_types__ = ( list, dict, set, Base, DeclarativeBase, BaseModelV2, BaseModelV1, ) # Dynamically generated classes for tracking dataclass mutations. __dataclass_proxies__: dict[type, type] = {} def __new__(cls, wrapped: Any, *args, **kwargs) -> MutableProxy: """Create a proxy instance for a mutable object that tracks changes. Args: wrapped: The object to proxy. *args: Other args passed to MutableProxy (ignored). **kwargs: Other kwargs passed to MutableProxy (ignored). Returns: The proxy instance. """ if dataclasses.is_dataclass(wrapped): wrapped_cls = type(wrapped) wrapper_cls_name = wrapped_cls.__name__ + cls.__name__ # Find the associated class if wrapper_cls_name not in cls.__dataclass_proxies__: # Create a new class that has the __dataclass_fields__ defined cls.__dataclass_proxies__[wrapper_cls_name] = type( wrapper_cls_name, (cls,), { dataclasses._FIELDS: getattr( # pyright: ignore [reportAttributeAccessIssue] wrapped_cls, dataclasses._FIELDS, # pyright: ignore [reportAttributeAccessIssue] ), }, ) cls = cls.__dataclass_proxies__[wrapper_cls_name] return super().__new__(cls) def __init__(self, wrapped: Any, state: BaseState, field_name: str): """Create a proxy for a mutable object that tracks changes. Args: wrapped: The object to proxy. state: The state to mark dirty when the object is changed. field_name: The name of the field on the state associated with the wrapped object. """ super().__init__(wrapped) self._self_state = state self._self_field_name = field_name def __repr__(self) -> str: """Get the representation of the wrapped object. Returns: The representation of the wrapped object. """ return f"{type(self).__name__}({self.__wrapped__})" def _mark_dirty( self, wrapped: Callable | None = None, instance: BaseState | None = None, args: tuple = (), kwargs: dict | None = None, ) -> Any: """Mark the state as dirty, then call a wrapped function. Intended for use with `FunctionWrapper` from the `wrapt` library. Args: wrapped: The wrapped function. instance: The instance of the wrapped function. args: The args for the wrapped function. kwargs: The kwargs for the wrapped function. Returns: The result of the wrapped function. """ self._self_state.dirty_vars.add(self._self_field_name) self._self_state._mark_dirty() if wrapped is not None: return wrapped(*args, **(kwargs or {})) @classmethod def _is_mutable_type(cls, value: Any) -> bool: """Check if a value is of a mutable type and should be wrapped. Args: value: The value to check. Returns: Whether the value is of a mutable type. """ return isinstance(value, cls.__mutable_types__) or ( dataclasses.is_dataclass(value) and not isinstance(value, Var) ) @staticmethod def _is_called_from_dataclasses_internal() -> bool: """Check if the current function is called from dataclasses helper. Returns: Whether the current function is called from dataclasses internal code. """ # Walk up the stack a bit to see if we are called from dataclasses # internal code, for example `asdict` or `astuple`. frame = inspect.currentframe() for _ in range(5): # Why not `inspect.stack()` -- this is much faster! if not (frame := frame and frame.f_back): break if inspect.getfile(frame) == dataclasses.__file__: return True return False def _wrap_recursive(self, value: Any) -> Any: """Wrap a value recursively if it is mutable. Args: value: The value to wrap. Returns: The wrapped value. """ # When called from dataclasses internal code, return the unwrapped value if self._is_called_from_dataclasses_internal(): return value # Recursively wrap mutable types, but do not re-wrap MutableProxy instances. if self._is_mutable_type(value) and not isinstance(value, MutableProxy): base_cls = globals()[self.__base_proxy__] return base_cls( wrapped=value, state=self._self_state, field_name=self._self_field_name, ) return value def _wrap_recursive_decorator( self, wrapped: Callable, instance: BaseState, args: list, kwargs: dict ) -> Any: """Wrap a function that returns a possibly mutable value. Intended for use with `FunctionWrapper` from the `wrapt` library. Args: wrapped: The wrapped function. instance: The instance of the wrapped function. args: The args for the wrapped function. kwargs: The kwargs for the wrapped function. Returns: The result of the wrapped function (possibly wrapped in a MutableProxy). """ return self._wrap_recursive(wrapped(*args, **kwargs)) def __getattr__(self, __name: str) -> Any: """Get the attribute on the proxied object and return a proxy if mutable. Args: __name: The name of the attribute. Returns: The attribute value. """ value = super().__getattr__(__name) if callable(value): if __name in self.__mark_dirty_attrs__: # Wrap special callables, like "append", which should mark state dirty. value = wrapt.FunctionWrapper(value, self._mark_dirty) if __name in self.__wrap_mutable_attrs__: # Wrap methods that may return mutable objects tied to the state. value = wrapt.FunctionWrapper( value, self._wrap_recursive_decorator, ) if ( isinstance(self.__wrapped__, Base) and __name not in self.__never_wrap_base_attrs__ and hasattr(value, "__func__") ): # Wrap methods called on Base subclasses, which might do _anything_ return wrapt.FunctionWrapper( functools.partial(value.__func__, self), # pyright: ignore [reportFunctionMemberAccess] self._wrap_recursive_decorator, ) if self._is_mutable_type(value) and __name not in ( "__wrapped__", "_self_state", "__dict__", ): # Recursively wrap mutable attribute values retrieved through this proxy. return self._wrap_recursive(value) return value def __getitem__(self, key: Any) -> Any: """Get the item on the proxied object and return a proxy if mutable. Args: key: The key of the item. Returns: The item value. """ value = super().__getitem__(key) if isinstance(key, slice) and isinstance(value, list): return [self._wrap_recursive(item) for item in value] # Recursively wrap mutable items retrieved through this proxy. return self._wrap_recursive(value) def __iter__(self) -> Any: """Iterate over the proxied object and return a proxy if mutable. Yields: Each item value (possibly wrapped in MutableProxy). """ for value in super().__iter__(): # Recursively wrap mutable items retrieved through this proxy. yield self._wrap_recursive(value) def __delattr__(self, name: str): """Delete the attribute on the proxied object and mark state dirty. Args: name: The name of the attribute. """ self._mark_dirty(super().__delattr__, args=(name,)) def __delitem__(self, key: str): """Delete the item on the proxied object and mark state dirty. Args: key: The key of the item. """ self._mark_dirty(super().__delitem__, args=(key,)) def __setitem__(self, key: str, value: Any): """Set the item on the proxied object and mark state dirty. Args: key: The key of the item. value: The value of the item. """ self._mark_dirty(super().__setitem__, args=(key, value)) def __setattr__(self, name: str, value: Any): """Set the attribute on the proxied object and mark state dirty. If the attribute starts with "_self_", then the state is NOT marked dirty as these are internal proxy attributes. Args: name: The name of the attribute. value: The value of the attribute. """ if name.startswith("_self_"): # Special case attributes of the proxy itself, not applied to the wrapped object. super().__setattr__(name, value) return self._mark_dirty(super().__setattr__, args=(name, value)) def __copy__(self) -> Any: """Return a copy of the proxy. Returns: A copy of the wrapped object, unconnected to the proxy. """ return copy.copy(self.__wrapped__) def __deepcopy__(self, memo: dict[int, Any] | None = None) -> Any: """Return a deepcopy of the proxy. Args: memo: The memo dict to use for the deepcopy. Returns: A deepcopy of the wrapped object, unconnected to the proxy. """ return copy.deepcopy(self.__wrapped__, memo=memo) def __reduce_ex__(self, protocol_version: SupportsIndex): """Get the state for redis serialization. This method is called by cloudpickle to serialize the object. It explicitly serializes the wrapped object, stripping off the mutable proxy. Args: protocol_version: The protocol version. Returns: Tuple of (wrapped class, empty args, class __getstate__) """ return self.__wrapped__.__reduce_ex__(protocol_version) @serializer def serialize_mutable_proxy(mp: MutableProxy): """Return the wrapped value of a MutableProxy. Args: mp: The MutableProxy to serialize. Returns: The wrapped object. """ return mp.__wrapped__ _orig_json_encoder_default = json.JSONEncoder.default def _json_encoder_default_wrapper(self: json.JSONEncoder, o: Any) -> Any: """Wrap JSONEncoder.default to handle MutableProxy objects. Args: self: the JSONEncoder instance. o: the object to serialize. Returns: A JSON-able object. """ try: return o.__wrapped__ except AttributeError: pass return _orig_json_encoder_default(self, o) json.JSONEncoder.default = _json_encoder_default_wrapper class ImmutableMutableProxy(MutableProxy): """A proxy for a mutable object that tracks changes. This wrapper comes from StateProxy, and will raise an exception if an attempt is made to modify the wrapped object when the StateProxy is immutable. """ # Ensure that recursively wrapped proxies use ImmutableMutableProxy as base. __base_proxy__ = "ImmutableMutableProxy" def _mark_dirty( self, wrapped: Callable | None = None, instance: BaseState | None = None, args: tuple = (), kwargs: dict | None = None, ) -> Any: """Raise an exception when an attempt is made to modify the object. Intended for use with `FunctionWrapper` from the `wrapt` library. Args: wrapped: The wrapped function. instance: The instance of the wrapped function. args: The args for the wrapped function. kwargs: The kwargs for the wrapped function. Returns: The result of the wrapped function. Raises: ImmutableStateError: if the StateProxy is not mutable. """ if not self._self_state._is_mutable(): raise ImmutableStateError( "Background task StateProxy is immutable outside of a context " "manager. Use `async with self` to modify state." ) return super()._mark_dirty( wrapped=wrapped, instance=instance, args=args, kwargs=kwargs )