diff --git a/news/+get-state-mixin.feature.md b/news/+get-state-mixin.feature.md new file mode 100644 index 00000000000..93f3e72c024 --- /dev/null +++ b/news/+get-state-mixin.feature.md @@ -0,0 +1 @@ +`get_state()` accepts a state mixin and returns the state it was mixed into, so shared handlers can reach it without naming the concrete class: `await self.get_state(AuthMixin)`. A `BaseStateToken` whose `cls` is a mixin addresses that state in the state manager too. Both raise `StateValueError` when no state, or more than one unrelated state, uses the mixin. diff --git a/packages/reflex-base/news/+registry-mixin-lookup.feature.md b/packages/reflex-base/news/+registry-mixin-lookup.feature.md new file mode 100644 index 00000000000..d375e2a3a34 --- /dev/null +++ b/packages/reflex-base/news/+registry-mixin-lookup.feature.md @@ -0,0 +1 @@ +Add `RegistrationContext.get_states_implementing()`, which lists the registered states that inherit from a state mixin, and `RegistrationContext.resolve_implementation()`, which returns the single state the mixin was mixed into and raises `StateValueError` when there is none or the choice is ambiguous. diff --git a/packages/reflex-base/src/reflex_base/registry.py b/packages/reflex-base/src/reflex_base/registry.py index 05fb5c694da..7642bc6c90c 100644 --- a/packages/reflex-base/src/reflex_base/registry.py +++ b/packages/reflex-base/src/reflex_base/registry.py @@ -3,7 +3,7 @@ from __future__ import annotations import dataclasses -from typing import TYPE_CHECKING, Any, Self +from typing import TYPE_CHECKING, Any, Self, TypeVar from reflex_base.context.base import BaseContext from reflex_base.utils.exceptions import ReflexRuntimeError, StateValueError @@ -19,6 +19,8 @@ from reflex_base.utils.types import ArgsSpec from reflex_base.vars.base import Var +STATE_TYPE = TypeVar("STATE_TYPE", bound="BaseState") + def _default_bundled_libraries() -> list[str]: """Return the initial set of bundled libraries for a new context. @@ -61,6 +63,19 @@ class RegistrationContext(BaseContext): default_factory=dict, repr=False, ) + # Memoized lookups, invalidated whenever a base state is registered. + _implementations: dict[type[BaseState], tuple[type[BaseState], ...]] = ( + dataclasses.field( + default_factory=dict, + repr=False, + ) + ) + _resolved_implementations: dict[type[BaseState], type[BaseState]] = ( + dataclasses.field( + default_factory=dict, + repr=False, + ) + ) _config: Config | None = dataclasses.field(default=None, repr=False) decorated_pages: list[tuple[Callable, dict[str, Any]]] = dataclasses.field( default_factory=list, @@ -215,6 +230,8 @@ def _register_base_state(self, state_cls: type[BaseState]) -> type[BaseState]: The registered base state class. """ self.base_states[state_cls.get_full_name()] = state_cls + self._implementations.clear() + self._resolved_implementations.clear() for event_handler in state_cls.event_handlers.values(): self._register_event_handler(event_handler, states=(state_cls,)) if (parent_state := state_cls.get_parent_state()) is not None: @@ -287,3 +304,66 @@ def get_substates( return self.base_state_substates.setdefault( base_state_cls.get_full_name(), set() ) + + def get_states_implementing( + self, interface: type[STATE_TYPE] + ) -> tuple[type[STATE_TYPE], ...]: + """Get every registered state class that is a subclass of the given interface. + + Args: + interface: The state mixin (or state class) to look up implementations of. + + Returns: + The registered states implementing the interface, in registration order. + """ + if (cached := self._implementations.get(interface)) is not None: + return cached # pyright: ignore [reportReturnType] + implementations = tuple( + state_cls + for state_cls in self.base_states.values() + if issubclass(state_cls, interface) + ) + self._implementations[interface] = implementations + return implementations + + def resolve_implementation(self, interface: type[STATE_TYPE]) -> type[STATE_TYPE]: + """Get the single registered state class the given interface was mixed into. + + Args: + interface: The state mixin (or state class) to resolve. + + Returns: + The state class the interface was mixed into. + + Raises: + StateValueError: If no or more than one state implements the interface. + """ + if (cached := self._resolved_implementations.get(interface)) is not None: + return cached # pyright: ignore [reportReturnType] + implementations = self.get_states_implementing(interface) + # Keep only where the interface entered the tree: an implementation + # inheriting from another implementation is not a candidate. + roots = [ + state_cls + for state_cls in implementations + if not any( + other is not state_cls and issubclass(state_cls, other) + for other in implementations + ) + ] + if not roots: + msg = ( + f"No registered state implements {interface.__name__}. " + "Define a state which inherits from it before resolving it." + ) + raise StateValueError(msg) + if len(roots) > 1: + names = ", ".join(sorted(state.get_full_name() for state in roots)) + msg = ( + f"{interface.__name__} is implemented by more than one state: {names}. " + "Resolving it is ambiguous, pass the wanted state class explicitly." + ) + raise StateValueError(msg) + resolved = roots[0] + self._resolved_implementations[interface] = resolved + return resolved diff --git a/reflex/istate/manager/__init__.py b/reflex/istate/manager/__init__.py index 156c5af687e..4846da3e62e 100644 --- a/reflex/istate/manager/__init__.py +++ b/reflex/istate/manager/__init__.py @@ -7,11 +7,12 @@ from abc import ABC, abstractmethod from collections.abc import AsyncIterator from types import TracebackType -from typing import TYPE_CHECKING, Generic, TypedDict, overload +from typing import TYPE_CHECKING, Generic, TypedDict, cast, overload from reflex_base import constants from reflex_base.config import get_config from reflex_base.event import Event +from reflex_base.registry import RegistrationContext from reflex_base.utils.exceptions import InvalidStateManagerModeError from typing_extensions import ReadOnly, Unpack, deprecated @@ -98,7 +99,7 @@ def create(cls): @staticmethod def _coerce_token(token: StateToken[TOKEN_TYPE] | str) -> StateToken[TOKEN_TYPE]: - """Convert a legacy string token to a StateToken if needed. + """Normalize a token: convert legacy string tokens and resolve state mixins. Args: token: The token, either a StateToken or legacy string. @@ -111,6 +112,12 @@ def _coerce_token(token: StateToken[TOKEN_TYPE] | str) -> StateToken[TOKEN_TYPE] from reflex.state import State return BaseStateToken.from_legacy_token(token, root_state=State) # type: ignore[return-value] + # The store is never keyed on a mixin: resolve it to the state it was mixed into. + if getattr(token.cls, "_mixin", False): + resolved = RegistrationContext.get().resolve_implementation( + cast("type[BaseState]", token.cls) + ) + return token.with_cls(cast("type[TOKEN_TYPE]", resolved)) return token @overload diff --git a/reflex/state.py b/reflex/state.py index bb7fbca3546..0c2633b2ed0 100644 --- a/reflex/state.py +++ b/reflex/state.py @@ -1793,11 +1793,15 @@ async def get_state(self, state_cls: type[T_STATE]) -> T_STATE: Allows for arbitrary access to sibling states from within an event handler. Args: - state_cls: The class of the state. + state_cls: The class of the state, or a state mixin to resolve. Returns: The instance of state_cls associated with this state's client_token. """ + # A mixin has no place in the state tree of its own: resolve it. + if state_cls._mixin: + state_cls = RegistrationContext.get().resolve_implementation(state_cls) + # Fast case - if this state instance is already cached, get_substate from root state. try: return self._get_state_from_cache(state_cls) diff --git a/tests/units/reflex_base/test_registry.py b/tests/units/reflex_base/test_registry.py index 82490effd6c..17c37474357 100644 --- a/tests/units/reflex_base/test_registry.py +++ b/tests/units/reflex_base/test_registry.py @@ -424,6 +424,130 @@ def test_bundled_libraries_isolated_between_contexts(): assert "some-extra-lib" not in ctx_b.bundled_libraries +def test_resolve_implementation(clean_registration_context: RegistrationContext): + """resolve_implementation finds the state a mixin was mixed into. + + Args: + clean_registration_context: A fresh, empty registration context. + """ + from reflex.state import BaseState + + class ResolveMixin(BaseState, mixin=True): + pass + + class ResolveImpl(ResolveMixin, BaseState): + pass + + assert clean_registration_context.resolve_implementation(ResolveMixin) is ( + ResolveImpl + ) + assert clean_registration_context.get_states_implementing(ResolveMixin) == ( + ResolveImpl, + ) + + +def test_resolve_implementation_ignores_subclasses( + clean_registration_context: RegistrationContext, +): + """Subclasses of an implementation do not make the resolution ambiguous. + + Args: + clean_registration_context: A fresh, empty registration context. + """ + from reflex.state import BaseState + + class SubclassMixin(BaseState, mixin=True): + pass + + class SubclassImpl(SubclassMixin, BaseState): + pass + + class SubclassChild(SubclassImpl): + pass + + assert set(clean_registration_context.get_states_implementing(SubclassMixin)) == { + SubclassImpl, + SubclassChild, + } + assert clean_registration_context.resolve_implementation(SubclassMixin) is ( + SubclassImpl + ) + + +def test_resolve_implementation_without_implementation( + clean_registration_context: RegistrationContext, +): + """Resolving a mixin no state implements raises StateValueError. + + Args: + clean_registration_context: A fresh, empty registration context. + """ + from reflex.state import BaseState + + class UnusedMixin(BaseState, mixin=True): + pass + + assert clean_registration_context.get_states_implementing(UnusedMixin) == () + with pytest.raises(StateValueError, match="No registered state implements"): + clean_registration_context.resolve_implementation(UnusedMixin) + + +def test_resolve_implementation_ambiguous( + clean_registration_context: RegistrationContext, +): + """Resolving a mixin used by two unrelated states raises StateValueError. + + Args: + clean_registration_context: A fresh, empty registration context. + """ + from reflex.state import BaseState + + class AmbiguousMixin(BaseState, mixin=True): + pass + + class AmbiguousImpl1(AmbiguousMixin, BaseState): + pass + + class AmbiguousImpl2(AmbiguousMixin, BaseState): + pass + + with pytest.raises(StateValueError, match="implemented by more than one state"): + clean_registration_context.resolve_implementation(AmbiguousMixin) + + +def test_resolve_implementation_cache_is_invalidated( + clean_registration_context: RegistrationContext, +): + """Registering a state invalidates previously cached lookups. + + Args: + clean_registration_context: A fresh, empty registration context. + """ + from reflex.state import BaseState + + class CacheMixin(BaseState, mixin=True): + pass + + class CacheImpl(CacheMixin, BaseState): + pass + + assert clean_registration_context.resolve_implementation(CacheMixin) is CacheImpl + # Repeated lookups are served from the cache. + assert clean_registration_context.get_states_implementing( + CacheMixin + ) is clean_registration_context.get_states_implementing(CacheMixin) + + class CacheImpl2(CacheMixin, BaseState): + pass + + assert set(clean_registration_context.get_states_implementing(CacheMixin)) == { + CacheImpl, + CacheImpl2, + } + with pytest.raises(StateValueError, match="implemented by more than one state"): + clean_registration_context.resolve_implementation(CacheMixin) + + def test_reset_compile_caches_empties_per_compile_maps( clean_registration_context: RegistrationContext, ): diff --git a/tests/units/test_state.py b/tests/units/test_state.py index 5f6fd7cceff..3dc3d37a95b 100644 --- a/tests/units/test_state.py +++ b/tests/units/test_state.py @@ -7120,6 +7120,120 @@ def test_setstate_drops_the_legacy_router_entry(): assert state.router.session.client_token == "" +async def test_get_state_mixin(clean_registration_context): + """get_state resolves a mixin to the state it was applied to. + + Args: + clean_registration_context: A fresh, empty registration context. + """ + + class MixinRootState(BaseState): + pass + + class GreetMixin(BaseState, mixin=True): + greeting: str = "hi" + + class GreetState(GreetMixin, MixinRootState): + pass + + class UnrelatedState(MixinRootState): + pass + + root = MixinRootState(_reflex_internal_init=True) # pyright: ignore [reportCallIssue] + unrelated = await root.get_state(UnrelatedState) + + greet_state = await unrelated.get_state(GreetMixin) + assert isinstance(greet_state, GreetState) + # Resolving the mixin and naming the state directly yield the same instance. + assert greet_state is await unrelated.get_state(GreetState) + assert greet_state.greeting == "hi" + + +async def test_get_state_mixin_ambiguous(clean_registration_context): + """get_state raises if more than one state implements the mixin. + + Args: + clean_registration_context: A fresh, empty registration context. + """ + + class AmbiguousRootState(BaseState): + pass + + class AmbiguousMixin(BaseState, mixin=True): + pass + + class AmbiguousState1(AmbiguousMixin, AmbiguousRootState): + pass + + class AmbiguousState2(AmbiguousMixin, AmbiguousRootState): + pass + + root = AmbiguousRootState(_reflex_internal_init=True) # pyright: ignore [reportCallIssue] + with pytest.raises(StateValueError, match="implemented by more than one state"): + await root.get_state(AmbiguousMixin) + + +async def test_state_manager_mixin_token( + clean_registration_context, state_manager: StateManager, token: str +): + """A token carrying a mixin resolves to the state implementing it. + + Args: + clean_registration_context: A fresh, empty registration context. + state_manager: A state manager instance. + token: A token. + """ + + class TokenRootState(BaseState): + pass + + class NoteMixin(BaseState, mixin=True): + note: str = "" + + class NoteState(NoteMixin, TokenRootState): + pass + + async with state_manager.modify_state( + BaseStateToken(ident=token, cls=NoteMixin) + ) as root: + assert isinstance(root, TokenRootState) + (await root.get_state(NoteMixin)).note = "written" + + if environment.REFLEX_OPLOCK_ENABLED.get(): + await state_manager.close() + + # The concrete token addresses the very same state. + root = await state_manager.get_state(BaseStateToken(ident=token, cls=NoteState)) + assert (await root.get_state(NoteState)).note == "written" + + +async def test_state_manager_mixin_token_ambiguous( + clean_registration_context, state_manager: StateManager, token: str +): + """A token carrying an ambiguous mixin raises instead of guessing. + + Args: + clean_registration_context: A fresh, empty registration context. + state_manager: A state manager instance. + token: A token. + """ + + class AmbiguousTokenRootState(BaseState): + pass + + class TwiceUsedMixin(BaseState, mixin=True): + pass + + class FirstUserState(TwiceUsedMixin, AmbiguousTokenRootState): + pass + + class SecondUserState(TwiceUsedMixin, AmbiguousTokenRootState): + pass + + with pytest.raises(StateValueError, match="implemented by more than one state"): + await state_manager.get_state(BaseStateToken(ident=token, cls=TwiceUsedMixin)) + + def test_previous_release_pickle_keys_are_reserved(): """A field cannot take the name older pickles kept the backend vars under.""" with pytest.raises(StateValueError, match="_backend_vars"):