From ea90a3ebfa05b05e0906f613803ffcc153935083 Mon Sep 17 00:00:00 2001
From: benedikt-bartscher
 <31854409+benedikt-bartscher@users.noreply.github.com>
Date: Thu, 12 Dec 2024 11:07:59 +0100
Subject: [PATCH] minor pickle improvement (#4499)

---
 reflex/state.py | 16 ++++++++++++++--
 1 file changed, 14 insertions(+), 2 deletions(-)

diff --git a/reflex/state.py b/reflex/state.py
index c73b1ac46..b4daa3e69 100644
--- a/reflex/state.py
+++ b/reflex/state.py
@@ -2122,14 +2122,26 @@ class BaseState(Base, ABC, extra=pydantic.Extra.allow):
             state["__dict__"].pop("router", None)
             state["__dict__"].pop("router_data", None)
         # Never serialize parent_state or substates.
-        state["__dict__"]["parent_state"] = None
-        state["__dict__"]["substates"] = {}
+        state["__dict__"].pop("parent_state", None)
+        state["__dict__"].pop("substates", None)
         state["__dict__"].pop("_was_touched", None)
         # Remove all inherited vars.
         for inherited_var_name in self.inherited_vars:
             state["__dict__"].pop(inherited_var_name, None)
         return state
 
+    def __setstate__(self, state: dict[str, Any]):
+        """Set the state from redis deserialization.
+
+        This method is called by pickle to deserialize the object.
+
+        Args:
+            state: The state dict for deserialization.
+        """
+        state["__dict__"]["parent_state"] = None
+        state["__dict__"]["substates"] = {}
+        super().__setstate__(state)
+
     def _check_state_size(
         self,
         pickle_state_size: int,