[ENG-3970] When normal pickle fails, try dill

If dill is not installed, suggest that the user `pip install` it.

Fix #4147
This commit is contained in:
Masen Furer 2024-10-24 13:52:49 -07:00
parent 9d29e7f3ee
commit eec30984ec
No known key found for this signature in database
GPG Key ID: B0008AD22B3B3A95
3 changed files with 38 additions and 4 deletions

View File

@ -65,6 +65,7 @@ pytest = ">=7.1.2,<9.0"
pytest-mock = ">=3.10.0,<4.0"
pyright = ">=1.1.229,<1.1.335"
darglint = ">=1.8.1,<2.0"
dill = ">=0.3.8"
toml = ">=0.10.2,<1.0"
pytest-asyncio = ">=0.24.0"
pytest-cov = ">=4.0.0,<6.0"

View File

@ -2053,12 +2053,21 @@ class BaseState(Base, ABC, extra=pydantic.Extra.allow):
"""
try:
return pickle.dumps((self._to_schema(), self))
except pickle.PicklingError:
console.warn(
except (pickle.PicklingError, AttributeError):
error = (
f"Failed to serialize state {self.get_full_name()} due to unpicklable object. "
"This state will not be persisted."
"This state will not be persisted. "
)
return b""
try:
import dill
return dill.dumps((self._to_schema(), self))
except ImportError:
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)
return b""
@classmethod
def _deserialize(

View File

@ -3346,3 +3346,27 @@ async def test_deserialize_gc_state_disk(token):
assert s.num == 43
c = await root.get_state(Child)
assert c.foo == "bar"
class Obj(Base):
"""A object containing a callable for testing fallback pickle."""
_f: Callable
def test_fallback_pickle():
"""Test that state serialization will fall back to dill."""
class DillState(BaseState):
_o: Obj | None = None
_f: Callable | None = None
state = DillState(_reflex_internal_init=True) # type: ignore
state._o = Obj(_f=lambda: 42)
state._f = lambda: 420
pk = state._serialize()
unpickled_state = BaseState._deserialize(pk)
assert unpickled_state._f() == 420
assert unpickled_state._o._f() == 42