Include original pickle error message for better debugging

When the pickling throws a warning and dill is not installed, include the
original pickle error.

Add a test case for an object that even dill cannot pickle to ensure error path
is hit as expected.
This commit is contained in:
Masen Furer 2024-10-24 14:08:24 -07:00
parent 6ec1809c29
commit cc6beabc07
No known key found for this signature in database
GPG Key ID: B0008AD22B3B3A95
2 changed files with 13 additions and 2 deletions

View File

@ -2053,7 +2053,7 @@ class BaseState(Base, ABC, extra=pydantic.Extra.allow):
"""
try:
return pickle.dumps((self._to_schema(), self))
except (pickle.PicklingError, AttributeError):
except (pickle.PicklingError, AttributeError) as og_pickle_error:
error = (
f"Failed to serialize state {self.get_full_name()} due to unpicklable object. "
"This state will not be persisted. "
@ -2063,7 +2063,10 @@ class BaseState(Base, ABC, extra=pydantic.Extra.allow):
return dill.dumps((self._to_schema(), self))
except ImportError:
error += "Consider `pip install 'dill>=0.3.8'` for more exotic serialization support. "
error += (
f"Pickle error: {og_pickle_error}. "
"Consider `pip install 'dill>=0.3.8'` for more exotic serialization support."
)
except (pickle.PicklingError, TypeError, ValueError) as ex:
error += f"Dill was also unable to pickle the state: {ex}"
console.warn(error)

View File

@ -3360,6 +3360,7 @@ def test_fallback_pickle():
class DillState(BaseState):
_o: Obj | None = None
_f: Callable | None = None
_g: Any = None
state = DillState(_reflex_internal_init=True) # type: ignore
state._o = Obj(_f=lambda: 42)
@ -3370,3 +3371,10 @@ def test_fallback_pickle():
unpickled_state = BaseState._deserialize(pk)
assert unpickled_state._f() == 420
assert unpickled_state._o._f() == 42
# Some object, like generator, are still unpicklable with dill.
state._g = (i for i in range(10))
pk = state._serialize()
assert len(pk) == 0
with pytest.raises(EOFError):
BaseState._deserialize(pk)