use cyscale

async_substrate_interface/type_registry.py

@@ -1,22 +1,63 @@
-from typing import Union
 from collections import namedtuple
-from bt_decode import (
-    NeuronInfo,
-    NeuronInfoLite,
-    DelegateInfo,
-    StakeInfo,
-    SubnetHyperparameters,
-    SubnetInfo,
-    SubnetInfoV2,
-    encode,
-)
+from typing import TYPE_CHECKING, Union
+
 from scalecodec import ss58_decode
+from scalecodec.base import ScaleBytes, RuntimeConfigurationObject
+from scalecodec.type_registry import load_type_registry_preset
+
+if TYPE_CHECKING:
+    from async_substrate_interface.types import Runtime
+
+# Bittensor-specific types needed for legacy (pre-V15) blocks
+_BITTENSOR_LEGACY_TYPES = {
+    "types": {
+        "Balance": "u64",
+        "StakeInfo": {
+            "type": "struct",
+            "type_mapping": [
+                ["hotkey", "AccountId"],
+                ["coldkey", "AccountId"],
+                ["stake", "Compact<u64>"],
+            ],
+        },
+    }
+}
+
+_legacy_rc: RuntimeConfigurationObject | None = None
+
+
+def _get_legacy_rc(ss58_format: int | None = None) -> RuntimeConfigurationObject:
+    """Return a lazily-initialised RC with legacy + Bittensor types."""
+    global _legacy_rc
+    if _legacy_rc is None:
+        rc = RuntimeConfigurationObject(ss58_format=ss58_format)
+        rc.update_type_registry(load_type_registry_preset(name="legacy"))
+        rc.update_type_registry(_BITTENSOR_LEGACY_TYPES)
+        _legacy_rc = rc
+    return _legacy_rc
+
+
+def _legacy_decode(type_string: str, raw_bytes: bytes, ss58_format: int):
+    """Decode raw_bytes using the legacy scalecodec type registry."""
+    rc = _get_legacy_rc(ss58_format=ss58_format)
+    obj = rc.create_scale_object(type_string, data=ScaleBytes(raw_bytes))
+    return obj.decode()
+
+
+def _cyscale_decode(type_name: str, raw_bytes: bytes, runtime: "Runtime"):
+    """Decode raw_bytes as type_name using cyscale's portable registry."""
+    type_id = runtime.registry_type_map.get(type_name)
+    if type_id is None:
+        raise ValueError(f"Type '{type_name}' not found in registry_type_map")
+    return runtime.runtime_config.batch_decode([f"scale_info::{type_id}"], [raw_bytes])[
+        0
+    ]
 
 
 def stake_info_decode_vec_legacy_compatibility(
-    item,
+    raw_bytes: bytes, runtime: "Runtime"
 ) -> list[dict[str, Union[str, int, bytes, bool]]]:
-    stake_infos: list[StakeInfo] = StakeInfo.decode_vec(item)
+    ss58_format = runtime.ss58_format
     NewStakeInfo = namedtuple(
         "NewStakeInfo",
         [
@@ -30,21 +71,21 @@ def stake_info_decode_vec_legacy_compatibility(
             "is_registered",
         ],
     )
-    output = []
-    for stake_info in stake_infos:
-        output.append(
-            NewStakeInfo(
-                0,
-                stake_info.hotkey,
-                stake_info.coldkey,
-                stake_info.stake,
-                0,
-                0,
-                0,
-                False,
-            )
-        )
-    return output
+
+    stake_infos = _legacy_decode("Vec<StakeInfo>", raw_bytes, ss58_format=ss58_format)
+    return [
+        NewStakeInfo(
+            0,
+            si["hotkey"],
+            si["coldkey"],
+            si["stake"],
+            0,
+            0,
+            0,
+            False,
+        )._asdict()
+        for si in stake_infos
+    ]
 
 
 def preprocess_get_stake_info_for_coldkeys(addrs):
@@ -59,6 +100,45 @@ def preprocess_get_stake_info_for_coldkeys(addrs):
     return output
 
 
+def _encode_vec_u8(params, runtime: "Runtime") -> bytes:
+    """Encode a single SS58 address as Vec<u8>."""
+    addr = (
+        params
+        if isinstance(params, str)
+        else params[0]
+        if isinstance(params, list)
+        else params["coldkey_account"]
+    )
+    raw = bytes.fromhex(ss58_decode(addr))
+    return bytes(
+        runtime.runtime_config.create_scale_object("Vec<u8>").encode(list(raw)).data
+    )
+
+
+def _encode_vec_u8_coldkey(params, runtime: "Runtime") -> bytes:
+    """Encode a coldkey account as Vec<u8>, handling list or dict input."""
+    if isinstance(params, list):
+        addr = params[0]
+    elif isinstance(params, dict):
+        addr = params["coldkey_account"]
+    else:
+        addr = params
+    raw = bytes.fromhex(ss58_decode(addr))
+    return bytes(
+        runtime.runtime_config.create_scale_object("Vec<u8>").encode(list(raw)).data
+    )
+
+
+def _encode_vec_vec_u8(params, runtime: "Runtime") -> bytes:
+    """Encode multiple coldkey accounts as Vec<Vec<u8>>."""
+    preprocessed = preprocess_get_stake_info_for_coldkeys(params)
+    return bytes(
+        runtime.runtime_config.create_scale_object("Vec<Vec<u8>>")
+        .encode(preprocessed)
+        .data
+    )
+
+
 _TYPE_REGISTRY: dict[str, dict] = {
     "types": {
         "Balance": "u64",  # Need to override default u128
@@ -73,16 +153,18 @@ def preprocess_get_stake_info_for_coldkeys(addrs):
                             "type": "Vec<u8>",
                         },
                     ],
-                    "encoder": lambda addr, reg: encode(
-                        "Vec<u8>", reg, list(bytes.fromhex(ss58_decode(addr)))
-                    ),
+                    "encoder": _encode_vec_u8,
                     "type": "Vec<u8>",
-                    "decoder": DelegateInfo.decode_delegated,
+                    "decoder": lambda raw_bytes, runtime: _cyscale_decode(
+                        "Vec<DelegateInfo>", raw_bytes, runtime
+                    ),
                 },
                 "get_delegates": {
                     "params": [],
                     "type": "Vec<u8>",
-                    "decoder": DelegateInfo.decode_vec,
+                    "decoder": lambda raw_bytes, runtime: _cyscale_decode(
+                        "Vec<DelegateInfo>", raw_bytes, runtime
+                    ),
                 },
             }
         },
@@ -100,7 +182,9 @@ def preprocess_get_stake_info_for_coldkeys(addrs):
                         },
                     ],
                     "type": "Vec<u8>",
-                    "decoder": NeuronInfoLite.decode,
+                    "decoder": lambda raw_bytes, runtime: _cyscale_decode(
+                        "NeuronInfoLite", raw_bytes, runtime
+                    ),
                 },
                 "get_neurons_lite": {
                     "params": [
@@ -110,7 +194,9 @@ def preprocess_get_stake_info_for_coldkeys(addrs):
                         },
                     ],
                     "type": "Vec<u8>",
-                    "decoder": NeuronInfoLite.decode_vec,
+                    "decoder": lambda raw_bytes, runtime: _cyscale_decode(
+                        "Vec<NeuronInfoLite>", raw_bytes, runtime
+                    ),
                 },
                 "get_neuron": {
                     "params": [
@@ -124,7 +210,9 @@ def preprocess_get_stake_info_for_coldkeys(addrs):
                         },
                     ],
                     "type": "Vec<u8>",
-                    "decoder": NeuronInfo.decode,
+                    "decoder": lambda raw_bytes, runtime: _cyscale_decode(
+                        "NeuronInfo", raw_bytes, runtime
+                    ),
                 },
                 "get_neurons": {
                     "params": [
@@ -134,7 +222,9 @@ def preprocess_get_stake_info_for_coldkeys(addrs):
                         },
                     ],
                     "type": "Vec<u8>",
-                    "decoder": NeuronInfo.decode_vec,
+                    "decoder": lambda raw_bytes, runtime: _cyscale_decode(
+                        "Vec<NeuronInfo>", raw_bytes, runtime
+                    ),
                 },
             }
         },
@@ -148,19 +238,7 @@ def preprocess_get_stake_info_for_coldkeys(addrs):
                         },
                     ],
                     "type": "Vec<u8>",
-                    "encoder": lambda addr, reg: encode(
-                        "Vec<u8>",
-                        reg,
-                        list(
-                            bytes.fromhex(
-                                ss58_decode(
-                                    addr[0]
-                                    if isinstance(addr, list)
-                                    else addr["coldkey_account"]
-                                )
-                            )
-                        ),
-                    ),
+                    "encoder": _encode_vec_u8_coldkey,
                     "decoder": stake_info_decode_vec_legacy_compatibility,
                 },
                 "get_stake_info_for_coldkeys": {
@@ -171,12 +249,10 @@ def preprocess_get_stake_info_for_coldkeys(addrs):
                         },
                     ],
                     "type": "Vec<u8>",
-                    "encoder": lambda addrs, reg: encode(
-                        "Vec<Vec<u8>>",
-                        reg,
-                        preprocess_get_stake_info_for_coldkeys(addrs),
+                    "encoder": _encode_vec_vec_u8,
+                    "decoder": lambda raw_bytes, runtime: _cyscale_decode(
+                        "Vec<(Vec<u8>, Vec<StakeInfo>)>", raw_bytes, runtime
                     ),
-                    "decoder": StakeInfo.decode_vec_tuple_vec,
                 },
             },
         },
@@ -190,7 +266,9 @@ def preprocess_get_stake_info_for_coldkeys(addrs):
                         },
                     ],
                     "type": "Vec<u8>",
-                    "decoder": SubnetHyperparameters.decode_option,
+                    "decoder": lambda raw_bytes, runtime: _cyscale_decode(
+                        "Option<SubnetHyperparameters>", raw_bytes, runtime
+                    ),
                 },
                 "get_subnet_info": {
                     "params": [
@@ -200,7 +278,9 @@ def preprocess_get_stake_info_for_coldkeys(addrs):
                         },
                     ],
                     "type": "Vec<u8>",
-                    "decoder": SubnetInfo.decode_option,
+                    "decoder": lambda raw_bytes, runtime: _cyscale_decode(
+                        "Option<SubnetInfo>", raw_bytes, runtime
+                    ),
                 },
                 "get_subnet_info_v2": {
                     "params": [
@@ -210,17 +290,23 @@ def preprocess_get_stake_info_for_coldkeys(addrs):
                         },
                     ],
                     "type": "Vec<u8>",
-                    "decoder": SubnetInfoV2.decode_option,
+                    "decoder": lambda raw_bytes, runtime: _cyscale_decode(
+                        "Option<SubnetInfoV2>", raw_bytes, runtime
+                    ),
                 },
                 "get_subnets_info": {
                     "params": [],
                     "type": "Vec<u8>",
-                    "decoder": SubnetInfo.decode_vec_option,
+                    "decoder": lambda raw_bytes, runtime: _cyscale_decode(
+                        "Vec<Option<SubnetInfo>>", raw_bytes, runtime
+                    ),
                 },
                 "get_subnets_info_v2": {
                     "params": [],
                     "type": "Vec<u8>",
-                    "decoder": SubnetInfo.decode_vec_option,
+                    "decoder": lambda raw_bytes, runtime: _cyscale_decode(
+                        "Vec<Option<SubnetInfo>>", raw_bytes, runtime
+                    ),
                 },
             }
         },