raise StateSerializationError if the state cannot be serialized

This commit is contained in:
Benedikt Bartscher 2024-11-28 21:46:35 +01:00
parent 39cdce6960
commit 5be7ce8a0d
No known key found for this signature in database
2 changed files with 14 additions and 0 deletions

View File

@ -97,6 +97,7 @@ from reflex.utils.exceptions import (
ReflexRuntimeError, ReflexRuntimeError,
SetUndefinedStateVarError, SetUndefinedStateVarError,
StateSchemaMismatchError, StateSchemaMismatchError,
StateSerializationError,
StateTooLargeError, StateTooLargeError,
) )
from reflex.utils.exec import is_testing_env from reflex.utils.exec import is_testing_env
@ -2177,8 +2178,12 @@ class BaseState(Base, ABC, extra=pydantic.Extra.allow):
Returns: Returns:
The serialized state. The serialized state.
Raises:
StateSerializationError: If the state cannot be serialized.
""" """
payload = b"" payload = b""
error = ""
try: try:
payload = pickle.dumps((self._to_schema(), self)) payload = pickle.dumps((self._to_schema(), self))
except HANDLED_PICKLE_ERRORS as og_pickle_error: except HANDLED_PICKLE_ERRORS as og_pickle_error:
@ -2198,8 +2203,13 @@ class BaseState(Base, ABC, extra=pydantic.Extra.allow):
except HANDLED_PICKLE_ERRORS as ex: except HANDLED_PICKLE_ERRORS as ex:
error += f"Dill was also unable to pickle the state: {ex}" error += f"Dill was also unable to pickle the state: {ex}"
console.warn(error) console.warn(error)
if environment.REFLEX_PERF_MODE.get() != PerformanceMode.OFF: if environment.REFLEX_PERF_MODE.get() != PerformanceMode.OFF:
self._check_state_size(len(payload)) self._check_state_size(len(payload))
if not payload:
raise StateSerializationError(error)
return payload return payload
@classmethod @classmethod

View File

@ -155,6 +155,10 @@ class StateTooLargeError(ReflexError):
"""Raised when the state is too large to be serialized.""" """Raised when the state is too large to be serialized."""
class StateSerializationError(ReflexError):
"""Raised when the state cannot be serialized."""
class SystemPackageMissingError(ReflexError): class SystemPackageMissingError(ReflexError):
"""Raised when a system package is missing.""" """Raised when a system package is missing."""