MutableProxy: wrap dataclass and BaseModel methods (#5979)

When calling a method on an arbitrary wrapped object, rebind it's `self` as the
mutable proxy so changes made inside the method are also tracked.

(Previously only wrapped `Base` instances had in-method tracking).

Author: masenf

reflex/istate/proxy.py

@@ -553,18 +553,17 @@ def __getattr__(self, __name: str) -> Any:
                 value = wrapt.FunctionWrapper(value, self._mark_dirty)
 
             if __name in self.__wrap_mutable_attrs__:
-                # Wrap methods that may return mutable objects tied to the state.
+                # Wrap special methods that may return mutable objects tied to the state.
                 value = wrapt.FunctionWrapper(
                     value,
                     self._wrap_recursive_decorator,  # pyright: ignore[reportArgumentType]
                 )
 
             if (
-                isinstance(self.__wrapped__, Base)
-                and __name not in NEVER_WRAP_BASE_ATTRS
-                and hasattr(value, "__func__")
-            ):
-                # Wrap methods called on Base subclasses, which might do _anything_
+                not isinstance(self.__wrapped__, Base)
+                or __name not in NEVER_WRAP_BASE_ATTRS
+            ) and hasattr(value, "__func__"):
+                # Wrap methods which might do _anything_
                 return wrapt.FunctionWrapper(
                     functools.partial(value.__func__, self),  # pyright: ignore [reportFunctionMemberAccess, reportAttributeAccessIssue]
                     self._wrap_recursive_decorator,  # pyright: ignore[reportArgumentType]

tests/units/test_state.py

@@ -2055,6 +2055,22 @@ class ModelDC:
     foo: str = "bar"
     ls: list[dict] = dataclasses.field(default_factory=list)
 
+    def set_foo(self, val: str):
+        """Set the attribute foo.
+
+        Args:
+            val: The value to set.
+        """
+        self.foo = val
+
+    def double_foo(self) -> str:
+        """Concatenate foo with foo.
+
+        Returns:
+            foo + foo
+        """
+        return self.foo + self.foo
+
 
 @pytest.mark.asyncio
 async def test_state_proxy(
@@ -3806,12 +3822,44 @@ class ModelV1(BaseModelV1):
 
     foo: str = "bar"
 
+    def set_foo(self, val: str):
+        """Set the attribute foo.
+
+        Args:
+            val: The value to set.
+        """
+        self.foo = val
+
+    def double_foo(self) -> str:
+        """Concatenate foo with foo.
+
+        Returns:
+            foo + foo
+        """
+        return self.foo + self.foo
+
 
 class ModelV2(BaseModelV2):
     """A pydantic BaseModel v2."""
 
     foo: str = "bar"
 
+    def set_foo(self, val: str):
+        """Set the attribute foo.
+
+        Args:
+            val: The value to set.
+        """
+        self.foo = val
+
+    def double_foo(self) -> str:
+        """Concatenate foo with foo.
+
+        Returns:
+            foo + foo
+        """
+        return self.foo + self.foo
+
 
 class PydanticState(rx.State):
     """A state with pydantic BaseModel vars."""
@@ -3828,26 +3876,46 @@ def test_mutable_models():
     state.v1.foo = "baz"
     assert state.dirty_vars == {"v1"}
     state.dirty_vars.clear()
+    state.v1.set_foo("quuc")
+    assert state.dirty_vars == {"v1"}
+    state.dirty_vars.clear()
+    assert state.v1.double_foo() == "quucquuc"
+    assert state.dirty_vars == set()
+    state.v1.copy(update={"foo": "larp"})
+    assert state.dirty_vars == set()
 
     assert isinstance(state.v2, MutableProxy)
     state.v2.foo = "baz"
     assert state.dirty_vars == {"v2"}
     state.dirty_vars.clear()
+    state.v2.set_foo("quuc")
+    assert state.dirty_vars == {"v2"}
+    state.dirty_vars.clear()
+    assert state.v2.double_foo() == "quucquuc"
+    assert state.dirty_vars == set()
+    state.v2.model_copy(update={"foo": "larp"})
+    assert state.dirty_vars == set()
 
     assert isinstance(state.dc, MutableProxy)
     state.dc.foo = "baz"
     assert state.dirty_vars == {"dc"}
     state.dirty_vars.clear()
     assert state.dirty_vars == set()
+    state.dc.set_foo("quuc")
+    assert state.dirty_vars == {"dc"}
+    state.dirty_vars.clear()
+    assert state.dirty_vars == set()
+    assert state.dc.double_foo() == "quucquuc"
+    assert state.dirty_vars == set()
     state.dc.ls.append({"hi": "reflex"})
     assert state.dirty_vars == {"dc"}
     state.dirty_vars.clear()
     assert state.dirty_vars == set()
-    assert dataclasses.asdict(state.dc) == {"foo": "baz", "ls": [{"hi": "reflex"}]}
-    assert dataclasses.astuple(state.dc) == ("baz", [{"hi": "reflex"}])
+    assert dataclasses.asdict(state.dc) == {"foo": "quuc", "ls": [{"hi": "reflex"}]}
+    assert dataclasses.astuple(state.dc) == ("quuc", [{"hi": "reflex"}])
     # creating a new instance shouldn't mark the state dirty
-    assert dataclasses.replace(state.dc, foo="quuc") == ModelDC(
-        foo="quuc", ls=[{"hi": "reflex"}]
+    assert dataclasses.replace(state.dc, foo="larp") == ModelDC(
+        foo="larp", ls=[{"hi": "reflex"}]
     )
     assert state.dirty_vars == set()