unpack

Author: adhami3310

reflex/app.py

@@ -1759,7 +1759,7 @@ async def process(
         })
         # Get the state for the session exclusively.
         async with app.state_manager.modify_state(
-            event.substate_token, context=event
+            event.substate_token, event=event
         ) as state:
             # When this is a brand new instance of the state, signal the
             # frontend to reload before processing it.

reflex/istate/manager/__init__.py

@@ -4,6 +4,9 @@
 import dataclasses
 from abc import ABC, abstractmethod
 from collections.abc import AsyncIterator
+from typing import TypedDict
+
+from typing_extensions import ReadOnly, Unpack
 
 from reflex import constants
 from reflex.config import get_config
@@ -13,6 +16,15 @@
 from reflex.utils.exceptions import InvalidStateManagerModeError
 
 
+class StateModificationContext(TypedDict, total=False):
+    """The context for modifying state."""
+
+    event: ReadOnly[Event]
+
+
+EmptyContext = StateModificationContext()
+
+
 @dataclasses.dataclass
 class StateManager(ABC):
     """A class to manage many client states."""
@@ -73,26 +85,29 @@ async def get_state(self, token: str) -> BaseState:
 
     @abstractmethod
     async def set_state(
-        self, token: str, state: BaseState, *, context: Event | None = None
+        self,
+        token: str,
+        state: BaseState,
+        **context: Unpack[StateModificationContext],
     ):
         """Set the state for a token.
 
         Args:
             token: The token to set the state for.
             state: The state to set.
-            context: The event context.
+            context: The state modification context.
         """
 
     @abstractmethod
     @contextlib.asynccontextmanager
     async def modify_state(
-        self, token: str, *, context: Event | None = None
+        self, token: str, **context: Unpack[StateModificationContext]
     ) -> AsyncIterator[BaseState]:
         """Modify the state for a token while holding exclusive lock.
 
         Args:
             token: The token to modify the state for.
-            context: The event context.
+            context: The state modification context.
 
         Yields:
             The state for the token.

reflex/istate/manager/disk.py

@@ -9,11 +9,14 @@
 from hashlib import md5
 from pathlib import Path
 
-from typing_extensions import override
+from typing_extensions import Unpack, override
 
 from reflex.environment import environment
-from reflex.event import Event
-from reflex.istate.manager import StateManager, _default_token_expiration
+from reflex.istate.manager import (
+    StateManager,
+    StateModificationContext,
+    _default_token_expiration,
+)
 from reflex.state import BaseState, _split_substate_key, _substate_key
 from reflex.utils import console, path_ops, prerequisites
 from reflex.utils.misc import run_in_thread
@@ -303,14 +306,14 @@ async def _schedule_process_write_queue(self):
 
     @override
     async def set_state(
-        self, token: str, state: BaseState, *, context: Event | None = None
+        self, token: str, state: BaseState, **context: Unpack[StateModificationContext]
     ):
         """Set the state for a token.
 
         Args:
             token: The token to set the state for.
             state: The state to set.
-            context: The event context.
+            context: The state modification context.
         """
         client_token, _ = _split_substate_key(token)
         if self._write_debounce_seconds > 0:
@@ -330,13 +333,13 @@ async def set_state(
     @override
     @contextlib.asynccontextmanager
     async def modify_state(
-        self, token: str, *, context: Event | None = None
+        self, token: str, **context: Unpack[StateModificationContext]
     ) -> AsyncIterator[BaseState]:
         """Modify the state for a token while holding exclusive lock.
 
         Args:
             token: The token to modify the state for.
-            context: The event context.
+            context: The state modification context.
 
         Yields:
             The state for the token.
@@ -351,7 +354,7 @@ async def modify_state(
         async with self._states_locks[client_token]:
             state = await self.get_state(token)
             yield state
-            await self.set_state(token, state, context=context)
+            await self.set_state(token, state, **context)
 
     async def close(self):
         """Close the state manager, flushing any pending writes to disk."""

reflex/istate/manager/memory.py

@@ -5,10 +5,9 @@
 import dataclasses
 from collections.abc import AsyncIterator
 
-from typing_extensions import override
+from typing_extensions import Unpack, override
 
-from reflex.event import Event
-from reflex.istate.manager import StateManager
+from reflex.istate.manager import StateManager, StateModificationContext
 from reflex.state import BaseState, _split_substate_key
 
 
@@ -45,28 +44,31 @@ async def get_state(self, token: str) -> BaseState:
 
     @override
     async def set_state(
-        self, token: str, state: BaseState, *, context: Event | None = None
+        self,
+        token: str,
+        state: BaseState,
+        **context: Unpack[StateModificationContext],
     ):
         """Set the state for a token.
 
         Args:
             token: The token to set the state for.
             state: The state to set.
-            context: The event context.
+            context: The state modification context.
         """
         token = _split_substate_key(token)[0]
         self.states[token] = state
 
     @override
     @contextlib.asynccontextmanager
     async def modify_state(
-        self, token: str, *, context: Event | None = None
+        self, token: str, **context: Unpack[StateModificationContext]
     ) -> AsyncIterator[BaseState]:
         """Modify the state for a token while holding exclusive lock.
 
         Args:
             token: The token to modify the state for.
-            context: The event context.
+            context: The state modification context.
 
         Yields:
             The state for the token.

reflex/istate/manager/redis.py

@@ -10,12 +10,15 @@
 from redis import ResponseError
 from redis.asyncio import Redis
 from redis.asyncio.client import PubSub
-from typing_extensions import override
+from typing_extensions import Unpack, override
 
 from reflex.config import get_config
 from reflex.environment import environment
-from reflex.event import Event
-from reflex.istate.manager import StateManager, _default_token_expiration
+from reflex.istate.manager import (
+    StateManager,
+    StateModificationContext,
+    _default_token_expiration,
+)
 from reflex.state import BaseState, _split_substate_key, _substate_key
 from reflex.utils import console
 from reflex.utils.exceptions import (
@@ -264,7 +267,7 @@ async def set_state(
         state: BaseState,
         *,
         lock_id: bytes | None = None,
-        context: Event | None = None,
+        **context: Unpack[StateModificationContext],
     ):
         """Set the state for a token.
 
@@ -287,7 +290,11 @@ async def set_state(
                 f"Lock expired for token {token} while processing. Consider increasing "
                 f"`app.state_manager.lock_expiration` (currently {self.lock_expiration}) "
                 "or use `@rx.event(background=True)` decorator for long-running tasks."
-                + (f" Happened in event: {context.name}" if context is not None else "")
+                + (
+                    f" Happened in event: {event.name}"
+                    if (event := context.get("event")) is not None
+                    else ""
+                )
             )
             raise LockExpiredError(msg)
         if lock_id is not None:
@@ -299,8 +306,8 @@ async def set_state(
                     f"Lock for token {token} was held too long {time_taken=}s, "
                     f"use `@rx.event(background=True)` decorator for long-running tasks."
                     + (
-                        f" Happened in event: {context.name}"
-                        if context is not None
+                        f" Happened in event: {event.name}"
+                        if (event := context.get("event")) is not None
                         else ""
                     ),
                     dedupe=True,
@@ -319,7 +326,7 @@ async def set_state(
                     _substate_key(client_token, substate),
                     substate,
                     lock_id=lock_id,
-                    context=context,
+                    **context,
                 ),
                 name=f"reflex_set_state|{client_token}|{substate.get_full_name()}",
             )
@@ -342,21 +349,21 @@ async def set_state(
     @override
     @contextlib.asynccontextmanager
     async def modify_state(
-        self, token: str, *, context: Event | None = None
+        self, token: str, **context: Unpack[StateModificationContext]
     ) -> AsyncIterator[BaseState]:
         """Modify the state for a token while holding exclusive lock.
 
         Args:
             token: The token to modify the state for.
-            context: The event context.
+            context: The state modification context.
 
         Yields:
             The state for the token.
         """
         async with self._lock(token) as lock_id:
             state = await self.get_state(token)
             yield state
-            await self.set_state(token, state, lock_id=lock_id, context=context)
+            await self.set_state(token, state, lock_id=lock_id, **context)
 
     @staticmethod
     def _lock_key(token: str) -> bytes: