From 307b487b4abf266cd731ba9d8d2444cfebebd17d Mon Sep 17 00:00:00 2001
From: samtin0x <40127309+samtin0x@users.noreply.github.com>
Date: Thu, 11 Jul 2024 20:38:36 +0200
Subject: [PATCH] Custom json decoder for bytearray serialization (#192)

* Implement custom json decoder for bytearray serialization

* Update subaccount examples with custom JSON decoding

* Rename Node client custom encoder method
---
 v4-client-py-v2/dydx_v4_client/node/client.py | 43 ++++++++++++++++++-
 .../examples/validator_get_example.py         |  6 ++-
 2 files changed, 46 insertions(+), 3 deletions(-)

diff --git a/v4-client-py-v2/dydx_v4_client/node/client.py b/v4-client-py-v2/dydx_v4_client/node/client.py
index 95742ef3..6d3098c0 100644
--- a/v4-client-py-v2/dydx_v4_client/node/client.py
+++ b/v4-client-py-v2/dydx_v4_client/node/client.py
@@ -1,7 +1,11 @@
+import base64
+import json
 from dataclasses import dataclass
+from typing import Union, Dict, Any
 
 import grpc
-from google.protobuf.message import Message
+from google._upb._message import Message
+from google.protobuf.json_format import MessageToDict, MessageToJson
 from typing_extensions import List, Optional, Self
 from v4_proto.cosmos.auth.v1beta1 import query_pb2_grpc as auth
 from v4_proto.cosmos.auth.v1beta1.auth_pb2 import BaseAccount
@@ -77,10 +81,47 @@
 from dydx_v4_client.wallet import Wallet
 
 
+class CustomJSONDecoder:
+    def __init__(self):
+        self.decoder = json.JSONDecoder(object_hook=self.decode_dict)
+
+    def decode(self, json_string):
+        return self.decoder.decode(json_string)
+
+    @staticmethod
+    def decode_base64(value):
+        if isinstance(value, str):
+            try:
+                return list(base64.b64decode(value))
+            except (base64.binascii.Error, ValueError):
+                return value
+        return value
+
+    def decode_dict(self, data):
+        if isinstance(data, dict):
+            return {k: self.decode_base64(v) for k, v in data.items()}
+        return data
+
+
 @dataclass
 class QueryNodeClient:
     channel: grpc.Channel
 
+    @staticmethod
+    def transcode_response(response: Message) -> Union[Dict[str, Any], List[Any]]:
+        """
+        Encodes the response using the custom JSON encoder.
+
+        Args:
+            response (Message): The response message to encode.
+
+        Returns:
+            Union[Dict[str, Any], List[Any]]: The encoded response.
+        """
+        response_dict = MessageToDict(response)
+        json_string = json.dumps(response_dict)
+        return CustomJSONDecoder().decode(json_string)
+
     async def get_account_balances(
         self, address: str
     ) -> bank_query.QueryAllBalancesResponse:
diff --git a/v4-client-py-v2/examples/validator_get_example.py b/v4-client-py-v2/examples/validator_get_example.py
index 5ec72179..2c4cefcc 100644
--- a/v4-client-py-v2/examples/validator_get_example.py
+++ b/v4-client-py-v2/examples/validator_get_example.py
@@ -66,16 +66,18 @@ async def test():
 
     try:
         subaccount = await node_client.get_subaccount(TEST_ADDRESS, 0)
+        decoded = node_client.transcode_response(subaccount)
         print("Get Subaccount:")
-        print(subaccount)
+        print(decoded)
     except Exception as e:
         print("Error in get_subaccount:")
         print(f"Error: {e}")
 
     try:
         subaccounts = await node_client.get_subaccounts()
+        decoded = node_client.transcode_response(subaccounts)
         print("Get Subaccounts:")
-        print(subaccounts)
+        print(decoded)
     except Exception as e:
         print("Error in get_subaccounts:")
         print(f"Error: {e}")