Add in selective type conversion

Author: thewhaleking

async_substrate_interface/async_substrate.py

@@ -734,6 +734,7 @@ def __init__(
         _mock: bool = False,
         _log_raw_websockets: bool = False,
         ws_shutdown_timer: float = 5.0,
+        decode_ss58: bool = False,
     ):
         """
         The asyncio-compatible version of the subtensor interface commands we use in bittensor. It is important to
@@ -753,10 +754,15 @@ def __init__(
             _mock: whether to use mock version of the subtensor interface
             _log_raw_websockets: whether to log raw websocket requests during RPC requests
             ws_shutdown_timer: how long after the last connection your websocket should close
+            decode_ss58: Whether to decode AccountIds to SS58 or leave them in raw bytes tuples.
 
         """
         super().__init__(
-            type_registry, type_registry_preset, use_remote_preset, ss58_format
+            type_registry,
+            type_registry_preset,
+            use_remote_preset,
+            ss58_format,
+            decode_ss58,
         )
         self.max_retries = max_retries
         self.retry_timeout = retry_timeout
@@ -996,12 +1002,15 @@ async def decode_scale(
                 runtime = await self.init_runtime(block_hash=block_hash)
             if runtime.metadata_v15 is not None or force_legacy is True:
                 obj = decode_by_type_string(type_string, runtime.registry, scale_bytes)
-                try:
-                    type_str_int = int(type_string.split("::")[1])
-                    decoded_type_str = runtime.type_id_to_name[type_str_int]
-                    obj = convert_account_ids(obj, decoded_type_str)
-                except (ValueError, KeyError):
-                    pass
+                if self.decode_ss58:
+                    try:
+                        type_str_int = int(type_string.split("::")[1])
+                        decoded_type_str = runtime.type_id_to_name[type_str_int]
+                        obj = convert_account_ids(
+                            obj, decoded_type_str, runtime.ss58_format
+                        )
+                    except (ValueError, KeyError):
+                        pass
             else:
                 obj = legacy_scale_decode(type_string, scale_bytes, runtime)
         if return_scale_obj:
@@ -3480,6 +3489,7 @@ async def query_map(
                     value_type,
                     key_hashers,
                     ignore_decoding_errors,
+                    self.decode_ss58,
                 )
         return AsyncQueryMapResult(
             records=result,

async_substrate_interface/sync_substrate.py

@@ -44,6 +44,7 @@
     _bt_decode_to_dict_or_list,
     decode_query_map,
     legacy_scale_decode,
+    convert_account_ids,
 )
 from async_substrate_interface.utils.storage import StorageKey
 from async_substrate_interface.type_registry import _TYPE_REGISTRY
@@ -486,6 +487,7 @@ def __init__(
         retry_timeout: float = 60.0,
         _mock: bool = False,
         _log_raw_websockets: bool = False,
+        decode_ss58: bool = False,
     ):
         """
         The sync compatible version of the subtensor interface commands we use in bittensor. Use this instance only
@@ -503,10 +505,15 @@ def __init__(
             retry_timeout: how to long wait since the last ping to retry the RPC request
             _mock: whether to use mock version of the subtensor interface
             _log_raw_websockets: whether to log raw websocket requests during RPC requests
+            decode_ss58: Whether to decode AccountIds to SS58 or leave them in raw bytes tuples.
 
         """
         super().__init__(
-            type_registry, type_registry_preset, use_remote_preset, ss58_format
+            type_registry,
+            type_registry_preset,
+            use_remote_preset,
+            ss58_format,
+            decode_ss58,
         )
         self.max_retries = max_retries
         self.retry_timeout = retry_timeout
@@ -560,6 +567,7 @@ def initialize(self):
                 )
                 if ss58_prefix_constant:
                     self.ss58_format = ss58_prefix_constant.value
+                    self.runtime.ss58_format = ss58_prefix_constant.value
         self.initialized = True
 
     def __exit__(self, exc_type, exc_val, exc_tb):
@@ -693,6 +701,15 @@ def decode_scale(
                 obj = decode_by_type_string(
                     type_string, self.runtime.registry, scale_bytes
                 )
+                if self.decode_ss58:
+                    try:
+                        type_str_int = int(type_string.split("::")[1])
+                        decoded_type_str = self.runtime.type_id_to_name[type_str_int]
+                        obj = convert_account_ids(
+                            obj, decoded_type_str, self.ss58_format
+                        )
+                    except (ValueError, KeyError):
+                        pass
             else:
                 obj = legacy_scale_decode(type_string, scale_bytes, self.runtime)
         if return_scale_obj:
@@ -3010,6 +3027,7 @@ def query_map(
                     value_type,
                     key_hashers,
                     ignore_decoding_errors,
+                    self.decode_ss58,
                 )
         return QueryMapResult(
             records=result,

async_substrate_interface/types.py

@@ -553,8 +553,10 @@ def __init__(
         type_registry_preset: Optional[str] = None,
         use_remote_preset: bool = False,
         ss58_format: Optional[int] = None,
+        decode_ss58: bool = False,
     ):
         # We load a very basic RuntimeConfigurationObject that is only used for the initial metadata decoding
+        self.decode_ss58 = decode_ss58
         self.runtime_config = RuntimeConfigurationObject(ss58_format=ss58_format)
         self.ss58_format = ss58_format
         self.runtime_config.update_type_registry(load_type_registry_preset(name="core"))

async_substrate_interface/utils/decoding.py

@@ -81,6 +81,7 @@ def decode_query_map(
     value_type,
     key_hashers,
     ignore_decoding_errors,
+    decode_ss58: bool = False,
 ):
     def concat_hash_len(key_hasher: str) -> int:
         """
@@ -129,25 +130,29 @@ def concat_hash_len(key_hasher: str) -> int:
             # strip key_hashers to use as item key
             if len(param_types) - len(params) == 1:
                 item_key = dk[1]
-                if kts[kts.index(", ") + 2 : kts.index(")")] == "scale_info::0":
-                    item_key = ss58_encode(bytes(item_key[0]), runtime.ss58_format)
+                if decode_ss58:
+                    if kts[kts.index(", ") + 2 : kts.index(")")] == "scale_info::0":
+                        item_key = ss58_encode(bytes(item_key[0]), runtime.ss58_format)
 
             else:
                 item_key = tuple(
                     dk[key + 1] for key in range(len(params), len(param_types) + 1, 2)
                 )
-                # TODO handle decoding here, but first figure out what triggers this
 
         except Exception as _:
             if not ignore_decoding_errors:
                 raise
             item_key = None
-        try:
-            value_type_str_int = int(vts.split("::")[1])
-            decoded_type_str = runtime.type_id_to_name[value_type_str_int]
-            item_value = convert_account_ids(dv, decoded_type_str, runtime.ss58_format)
-        except (ValueError, KeyError) as e:
-            item_value = dv
+        item_value = dv
+        if decode_ss58:
+            try:
+                value_type_str_int = int(vts.split("::")[1])
+                decoded_type_str = runtime.type_id_to_name[value_type_str_int]
+                item_value = convert_account_ids(
+                    dv, decoded_type_str, runtime.ss58_format
+                )
+            except (ValueError, KeyError):
+                pass
         result.append([item_key, ScaleObj(item_value)])
     return result