expose rx.get_state() to get instance of state from anywhere (#3959)

* expose rx.get_state() to get instance of state from anywhere

* fix circular import and add read-only proxy

reflex/__init__.py

@@ -336,6 +336,7 @@
         "State",
         "dynamic",
     ],
+    "istate.wrappers": ["get_state"],
     "style": ["Style", "toggle_color_mode"],
     "utils.imports": ["ImportVar"],
     "utils.serializers": ["serializer"],

reflex/__init__.pyi

@@ -180,6 +180,7 @@ from .experimental import _x as _x
 from .istate.storage import Cookie as Cookie
 from .istate.storage import LocalStorage as LocalStorage
 from .istate.storage import SessionStorage as SessionStorage
+from .istate.wrappers import get_state as get_state
 from .middleware import Middleware as Middleware
 from .middleware import middleware as middleware
 from .model import Model as Model

reflex/istate/__init__.py

@@ -0,0 +1 @@
+"""This module will provide interfaces for the state."""

reflex/istate/proxy.py

@@ -0,0 +1,33 @@
+"""A module to hold state proxy classes."""
+
+from typing import Any
+
+from reflex.state import StateProxy
+
+
+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.")

reflex/istate/wrappers.py

@@ -0,0 +1,31 @@
+"""Wrappers for the state manager."""
+
+from typing import Any
+
+from reflex.istate.proxy import ReadOnlyStateProxy
+from reflex.state import (
+    _split_substate_key,
+    _substate_key,
+    get_state_manager,
+)
+
+
+async def get_state(token, state_cls: Any | None = None) -> ReadOnlyStateProxy:
+    """Get the instance of a state for a token.
+
+    Args:
+        token: The token for the state.
+        state_cls: The class of the state.
+
+    Returns:
+        A read-only proxy of the state instance.
+    """
+    mng = get_state_manager()
+    if state_cls is not None:
+        root_state = await mng.get_state(_substate_key(token, state_cls))
+    else:
+        root_state = await mng.get_state(token)
+        _, state_path = _split_substate_key(token)
+        state_cls = root_state.get_class_substate(tuple(state_path.split(".")))
+    instance = await root_state.get_state(state_cls)
+    return ReadOnlyStateProxy(instance)