raise StateSerializationError if the state cannot be serialized (#4453)

* raise StateSerializationError if the state cannot be serialized

* fix test
This commit is contained in:
benedikt-bartscher 2024-12-11 20:08:18 +01:00 committed by GitHub
parent 862d7ec807
commit 2ee201b520
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
3 changed files with 22 additions and 3 deletions

View File

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

View File

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

View File

@ -55,7 +55,11 @@ from reflex.state import (
)
from reflex.testing import chdir
from reflex.utils import format, prerequisites, types
from reflex.utils.exceptions import ReflexRuntimeError, SetUndefinedStateVarError
from reflex.utils.exceptions import (
ReflexRuntimeError,
SetUndefinedStateVarError,
StateSerializationError,
)
from reflex.utils.format import json_dumps
from reflex.vars.base import Var, computed_var
from tests.units.states.mutation import MutableSQLAModel, MutableTestState
@ -3433,8 +3437,9 @@ def test_fallback_pickle():
# Some object, like generator, are still unpicklable with dill.
state3 = DillState(_reflex_internal_init=True) # type: ignore
state3._g = (i for i in range(10))
pk3 = state3._serialize()
assert len(pk3) == 0
with pytest.raises(StateSerializationError):
_ = state3._serialize()
def test_typed_state() -> None: