diff --git a/news/6689.feature.md b/news/6689.feature.md new file mode 100644 index 00000000000..b549066792a --- /dev/null +++ b/news/6689.feature.md @@ -0,0 +1 @@ +Support using mutable state proxies as async context managers. diff --git a/reflex/istate/proxy.py b/reflex/istate/proxy.py index 776a519fe2c..67cbd7a94fb 100644 --- a/reflex/istate/proxy.py +++ b/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, Literal, NoReturn, SupportsIndex, TypeVar import wrapt from reflex_base.event import Event @@ -29,6 +29,13 @@ T_STATE = TypeVar("T_STATE", bound="BaseState") T = TypeVar("T") +AccessSpec = ( + tuple[Literal["attr"], str] + | tuple[Literal["item"], Any] + | tuple[Literal["iter"], None] +) +_ITER_ACCESS_SPEC: AccessSpec = ("iter", None) +_ROOT_ITER_PATH: tuple[AccessSpec, ...] = (_ITER_ACCESS_SPEC,) # Cached filename of the dataclasses module, used to detect reads originating # from `dataclasses.asdict`/`astuple` internals on the proxy read hot-path. @@ -446,7 +453,13 @@ 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[AccessSpec, ...] | None = None, + ): """Create a proxy for a mutable object that tracks changes. Args: @@ -454,10 +467,13 @@ def __init__(self, wrapped: Any, state: BaseState, field_name: str): 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 +483,73 @@ 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. + """ + if self._self_actx_state is not None: + msg = ( + "Mutable proxy is already mutable. Do not nest `async with proxy` " + "blocks." + ) + raise RuntimeError(msg) + context_state = self._self_state + self._self_actx_state = context_state + aenter_ok = False + try: + state = await context_state.__aenter__() + aenter_ok = True + refreshed_value = getattr(state, self._self_field_name) + for access_spec in self._self_path: + match access_spec: + case ("attr", access_key): + refreshed_value = getattr(refreshed_value, access_key) + case ("item", access_key): + refreshed_value = refreshed_value[access_key] + case _: + self._raise_refresh_error() + if ( + isinstance(refreshed_value, MutableProxy) + and self._self_field_name == refreshed_value._self_field_name + ): + super().__setattr__("__wrapped__", refreshed_value.__wrapped__) + self._self_state = refreshed_value._self_state + self._self_path = refreshed_value._self_path + else: + self._raise_refresh_error() + except (Exception, asyncio.CancelledError): + try: + if aenter_ok: + await context_state.__aexit__(*sys.exc_info()) + finally: + 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 + if context_state is None: + return + try: + await context_state.__aexit__(*exc_info) + finally: + self._self_actx_state = None + def _mark_dirty( self, wrapped: Callable | None = None, @@ -513,11 +596,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, new_path_segment: AccessSpec | None = None + ) -> Any: """Wrap a value recursively if it is mutable. Args: value: The value to wrap. + new_path_segment: Access path segment from this proxy to the value. Returns: The wrapped value. @@ -532,10 +618,18 @@ def _wrap_recursive(self, value: Any) -> Any: # Recursively wrap mutable types. if is_mutable_type(type(value)): base_cls = globals()[self.__base_proxy__] + path = self._self_path + if new_path_segment is not None: + path = ( + _ROOT_ITER_PATH + if new_path_segment is _ITER_ACCESS_SPEC and not path + else (*path, new_path_segment) + ) return base_cls( wrapped=value, state=self._self_state, field_name=self._self_field_name, + path=path, ) return value @@ -592,10 +686,12 @@ def __getattr__(self, __name: str) -> Any: if is_mutable_type(type(value)) and __name not in ( "__wrapped__", "_self_state", + "_self_path", + "_self_actx_state", "__dict__", ): # Recursively wrap mutable attribute values retrieved through this proxy. - return self._wrap_recursive(value) + return self._wrap_recursive(value, ("attr", __name)) return value @@ -610,9 +706,13 @@ def __getitem__(self, key: Any) -> Any: """ value = super().__getitem__(key) # pyright: ignore[reportAttributeAccessIssue] if isinstance(key, slice) and isinstance(value, list): - return [self._wrap_recursive(item) for item in value] + indices = range(len(self.__wrapped__))[key] + return [ + self._wrap_recursive(item, ("item", index)) + for item, index in zip(value, indices, strict=False) + ] # Recursively wrap mutable items retrieved through this proxy. - return self._wrap_recursive(value) + return self._wrap_recursive(value, ("item", key)) def __iter__(self) -> Any: """Iterate over the proxied object and return a proxy if mutable. @@ -622,7 +722,7 @@ def __iter__(self) -> Any: """ for value in super().__iter__(): # pyright: ignore[reportAttributeAccessIssue] # Recursively wrap mutable items retrieved through this proxy. - yield self._wrap_recursive(value) + yield self._wrap_recursive(value, _ITER_ACCESS_SPEC) def __delattr__(self, name: str): """Delete the attribute on the proxied object and mark state dirty. diff --git a/tests/units/test_state.py b/tests/units/test_state.py index b0aaf6e022f..181d4e95c7b 100644 --- a/tests/units/test_state.py +++ b/tests/units/test_state.py @@ -4954,6 +4954,12 @@ class MutableProxyState(BaseState): data: dict[str, list[int]] = {"a": [1], "b": [2]} +class IterableMutableProxyState(BaseState): + """A test state with mutable values that can be accessed by iteration.""" + + data: list[list[int]] = [[1]] + + @pytest.mark.asyncio async def test_rebind_mutable_proxy( token: str, attached_mock_event_context: EventContext @@ -5006,6 +5012,179 @@ 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) + await data_proxy.__aexit__(None, None, None) + 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 + with pytest.raises(RuntimeError, match="already mutable"): + async with data_proxy: + pass + 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] + + +@pytest.mark.asyncio +async def test_immutable_mutable_proxy_async_context_rejects_iter_proxy( + token: str, attached_mock_event_context: EventContext +) -> None: + """Iteration-sourced mutable proxies fail clearly as async context managers.""" + state_manager = attached_mock_event_context.state_manager + + async with state_manager.modify_state( + BaseStateToken(ident=token, cls=IterableMutableProxyState) + ) as state: + state.router = RouterData.from_router_data({ + "query": {}, + "token": token, + "sid": "test_sid", + }) + state_proxy = StateProxy(state) + [items_proxy] = state_proxy.data + + assert isinstance(items_proxy, ImmutableMutableProxy) + with pytest.raises(RuntimeError, match="Unable to refresh mutable proxy"): + async with items_proxy: + pass + + async with state_proxy: + state_proxy.data[0].append(2) + + async with state_manager.modify_state( + BaseStateToken(ident=token, cls=IterableMutableProxyState) + ) as state: + assert isinstance(state, IterableMutableProxyState) + assert state.data == [[1, 2]] + + +@pytest.mark.asyncio +async def test_immutable_mutable_proxy_async_context_recovers_from_enter_failure( + token: str, + attached_mock_event_context: EventContext, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """A failed owning StateProxy enter does not permanently block the proxy.""" + 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 + + original_aenter = StateProxy.__aenter__ + fail_once = True + + async def fail_first_enter(self: StateProxy) -> StateProxy: + nonlocal fail_once + if fail_once: + fail_once = False + raise asyncio.CancelledError + return await original_aenter(self) + + monkeypatch.setattr(StateProxy, "__aenter__", fail_first_enter) + + with pytest.raises(asyncio.CancelledError): + async with data_proxy: + pass + + assert data_proxy._self_actx_state is None + async with data_proxy as mutable_data: + mutable_data["a"].append(2) + + async with state_manager.modify_state( + BaseStateToken(ident=token, cls=MutableProxyState) + ) as state: + assert isinstance(state, MutableProxyState) + assert state.data["a"] == [1, 2] + + +@pytest.mark.asyncio +async def test_immutable_mutable_proxy_async_context_clears_state_when_cleanup_fails( + token: str, + attached_mock_event_context: EventContext, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """A failed cleanup while entering does not leave the proxy marked mutable.""" + state_manager = attached_mock_event_context.state_manager + + async with state_manager.modify_state( + BaseStateToken(ident=token, cls=IterableMutableProxyState) + ) as state: + state.router = RouterData.from_router_data({ + "query": {}, + "token": token, + "sid": "test_sid", + }) + state_proxy = StateProxy(state) + [items_proxy] = state_proxy.data + + async def fail_exit(self: StateProxy, *exc_info: Any) -> None: + await asyncio.sleep(0) + msg = "cleanup failed" + raise RuntimeError(msg) + + monkeypatch.setattr(StateProxy, "__aexit__", fail_exit) + + with pytest.raises(RuntimeError, match="cleanup failed"): + async with items_proxy: + pass + + assert items_proxy._self_actx_state is None + + 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