diff --git a/reflex/state.py b/reflex/state.py index 010af0bd1..80c56332e 100644 --- a/reflex/state.py +++ b/reflex/state.py @@ -1868,7 +1868,13 @@ class ImmutableMutableProxy(MutableProxy): to modify the wrapped object when the StateProxy is immutable. """ - def _mark_dirty(self, wrapped=None, instance=None, args=tuple(), kwargs=None): + def _mark_dirty( + self, + wrapped=None, + instance=None, + args=tuple(), + kwargs=None, + ) -> Any: """Raise an exception when an attempt is made to modify the object. Intended for use with `FunctionWrapper` from the `wrapt` library. @@ -1879,6 +1885,9 @@ class ImmutableMutableProxy(MutableProxy): 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. """ @@ -1887,6 +1896,6 @@ class ImmutableMutableProxy(MutableProxy): "Background task StateProxy is immutable outside of a context " "manager. Use `async with self` to modify state." ) - super()._mark_dirty( + return super()._mark_dirty( wrapped=wrapped, instance=instance, args=args, kwargs=kwargs ) diff --git a/tests/test_state.py b/tests/test_state.py index b96b263ce..c6537b55f 100644 --- a/tests/test_state.py +++ b/tests/test_state.py @@ -1625,7 +1625,7 @@ class BackgroundTaskState(State): """A state with a background task.""" order: List[str] = [] - dict_list: Dict[str, List[int]] = {"foo": []} + dict_list: Dict[str, List[int]] = {"foo": [1, 2, 3]} @rx.background async def background_task(self): @@ -1657,6 +1657,9 @@ class BackgroundTaskState(State): pass # update proxy instance async with self: + # Methods on ImmutableMutableProxy should return their wrapped return value. + assert self.dict_list.pop("foo") is not None + self.order.append("background_task:stop") self.other() # direct calling event handlers works in context self._private_method()