From 953495775d52ecc2c41c3b84ed34a11ae7183ef6 Mon Sep 17 00:00:00 2001
From: Masen Furer <m_github@0x26.net>
Date: Thu, 22 Feb 2024 10:14:14 -0800
Subject: [PATCH] [REF-2045] Implement __reduce_ex__ for MutableProxy  (#2688)

* test_state: augment modify_state test for writing MutableProxy

If the object contains a MutableProxy inside of it, then we get a pickling
error.

* Implement __reduce_ex__ for MutableProxy

Pass through `__reduce_ex__` onto the wrapped instance to strip it off when
cloudpickling to redis.

* base: get_value actually works with a str key

Unless the key isn't a field on the model, then it falls back to the previous
behavior of just returning the given key as is... why does it do this? I don't
know.
---
 reflex/base.py      |  4 ++++
 reflex/state.py     | 15 +++++++++++++++
 tests/test_state.py |  6 +++++-
 3 files changed, 24 insertions(+), 1 deletion(-)

diff --git a/reflex/base.py b/reflex/base.py
index 9efea60e0..36b150c5c 100644
--- a/reflex/base.py
+++ b/reflex/base.py
@@ -115,6 +115,10 @@ class Base(pydantic.BaseModel):
         Returns:
             The value of the field.
         """
+        if isinstance(key, str) and key in self.__fields__:
+            # Seems like this function signature was wrong all along?
+            # If the user wants a field that we know of, get it and pass it off to _get_value
+            key = getattr(self, key)
         return self._get_value(
             key,
             to_dict=True,
diff --git a/reflex/state.py b/reflex/state.py
index 7b377b4d0..f6bf2161b 100644
--- a/reflex/state.py
+++ b/reflex/state.py
@@ -2363,6 +2363,21 @@ class MutableProxy(wrapt.ObjectProxy):
         """
         return copy.deepcopy(self.__wrapped__, memo=memo)
 
+    def __reduce_ex__(self, protocol_version):
+        """Get the state for redis serialization.
+
+        This method is called by cloudpickle to serialize the object.
+
+        It explicitly serializes the wrapped object, stripping off the mutable proxy.
+
+        Args:
+            protocol_version: The protocol version.
+
+        Returns:
+            Tuple of (wrapped class, empty args, class __getstate__)
+        """
+        return self.__wrapped__.__reduce_ex__(protocol_version)
+
 
 @serializer
 def serialize_mutable_proxy(mp: MutableProxy) -> SerializedType:
diff --git a/tests/test_state.py b/tests/test_state.py
index dbcc1afce..565f1945b 100644
--- a/tests/test_state.py
+++ b/tests/test_state.py
@@ -1457,12 +1457,16 @@ async def test_state_manager_modify_state(
         token: A token.
         substate_token: A token + substate name for looking up in state manager.
     """
-    async with state_manager.modify_state(substate_token):
+    async with state_manager.modify_state(substate_token) as state:
         if isinstance(state_manager, StateManagerRedis):
             assert await state_manager.redis.get(f"{token}_lock")
         elif isinstance(state_manager, StateManagerMemory):
             assert token in state_manager._states_locks
             assert state_manager._states_locks[token].locked()
+        # Should be able to write proxy objects inside mutables
+        complex_1 = state.complex[1]
+        assert isinstance(complex_1, MutableProxy)
+        state.complex[3] = complex_1
     # lock should be dropped after exiting the context
     if isinstance(state_manager, StateManagerRedis):
         assert (await state_manager.redis.get(f"{token}_lock")) is None