From 163acf70a280e63c59237856b6ccf1811bd0e34b Mon Sep 17 00:00:00 2001
From: Masen Furer <m_github@0x26.net>
Date: Mon, 4 Nov 2024 10:11:51 -0800
Subject: [PATCH] [ENG-759] [ENG-1104] patch `json.dumps` to handle
 `__wrapped__` objects (#4166)

* [ENG-1104] patch `json.dumps` to handle `__wrapped__` objects

Unwrap proxied objects if the default serializer doesn't work.

* pre-commit fixup

* Skip default fallback logic when `cls` is specified

`cls` will provide its own default serialization mechanism, passing a `cls`
Encoder class is now also a way to opt-out of our patching shenanigans and just
use your own code.

This will work, provided the library doing the JSON encoding isn't also using
their own custom class.

* Override JSONEncoder.default instead of json.dumps

Many libraries (like httpx, used by openai), will use `from json import dumps`,
and if they do that before `reflex.state` gets imported, then they will get the
original dumps function instead of our patched one.

To workaround this, monkeypatch the `JSONEncoder.default` function instead.

This is also nicer behavior for custom subclasses of JSONEncoder; if someone
wants to opt-out of our monkeypatching, they can simply not call
`super().default(o)` in their subclass, which by default only raises a
TypeError.

---------

Co-authored-by: Nikhil Rao <nikhil@reflex.dev>
---
 reflex/state.py | 24 ++++++++++++++++++++++++
 1 file changed, 24 insertions(+)

diff --git a/reflex/state.py b/reflex/state.py
index ee0aa3e93..44ccebc30 100644
--- a/reflex/state.py
+++ b/reflex/state.py
@@ -8,6 +8,7 @@ import copy
 import dataclasses
 import functools
 import inspect
+import json
 import pickle
 import sys
 import uuid
@@ -3713,6 +3714,29 @@ def serialize_mutable_proxy(mp: MutableProxy):
     return mp.__wrapped__
 
 
+_orig_json_JSONEncoder_default = json.JSONEncoder.default
+
+
+def _json_JSONEncoder_default_wrapper(self: json.JSONEncoder, o: Any) -> Any:
+    """Wrap JSONEncoder.default to handle MutableProxy objects.
+
+    Args:
+        self: the JSONEncoder instance.
+        o: the object to serialize.
+
+    Returns:
+        A JSON-able object.
+    """
+    try:
+        return o.__wrapped__
+    except AttributeError:
+        pass
+    return _orig_json_JSONEncoder_default(self, o)
+
+
+json.JSONEncoder.default = _json_JSONEncoder_default_wrapper
+
+
 class ImmutableMutableProxy(MutableProxy):
     """A proxy for a mutable object that tracks changes.