Skip to content

Commit

Permalink
raise StateSerializationError if the state cannot be serialized (#4453)
Browse files Browse the repository at this point in the history
* raise StateSerializationError if the state cannot be serialized

* fix test
  • Loading branch information
benedikt-bartscher authored Dec 11, 2024
1 parent 862d7ec commit 2ee201b
Show file tree
Hide file tree
Showing 3 changed files with 22 additions and 3 deletions.
10 changes: 10 additions & 0 deletions reflex/state.py
Original file line number Diff line number Diff line change
Expand Up @@ -97,6 +97,7 @@
ReflexRuntimeError,
SetUndefinedStateVarError,
StateSchemaMismatchError,
StateSerializationError,
StateTooLargeError,
)
from reflex.utils.exec import is_testing_env
Expand Down Expand Up @@ -2193,8 +2194,12 @@ def _serialize(self) -> bytes:
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:
Expand All @@ -2214,8 +2219,13 @@ def _serialize(self) -> bytes:
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
Expand Down
4 changes: 4 additions & 0 deletions reflex/utils/exceptions.py
Original file line number Diff line number Diff line change
Expand Up @@ -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."""

Expand Down
11 changes: 8 additions & 3 deletions tests/units/test_state.py
Original file line number Diff line number Diff line change
Expand Up @@ -55,7 +55,11 @@
)
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
Expand Down Expand Up @@ -3433,8 +3437,9 @@ class DillState(BaseState):
# 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:
Expand Down

0 comments on commit 2ee201b

Please sign in to comment.