From 1a1fff07e9c8d62fd9c5d303ce76872d0df6fa5f Mon Sep 17 00:00:00 2001
From: Masen Furer <m_github@0x26.net>
Date: Thu, 31 Oct 2024 14:08:46 -0700
Subject: [PATCH] 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.
---
 reflex/state.py | 37 ++++++++++++-------------------------
 1 file changed, 12 insertions(+), 25 deletions(-)

diff --git a/reflex/state.py b/reflex/state.py
index 453e17e30..84465ddb0 100644
--- a/reflex/state.py
+++ b/reflex/state.py
@@ -3716,40 +3716,27 @@ def serialize_mutable_proxy(mp: MutableProxy):
     return mp.__wrapped__
 
 
-_orig_json_dumps = json.dumps
+_orig_json_JSONEncoder_default = json.JSONEncoder.default
 
 
-def _json_dumps_wrapper(*args, **kwargs):
-    """Wrap json.dumps to handle MutableProxy objects.
+def _json_JSONEncoder_default_wrapper(self: json.JSONEncoder, o: Any) -> Any:
+    """Wrap JSONEncoder.default to handle MutableProxy objects.
 
     Args:
-        *args: args for json.dumps.
-        **kwargs: kwargs for json.dumps.
+        self: the JSONEncoder instance.
+        o: the object to serialize.
 
     Returns:
-        The JSON string.
+        A JSON-able object.
     """
-    if "cls" not in kwargs:
-        _orig_default = kwargs.pop("default", None)
-
-    def _default(obj):
-        if _orig_default is not None:
-            try:
-                return _orig_default(obj)
-            except TypeError:
-                pass
-        try:
-            return obj.__wrapped__
-        except AttributeError:
-            pass
-        raise TypeError()
-
-    if "cls" not in kwargs:
-        kwargs["default"] = _default
-    return _orig_json_dumps(*args, **kwargs)
+    try:
+        return o.__wrapped__
+    except AttributeError:
+        pass
+    return _orig_json_JSONEncoder_default(self, o)
 
 
-json.dumps = _json_dumps_wrapper
+json.JSONEncoder.default = _json_JSONEncoder_default_wrapper
 
 
 class ImmutableMutableProxy(MutableProxy):