From 0ed7c5d96994337a7e2dbfde1b0e741d9bf108ed Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Thomas=20Brand=C3=A9ho?= <thomas.brandeho@gmail.com>
Date: Tue, 5 Nov 2024 07:21:59 -0800
Subject: [PATCH] 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        |  1 +
 reflex/__init__.pyi       |  1 +
 reflex/istate/__init__.py |  1 +
 reflex/istate/proxy.py    | 33 +++++++++++++++++++++++++++++++++
 reflex/istate/wrappers.py | 31 +++++++++++++++++++++++++++++++
 5 files changed, 67 insertions(+)
 create mode 100644 reflex/istate/__init__.py
 create mode 100644 reflex/istate/proxy.py
 create mode 100644 reflex/istate/wrappers.py

diff --git a/reflex/__init__.py b/reflex/__init__.py
index acba5936a..cfb971a99 100644
--- a/reflex/__init__.py
+++ b/reflex/__init__.py
@@ -336,6 +336,7 @@ _MAPPING: dict = {
         "State",
         "dynamic",
     ],
+    "istate.wrappers": ["get_state"],
     "style": ["Style", "toggle_color_mode"],
     "utils.imports": ["ImportVar"],
     "utils.serializers": ["serializer"],
diff --git a/reflex/__init__.pyi b/reflex/__init__.pyi
index a0f60ea57..2d22fe497 100644
--- a/reflex/__init__.pyi
+++ b/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
diff --git a/reflex/istate/__init__.py b/reflex/istate/__init__.py
new file mode 100644
index 000000000..d71d038f8
--- /dev/null
+++ b/reflex/istate/__init__.py
@@ -0,0 +1 @@
+"""This module will provide interfaces for the state."""
diff --git a/reflex/istate/proxy.py b/reflex/istate/proxy.py
new file mode 100644
index 000000000..8d6051cf2
--- /dev/null
+++ b/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.")
diff --git a/reflex/istate/wrappers.py b/reflex/istate/wrappers.py
new file mode 100644
index 000000000..7f010eb9e
--- /dev/null
+++ b/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)