From ed9c351141cc8bc78eb9e96b7c0132828cdaa044 Mon Sep 17 00:00:00 2001
From: Khaleel Al-Adhami <khaleel.aladhami@gmail.com>
Date: Fri, 10 Jan 2025 16:38:34 -0800
Subject: [PATCH] dang it darglint

---
 reflex/state.py            | 9 +++++++--
 reflex/utils/exceptions.py | 4 ++++
 2 files changed, 11 insertions(+), 2 deletions(-)

diff --git a/reflex/state.py b/reflex/state.py
index 9a2448819..16875bec3 100644
--- a/reflex/state.py
+++ b/reflex/state.py
@@ -104,6 +104,7 @@ from reflex.utils.exceptions import (
     LockExpiredError,
     ReflexRuntimeError,
     SetUndefinedStateVarError,
+    StateMismatchError,
     StateSchemaMismatchError,
     StateSerializationError,
     StateTooLargeError,
@@ -1551,11 +1552,14 @@ class BaseState(Base, ABC, extra=pydantic.Extra.allow):
 
         Returns:
             The instance of state_cls associated with this state's client_token.
+
+        Raises:
+            StateMismatchError: If the state instance is not of the expected type.
         """
         root_state = self._get_root_state()
         substate = root_state.get_substate(state_cls.get_full_name().split("."))
         if not isinstance(substate, state_cls):
-            raise ValueError(
+            raise StateMismatchError(
                 f"Searched for state {state_cls.get_full_name()} but found {substate}."
             )
         return substate
@@ -1571,6 +1575,7 @@ class BaseState(Base, ABC, extra=pydantic.Extra.allow):
 
         Raises:
             RuntimeError: If redis is not used in this backend process.
+            StateMismatchError: If the state instance is not of the expected type.
         """
         # Fetch all missing parent states from redis.
         parent_state_of_state_cls = await self._populate_parent_states(state_cls)
@@ -1591,7 +1596,7 @@ class BaseState(Base, ABC, extra=pydantic.Extra.allow):
         )
 
         if not isinstance(state_in_redis, state_cls):
-            raise ValueError(
+            raise StateMismatchError(
                 f"Searched for state {state_cls.get_full_name()} but found {state_in_redis}."
             )
 
diff --git a/reflex/utils/exceptions.py b/reflex/utils/exceptions.py
index bceadc977..339abcda1 100644
--- a/reflex/utils/exceptions.py
+++ b/reflex/utils/exceptions.py
@@ -163,6 +163,10 @@ class StateSerializationError(ReflexError):
     """Raised when the state cannot be serialized."""
 
 
+class StateMismatchError(ReflexError, ValueError):
+    """Raised when the state retrieved does not match the expected state."""
+
+
 class SystemPackageMissingError(ReflexError):
     """Raised when a system package is missing."""