WIP

Author: thewhaleking

async_substrate_interface/type_registry.py

@@ -13,17 +13,52 @@
 from scalecodec import ss58_decode
 
 
-def stake_info_decode_vec_legacy_compatibility(item) -> list[dict[str, Union[str, int, bytes, bool]]]:
+def stake_info_decode_vec_legacy_compatibility(
+    item,
+) -> list[dict[str, Union[str, int, bytes, bool]]]:
     stake_infos: list[StakeInfo] = StakeInfo.decode_vec(item)
-    NewStakeInfo = namedtuple("NewStakeInfo", ["netuid", "hotkey", "coldkey", "stake", "locked", "emission", "drain", "is_registered"])
+    NewStakeInfo = namedtuple(
+        "NewStakeInfo",
+        [
+            "netuid",
+            "hotkey",
+            "coldkey",
+            "stake",
+            "locked",
+            "emission",
+            "drain",
+            "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)
+            NewStakeInfo(
+                0,
+                stake_info.hotkey,
+                stake_info.coldkey,
+                stake_info.stake,
+                0,
+                0,
+                0,
+                False,
+            )
         )
     return output
 
 
+def preprocess_get_stake_info_for_coldkeys(addrs):
+    output = []
+    if isinstance(addrs[0], list):  # I think
+        for addr in addrs[0]:
+            output.append(list(bytes.fromhex(ss58_decode(addr))))
+    else:
+        if isinstance(addrs[0], dict):
+            for addr in addrs[0]["coldkey_accounts"]:
+                output.append(list(bytes.fromhex(ss58_decode(addr))))
+    return output
+
+
 _TYPE_REGISTRY: dict[str, dict] = {
     "types": {
         "Balance": "u64",  # Need to override default u128
@@ -114,7 +149,17 @@ def stake_info_decode_vec_legacy_compatibility(item) -> list[dict[str, Union[str
                     ],
                     "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"])))
+                        "Vec<u8>",
+                        reg,
+                        list(
+                            bytes.fromhex(
+                                ss58_decode(
+                                    addr[0]
+                                    if isinstance(addr, list)
+                                    else addr["coldkey_account"]
+                                )
+                            )
+                        ),
                     ),
                     "decoder": stake_info_decode_vec_legacy_compatibility,
                 },
@@ -129,7 +174,7 @@ def stake_info_decode_vec_legacy_compatibility(item) -> list[dict[str, Union[str
                     "encoder": lambda addrs, reg: encode(
                         "Vec<Vec<u8>>",
                         reg,
-                        [list(bytes.fromhex(ss58_decode(addr))) for addr in addrs],
+                        preprocess_get_stake_info_for_coldkeys(addrs),
                     ),
                     "decoder": StakeInfo.decode_vec_tuple_vec,
                 },

tests/integration_tests/test_substrate_interface.py

@@ -116,8 +116,13 @@ def test_query_map_with_odd_number_of_params():
 
 def test_old_runtime_calls():
     from bittensor import SubtensorApi
-    sub = SubtensorApi(network='archive', legacy_methods=True, async_subtensor=False)
+
+    sub = SubtensorApi(network="archive", legacy_methods=True, async_subtensor=False)
     # will pass
-    assert sub.get_stake_info_for_coldkey('5CQ6dMW8JZhKCZX9kWsZRqa3kZRKmNHxbPPVFEt6FgyvGv2G', 4943592)
+    assert sub.get_stake_info_for_coldkey(
+        "5CQ6dMW8JZhKCZX9kWsZRqa3kZRKmNHxbPPVFEt6FgyvGv2G", 4943592
+    )
     # needs to use legacy
-    assert sub.get_stake_info_for_coldkey('5CQ6dMW8JZhKCZX9kWsZRqa3kZRKmNHxbPPVFEt6FgyvGv2G', 4670227)
+    assert sub.get_stake_info_for_coldkey(
+        "5CQ6dMW8JZhKCZX9kWsZRqa3kZRKmNHxbPPVFEt6FgyvGv2G", 4670227
+    )