Address the #7215 review follow-ups in BaseStateMeta

packages/reflex-base/news/+state-metaclass-contract.docs.md

@@ -0,0 +1 @@
+Document the supported `BaseStateMeta` contract for custom state metaclasses, including the `state_root` and `mixin` class keywords and the reserved `_reflex_state_root` attribute.

packages/reflex-base/src/reflex_base/vars/base.py

@@ -4669,18 +4669,23 @@ def _reserved_state_members(root: BaseStateMeta) -> dict[str, Any]:
     return members
 
 
-def _validate_state_name(root: BaseStateMeta, name: str, value: Any = None) -> None:
+def _validate_state_name(
+    root: BaseStateMeta | None, name: str, value: Any = None
+) -> None:
     """Reject declarations that replace framework methods or bookkeeping.
 
     Args:
-        root: The state class whose namespace is reserved.
+        root: The state class whose namespace is reserved, or None for a state
+            outside any hierarchy, which reserves nothing.
         name: The declared or dynamically registered name.
         value: The raw class declaration, when available.
 
     Raises:
         StateValueError: If a declaration uses a reserved name.
         EventHandlerShadowsBuiltInStateMethodError: If a method overrides a builtin.
     """
+    if root is None:
+        return
     members = _reserved_state_members(root)
     if name not in members:
         return
@@ -4957,13 +4962,28 @@ def _is_tree_state(cls: Any) -> bool:
     return (
         isinstance(cls, BaseStateMeta)
         and not cls._mixin
-        and getattr(cls, "_reflex_state_root", cls) is not cls
+        and (root := cls._reflex_state_root) is not None
+        and root is not cls
     )
 
 
 @dataclass_transform(kw_only_default=True, field_specifiers=(field,))
 class BaseStateMeta(ABCMeta):
-    """Meta class for BaseState."""
+    """Metaclass of every state class, including ``rx.State``.
+
+    ``type(rx.State) is BaseStateMeta`` is a supported contract: a downstream
+    metaclass that customizes state classes must derive from this one.
+
+    It understands two class keywords:
+
+    - ``state_root=True`` makes the class the root of a state hierarchy; its
+      members are reserved and its subclasses may not shadow them.
+    - ``mixin=True`` marks a state mixin, which is never instantiated. A class
+      whose state bases are all mixins is a mixin too.
+
+    ``_reflex_state_root`` is a reserved name: it holds the root of the class's
+    hierarchy (the class itself on a root), or None outside any hierarchy.
+    """
 
     if TYPE_CHECKING:
         __inherited_fields__: Mapping[str, Field]
@@ -4973,9 +4993,9 @@ class BaseStateMeta(ABCMeta):
         # Whether this state class is a mixin and should not be instantiated.
         _mixin: bool = False
 
-        # The state declared with ``state_root=True`` that this class descends
-        # from; its namespace is reserved for the whole hierarchy.
-        _reflex_state_root: BaseStateMeta
+    # The state declared with ``state_root=True`` that this class descends
+    # from; its namespace is reserved for the whole hierarchy. None outside one.
+    _reflex_state_root: BaseStateMeta | None = None
 
     def __setattr__(cls, name: str, value: Any) -> None:
         """Set a class attribute, refusing to replace a state var.
@@ -5046,15 +5066,14 @@ def __new__(
             The new class.
         """
         lookup_order = _linearize_bases(bases)
-        for base in bases:
-            root = getattr(base, "_reflex_state_root", None)
-            if root is not None:
-                _validate_state_declaration(root, lookup_order, namespace)
-                break
-
         state_bases = [
             base for base in bases if issubclass(base, EvenMoreBasicBaseState)
         ]
+        for base in state_bases:
+            if (root := base._reflex_state_root) is not None:
+                _validate_state_declaration(root, lookup_order, namespace)
+                break
+
         mixin = mixin or (
             bool(state_bases) and all(base._mixin for base in state_bases)
         )

tests/units/reflex_base/vars/test_base.py

@@ -50,8 +50,10 @@
     Var,
     VarData,
     _global_vars,
+    _is_tree_state,
     _linearize_bases,
     _type_check_depth,
+    _validate_state_name,
     cached_property,
     cached_property_no_lock,
     computed_var,
@@ -861,22 +863,30 @@ class ShadowState(BaseState):
             _init_bookkeeping: int
 
 
-@pytest.mark.parametrize("state_mixin", [False, True])
-def test_reserved_mixin_var(state_mixin: bool, clean_registration_context):
-    """Reject collisions from both ordinary Python mixins and state mixins.
+def test_reserved_mixin_var(clean_registration_context):
+    """Reject a reserved var from an ordinary Python mixin where a state inherits it.
 
     Args:
-        state_mixin: Whether the mixin subclasses BaseState.
         clean_registration_context: An isolated state registry.
     """
+    mixin = type("Mixin", (), {"__module__": __name__, "_get_root_state": 7})
     with pytest.raises(StateValueError, match="_get_root_state"):
-        mixin = type(
+        type("MixedState", (mixin, BaseState), {"__module__": __name__})
+
+
+def test_reserved_state_mixin_var(clean_registration_context):
+    """Reject a reserved var where a state mixin declares it.
+
+    Args:
+        clean_registration_context: An isolated state registry.
+    """
+    with pytest.raises(StateValueError, match="_get_root_state"):
+        type(
             "Mixin",
-            (BaseState,) if state_mixin else (),
+            (BaseState,),
             {"__module__": __name__, "_get_root_state": 7},
-            **({"mixin": True} if state_mixin else {}),
+            mixin=True,
         )
-        type("MixedState", (mixin, BaseState), {"__module__": __name__})
 
 
 @pytest.mark.parametrize("state_mixin", [False, True])
@@ -1198,6 +1208,31 @@ class ShadowA(RootA):
             alpha: int = 3
 
 
+def test_only_states_below_a_root_are_tree_states():
+    """Treat neither a root nor a state outside any root's hierarchy as a tree state."""
+
+    class Rootless(EvenMoreBasicBaseState):
+        pass
+
+    class Root(EvenMoreBasicBaseState, state_root=True):
+        pass
+
+    class Child(Root):
+        pass
+
+    assert Rootless._reflex_state_root is None
+    assert not _is_tree_state(Rootless)
+    assert not _is_tree_state(Root)
+    assert _is_tree_state(Child)
+
+
+def test_no_root_reserves_no_names():
+    """Accept any name for a state outside every root's hierarchy."""
+    _validate_state_name(None, "_get_root_state", 7)
+    with pytest.raises(StateValueError, match="_get_root_state"):
+        _validate_state_name(BaseState, "_get_root_state", 7)
+
+
 def test_inherited_field_on_plain_model():
     """An inherited field of a model without a state tree reads from the model itself."""