From eec30984ec783bde7b9a7616513e9d211bf75d97 Mon Sep 17 00:00:00 2001
From: Masen Furer <m_github@0x26.net>
Date: Thu, 24 Oct 2024 13:52:49 -0700
Subject: [PATCH] [ENG-3970] When normal pickle fails, try dill

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

Fix #4147
---
 pyproject.toml            |  1 +
 reflex/state.py           | 17 +++++++++++++----
 tests/units/test_state.py | 24 ++++++++++++++++++++++++
 3 files changed, 38 insertions(+), 4 deletions(-)

diff --git a/pyproject.toml b/pyproject.toml
index 93f3c5d50..2635e1156 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -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"
diff --git a/reflex/state.py b/reflex/state.py
index 37bc06360..e80013418 100644
--- a/reflex/state.py
+++ b/reflex/state.py
@@ -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(
diff --git a/tests/units/test_state.py b/tests/units/test_state.py
index ebfeeb72c..23d799f6d 100644
--- a/tests/units/test_state.py
+++ b/tests/units/test_state.py
@@ -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