From 4ad69cba548d1357cc715bfdcd89cd490f204f34 Mon Sep 17 00:00:00 2001
From: Benjamin Himes <benjamindhimes@gmail.com>
Date: Wed, 22 Jan 2025 20:58:12 +0200
Subject: [PATCH] Remove ujson as requirement. Allow users to use whatever json
 module they have installed.

---
 async_substrate_interface/async_substrate.py |  7 +++----
 async_substrate_interface/sync_substrate.py  |  7 +++----
 async_substrate_interface/utils/__init__.py  | 18 ++++++++++++++++++
 pyproject.toml                               |  1 -
 4 files changed, 24 insertions(+), 9 deletions(-)

diff --git a/async_substrate_interface/async_substrate.py b/async_substrate_interface/async_substrate.py
index 9992f7a..b1f33a6 100644
--- a/async_substrate_interface/async_substrate.py
+++ b/async_substrate_interface/async_substrate.py
@@ -26,7 +26,6 @@
 from bt_decode import PortableRegistry, decode as decode_by_type_string, MetadataV15
 from scalecodec.base import ScaleBytes, ScaleType, RuntimeConfigurationObject
 from scalecodec.types import GenericCall, GenericRuntimeCallDefinition, GenericExtrinsic
-import ujson
 from websockets.asyncio.client import connect
 from websockets.exceptions import ConnectionClosed
 
@@ -43,7 +42,7 @@
     SubstrateMixin,
     Preprocessed,
 )
-from async_substrate_interface.utils import hex_to_bytes
+from async_substrate_interface.utils import hex_to_bytes, json
 from async_substrate_interface.utils.storage import StorageKey
 
 if TYPE_CHECKING:
@@ -575,7 +574,7 @@ async def shutdown(self):
     async def _recv(self) -> None:
         try:
             # TODO consider wrapping this in asyncio.wait_for and use that for the timeout logic
-            response = ujson.loads(await self.ws.recv(decode=False))
+            response = json.loads(await self.ws.recv(decode=False))
             self.last_received = time.time()
             async with self._lock:
                 # note that these 'subscriptions' are all waiting sent messages which have not received
@@ -617,7 +616,7 @@ async def send(self, payload: dict) -> int:
         self.id += 1
         # self._open_subscriptions += 1
         try:
-            await self.ws.send(ujson.dumps({**payload, **{"id": original_id}}))
+            await self.ws.send(json.dumps({**payload, **{"id": original_id}}))
             return original_id
         except (ConnectionClosed, ssl.SSLError, EOFError):
             async with self._lock:
diff --git a/async_substrate_interface/sync_substrate.py b/async_substrate_interface/sync_substrate.py
index 8773be1..196d83a 100644
--- a/async_substrate_interface/sync_substrate.py
+++ b/async_substrate_interface/sync_substrate.py
@@ -8,7 +8,6 @@
 from bt_decode import PortableRegistry, decode as decode_by_type_string, MetadataV15
 from scalecodec import GenericExtrinsic, GenericCall, GenericRuntimeCallDefinition
 from scalecodec.base import RuntimeConfigurationObject, ScaleBytes, ScaleType
-import ujson
 from websockets.sync.client import connect
 
 from async_substrate_interface.errors import (
@@ -24,7 +23,7 @@
     Preprocessed,
     ScaleObj,
 )
-from async_substrate_interface.utils import hex_to_bytes
+from async_substrate_interface.utils import hex_to_bytes, json
 from async_substrate_interface.utils.storage import StorageKey
 
 
@@ -1633,12 +1632,12 @@ def _make_rpc_request(
             item_id = 0
             for payload in payloads:
                 item_id += 1
-                ws.send(ujson.dumps({**payload["payload"], **{"id": item_id}}))
+                ws.send(json.dumps({**payload["payload"], **{"id": item_id}}))
                 request_manager.add_request(item_id, payload["id"])
 
             while True:
                 try:
-                    response = ujson.loads(
+                    response = json.loads(
                         ws.recv(timeout=self.retry_timeout, decode=False)
                     )
                 except TimeoutError:
diff --git a/async_substrate_interface/utils/__init__.py b/async_substrate_interface/utils/__init__.py
index f4644f8..40e26c5 100644
--- a/async_substrate_interface/utils/__init__.py
+++ b/async_substrate_interface/utils/__init__.py
@@ -1,3 +1,6 @@
+import importlib
+
+
 def hex_to_bytes(hex_str: str) -> bytes:
     """
     Converts a hex-encoded string into bytes. Handles 0x-prefixed and non-prefixed hex-encoded strings.
@@ -7,3 +10,18 @@ def hex_to_bytes(hex_str: str) -> bytes:
     else:
         bytes_result = bytes.fromhex(hex_str)
     return bytes_result
+
+
+def import_json_lib():
+    libs = ["ujson", "orjson", "simplejson", "json"]
+
+    for lib in libs:
+        try:
+            return importlib.import_module(lib)
+        except ImportError:
+            continue
+
+    raise ImportError("None of the specified JSON libraries are installed.")
+
+
+json = import_json_lib()
diff --git a/pyproject.toml b/pyproject.toml
index ed94536..6f82643 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -12,7 +12,6 @@ dependencies = [
   "bittensor-wallet>=2.1.3",
   "bt-decode==0.4.0",
   "scalecodec~=1.2.11",
-  "ujson",
   "websockets>=14.1",
   "xxhash"
 ]