support async context for mutable proxies

news/6689.feature.md

@@ -0,0 +1 @@
+Support using mutable state proxies as async context managers.

reflex/istate/proxy.py

@@ -12,7 +12,7 @@
 from collections.abc import Callable, Sequence
 from importlib.util import find_spec
 from types import MethodType
-from typing import TYPE_CHECKING, Any, SupportsIndex, TypeVar
+from typing import TYPE_CHECKING, Any, NoReturn, SupportsIndex, TypeVar
 
 import wrapt
 from reflex_base.event import Event
@@ -446,18 +446,27 @@ def __new__(cls, wrapped: Any, *args, **kwargs) -> MutableProxy:
             cls = cls.__dataclass_proxies__[wrapper_cls_name]
         return super().__new__(cls)  # pyright: ignore[reportArgumentType]
 
-    def __init__(self, wrapped: Any, state: BaseState, field_name: str):
+    def __init__(
+        self,
+        wrapped: Any,
+        state: BaseState,
+        field_name: str,
+        path: tuple[tuple[str, Any], ...] | None = None,
+    ):
         """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.
+            path: Access path from the state field to this wrapped object.
         """
         super().__init__(wrapped)
         self._self_state = state
         self._self_field_name = field_name
+        self._self_path = path or ()
+        self._self_actx_state = None
 
     def __repr__(self) -> str:
         """Get the representation of the wrapped object.
@@ -467,6 +476,56 @@ def __repr__(self) -> str:
         """
         return f"{type(self).__name__}({self.__wrapped__})"
 
+    async def __aenter__(self) -> Self:
+        """Enter the async context manager protocol through the bound state.
+
+        Returns:
+            This proxy refreshed from the current state field.
+        """
+        context_state = self._self_state
+        self._self_actx_state = context_state
+        state = await context_state.__aenter__()
+        try:
+            refreshed_value = getattr(state, self._self_field_name)
+            for access_kind, access_key in self._self_path:
+                refreshed_value = (
+                    getattr(refreshed_value, access_key)
+                    if access_kind == "attr"
+                    else refreshed_value[access_key]
+                )
+            if isinstance(refreshed_value, MutableProxy):
+                super().__setattr__("__wrapped__", refreshed_value.__wrapped__)
+                self._self_state = refreshed_value._self_state
+                self._self_field_name = refreshed_value._self_field_name
+                self._self_path = refreshed_value._self_path
+            else:
+                self._raise_refresh_error()
+        except (Exception, asyncio.CancelledError):
+            await context_state.__aexit__(*sys.exc_info())
+            self._self_actx_state = None
+            raise
+        return self
+
+    def _raise_refresh_error(self) -> NoReturn:
+        """Raise when this proxy cannot be refreshed from its state field."""
+        msg = (
+            "Unable to refresh mutable proxy from state field "
+            f"`{self._self_field_name}`."
+        )
+        raise RuntimeError(msg)
+
+    async def __aexit__(self, *exc_info: Any) -> None:
+        """Exit the async context manager protocol through the bound state.
+
+        Args:
+            exc_info: The exception info tuple.
+        """
+        context_state = self._self_actx_state or self._self_state
+        try:
+            await context_state.__aexit__(*exc_info)
+        finally:
+            self._self_actx_state = None
+
     def _mark_dirty(
         self,
         wrapped: Callable | None = None,
@@ -513,11 +572,14 @@ def _is_called_from_dataclasses_internal() -> bool:
                 return True
         return False
 
-    def _wrap_recursive(self, value: Any) -> Any:
+    def _wrap_recursive(
+        self, value: Any, path: tuple[tuple[str, Any], ...] | None = None
+    ) -> Any:
         """Wrap a value recursively if it is mutable.
 
         Args:
             value: The value to wrap.
+            path: Access path from the state field to this wrapped object.
 
         Returns:
             The wrapped value.
@@ -536,6 +598,7 @@ def _wrap_recursive(self, value: Any) -> Any:
                 wrapped=value,
                 state=self._self_state,
                 field_name=self._self_field_name,
+                path=path or self._self_path,
             )
         return value
 
@@ -592,10 +655,11 @@ def __getattr__(self, __name: str) -> Any:
         if is_mutable_type(type(value)) and __name not in (
             "__wrapped__",
             "_self_state",
+            "_self_path",
             "__dict__",
         ):
             # Recursively wrap mutable attribute values retrieved through this proxy.
-            return self._wrap_recursive(value)
+            return self._wrap_recursive(value, (*self._self_path, ("attr", __name)))
 
         return value
 
@@ -612,7 +676,7 @@ def __getitem__(self, key: Any) -> Any:
         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)
+        return self._wrap_recursive(value, (*self._self_path, ("item", key)))
 
     def __iter__(self) -> Any:
         """Iterate over the proxied object and return a proxy if mutable.

tests/units/test_state.py

@@ -5006,6 +5006,61 @@ async def test_rebind_mutable_proxy(
         assert state.data["b"] == [2, 3]
 
 
+@pytest.mark.asyncio
+async def test_immutable_mutable_proxy_async_context_manager(
+    token: str, attached_mock_event_context: EventContext
+) -> None:
+    """Mutable state proxies can enter the owning StateProxy context."""
+    state_manager = attached_mock_event_context.state_manager
+
+    async with state_manager.modify_state(
+        BaseStateToken(ident=token, cls=MutableProxyState)
+    ) as state:
+        state.router = RouterData.from_router_data({
+            "query": {},
+            "token": token,
+            "sid": "test_sid",
+        })
+        state_proxy = StateProxy(state)
+        data_proxy = state_proxy.data
+        items_proxy = data_proxy["a"]
+
+    assert isinstance(data_proxy, ImmutableMutableProxy)
+    assert isinstance(items_proxy, ImmutableMutableProxy)
+    with pytest.raises(ImmutableStateError):
+        data_proxy["a"].append(3)
+    with pytest.raises(ImmutableStateError):
+        items_proxy.append(3)
+
+    async with state_manager.modify_state(
+        BaseStateToken(ident=token, cls=MutableProxyState)
+    ) as state:
+        assert isinstance(state, MutableProxyState)
+        state.data["a"].append(2)
+
+    async with data_proxy as mutable_data:
+        assert mutable_data is data_proxy
+        data_proxy["a"].append(3)
+        mutable_data["b"].append(4)
+
+    async with items_proxy as mutable_items:
+        assert mutable_items is items_proxy
+        mutable_items.append(5)
+        assert items_proxy.__wrapped__ == [1, 2, 3, 5]
+
+    with pytest.raises(ImmutableStateError):
+        data_proxy["a"].append(6)
+    with pytest.raises(ImmutableStateError):
+        items_proxy.append(6)
+
+    async with state_manager.modify_state(
+        BaseStateToken(ident=token, cls=MutableProxyState)
+    ) as state:
+        assert isinstance(state, MutableProxyState)
+        assert state.data["a"] == [1, 2, 3, 5]
+        assert state.data["b"] == [2, 4]
+
+
 def test_override_base_method_skips_event_handler_wrapping():
     """A method marked with __override_base_method__ should not be wrapped as an EventHandler."""
     from reflex.state import _override_base_method