[Staging] add stake info runtime api (#1505)

* add stake info runtime api specs

* add StakeInfo to chaindata

* add util function for ss58 -> Vec<u8>

* use new util

* add stake info query

* run black

* modify helpers to decode w/only typestring as well

* typo

* use new tuple format

* run black

* make 7 instead of 6

* run black

bittensor/__init__.py

@@ -154,6 +154,28 @@ def turn_console_off():
                 },
             }
         },
+        "StakeInfoRuntimeApi": {
+            "methods": {
+                "get_stake_info_for_coldkey": {
+                    "params": [
+                        {
+                            "name": "coldkey_account_vec",
+                            "type": "Vec<u8>",
+                        },
+                    ],
+                    "type": "Vec<u8>",
+                },
+                "get_stake_info_for_coldkeys": {
+                    "params": [
+                        {
+                            "name": "coldkey_account_vecs",
+                            "type": "Vec<Vec<u8>>",
+                        },
+                    ],
+                    "type": "Vec<u8>",
+                },
+            },
+        },
     },
 }
 

bittensor/_subtensor/chain_data.py

@@ -139,6 +139,14 @@
                 ["ip_type", "u8"],
             ],
         },
+        "StakeInfo": {
+            "type": "struct",
+            "type_mapping": [
+                ["hotkey", "AccountId"],
+                ["coldkey", "AccountId"],
+                ["stake", "Compact<u64>"],
+            ],
+        },
     }
 }
 
@@ -149,6 +157,7 @@ class ChainDataType(Enum):
     DelegateInfo = 3
     NeuronInfoLite = 4
     DelegatedInfo = 5
+    StakeInfo = 7
 
 
 # Constants
@@ -162,6 +171,21 @@ def from_scale_encoding(
     type_name: ChainDataType,
     is_vec: bool = False,
     is_option: bool = False,
+) -> Optional[Dict]:
+    type_string = type_name.name
+    if type_name == ChainDataType.DelegatedInfo:
+        # DelegatedInfo is a tuple of (DelegateInfo, Compact<u64>)
+        type_string = f"({ChainDataType.DelegateInfo.name}, Compact<u64>)"
+    if is_option:
+        type_string = f"Option<{type_string}>"
+    if is_vec:
+        type_string = f"Vec<{type_string}>"
+
+    return from_scale_encoding_using_type_string(input, type_string)
+
+
+def from_scale_encoding_using_type_string(
+    input: Union[List[int], bytes, ScaleBytes], type_string: str
 ) -> Optional[Dict]:
     if isinstance(input, ScaleBytes):
         as_scale_bytes = input
@@ -180,15 +204,6 @@ def from_scale_encoding(
     rpc_runtime_config.update_type_registry(load_type_registry_preset("legacy"))
     rpc_runtime_config.update_type_registry(custom_rpc_type_registry)
 
-    type_string = type_name.name
-    if type_name == ChainDataType.DelegatedInfo:
-        # DelegatedInfo is a tuple of (DelegateInfo, Compact<u64>)
-        type_string = f"({ChainDataType.DelegateInfo.name}, Compact<u64>)"
-    if is_option:
-        type_string = f"Option<{type_string}>"
-    if is_vec:
-        type_string = f"Vec<{type_string}>"
-
     obj = rpc_runtime_config.create_scale_object(type_string, data=as_scale_bytes)
 
     return obj.decode()
@@ -646,6 +661,76 @@ def delegated_list_from_vec_u8(
         return decoded
 
 
+@dataclass
+class StakeInfo:
+    r"""
+    Dataclass for stake info.
+    """
+    hotkey_ss58: str  # Hotkey address
+    coldkey_ss58: str  # Coldkey address
+    stake: Balance  # Stake for the hotkey-coldkey pair
+
+    @classmethod
+    def fix_decoded_values(cls, decoded: Any) -> "StakeInfo":
+        r"""Fixes the decoded values."""
+
+        return cls(
+            hotkey_ss58=ss58_encode(decoded["hotkey"], bittensor.__ss58_format__),
+            coldkey_ss58=ss58_encode(decoded["coldkey"], bittensor.__ss58_format__),
+            stake=Balance.from_rao(decoded["stake"]),
+        )
+
+    @classmethod
+    def from_vec_u8(cls, vec_u8: List[int]) -> Optional["StakeInfo"]:
+        r"""Returns a StakeInfo object from a vec_u8."""
+        if len(vec_u8) == 0:
+            return None
+
+        decoded = from_scale_encoding(vec_u8, ChainDataType.StakeInfo)
+
+        if decoded is None:
+            return None
+
+        decoded = StakeInfo.fix_decoded_values(decoded)
+
+        return decoded
+
+    @classmethod
+    def list_of_tuple_from_vec_u8(
+        cls, vec_u8: List[int]
+    ) -> Dict[str, List["StakeInfo"]]:
+        r"""Returns a list of StakeInfo objects from a vec_u8."""
+        decoded: Optional[
+            List[Tuple(str, List[object])]
+        ] = from_scale_encoding_using_type_string(
+            input=vec_u8, type_string="Vec<(AccountId, Vec<StakeInfo>)>"
+        )
+
+        if decoded is None:
+            return {}
+
+        stake_map = {
+            ss58_encode(address=account_id, ss58_format=bittensor.__ss58_format__): [
+                StakeInfo.fix_decoded_values(d) for d in stake_info
+            ]
+            for account_id, stake_info in decoded
+        }
+
+        return stake_map
+
+    @classmethod
+    def list_from_vec_u8(cls, vec_u8: List[int]) -> List["StakeInfo"]:
+        r"""Returns a list of StakeInfo objects from a vec_u8."""
+        decoded = from_scale_encoding(vec_u8, ChainDataType.StakeInfo, is_vec=True)
+
+        if decoded is None:
+            return []
+
+        decoded = [StakeInfo.fix_decoded_values(d) for d in decoded]
+
+        return decoded
+
+
 @dataclass
 class SubnetInfo:
     r"""

bittensor/_subtensor/subtensor_impl.py

@@ -28,7 +28,13 @@
 
 from bittensor._subtensor.chain_data import custom_rpc_type_registry
 from bittensor.utils.balance import Balance
-from bittensor.utils import U16_NORMALIZED_FLOAT, U64_MAX, RAOPERTAO, U16_MAX
+from bittensor.utils import (
+    U16_NORMALIZED_FLOAT,
+    U64_MAX,
+    RAOPERTAO,
+    U16_MAX,
+    ss58_to_vec_u8,
+)
 from bittensor.utils.registration import POWSolution
 
 # Local imports.
@@ -37,6 +43,7 @@
     DelegateInfo,
     PrometheusInfo,
     SubnetInfo,
+    StakeInfo,
     NeuronInfoLite,
     axon_info,
     ProposalVoteData,
@@ -1569,8 +1576,7 @@ def make_substrate_call_with_retry(encoded_hotkey: List[int]):
                     params=params,
                 )
 
-        hotkey_bytes: bytes = bittensor.utils.ss58_address_to_bytes(hotkey_ss58)
-        encoded_hotkey: List[int] = [int(byte) for byte in hotkey_bytes]
+        encoded_hotkey = ss58_to_vec_u8(hotkey_ss58)
         json_body = make_substrate_call_with_retry(encoded_hotkey)
         result = json_body["result"]
 
@@ -1617,8 +1623,7 @@ def make_substrate_call_with_retry(encoded_coldkey: List[int]):
                     params=params,
                 )
 
-        coldkey_bytes: bytes = bittensor.utils.ss58_address_to_bytes(coldkey_ss58)
-        encoded_coldkey: List[int] = [int(byte) for byte in coldkey_bytes]
+        encoded_coldkey = ss58_to_vec_u8(coldkey_ss58)
         json_body = make_substrate_call_with_retry(encoded_coldkey)
         result = json_body["result"]
 
@@ -1627,6 +1632,59 @@ def make_substrate_call_with_retry(encoded_coldkey: List[int]):
 
         return DelegateInfo.delegated_list_from_vec_u8(result)
 
+    ###########################
+    #### Stake Information ####
+    ###########################
+
+    def get_stake_info_for_colkey(
+        self, coldkey_ss58: str, block: Optional[int] = None
+    ) -> List[Tuple[DelegateInfo, Balance]]:
+        """Returns the list of StakeInfo objects for this coldkey"""
+
+        encoded_coldkey = ss58_to_vec_u8(coldkey_ss58)
+
+        hex_bytes_result = self.query_runtime_api(
+            runtime_api="StakeInfoRuntimeApi",
+            method="get_stake_info_for_coldkey",
+            params=[encoded_coldkey],
+            block=block,
+        )
+
+        if hex_bytes_result == None:
+            return None
+
+        if hex_bytes_result.startswith("0x"):
+            bytes_result = bytes.fromhex(hex_bytes_result[2:])
+        else:
+            bytes_result = bytes.fromhex(hex_bytes_result)
+
+        return StakeInfo.list_from_vec_u8(bytes_result)
+
+    def get_stake_info_for_colkeys(
+        self, coldkey_ss58_list: List[str], block: Optional[int] = None
+    ) -> List[Tuple[DelegateInfo, Balance]]:
+        """Returns the list of StakeInfo objects for all coldkeys in the list."""
+        encoded_coldkeys = [
+            ss58_to_vec_u8(coldkey_ss58) for coldkey_ss58 in coldkey_ss58_list
+        ]
+
+        hex_bytes_result = self.query_runtime_api(
+            runtime_api="StakeInfoRuntimeApi",
+            method="get_stake_info_for_coldkeys",
+            params=[encoded_coldkeys],
+            block=block,
+        )
+
+        if hex_bytes_result == None:
+            return None
+
+        if hex_bytes_result.startswith("0x"):
+            bytes_result = bytes.fromhex(hex_bytes_result[2:])
+        else:
+            bytes_result = bytes.fromhex(hex_bytes_result)
+
+        return StakeInfo.list_of_tuple_from_vec_u8(bytes_result)
+
     ########################################
     #### Neuron information per subnet ####
     ########################################

bittensor/utils/__init__.py

@@ -17,6 +17,12 @@
 U64_MAX = 18446744073709551615
 
 
+def ss58_to_vec_u8(ss58_address: str) -> List[int]:
+    ss58_bytes: bytes = bittensor.utils.ss58_address_to_bytes(ss58_address)
+    encoded_address: List[int] = [int(byte) for byte in ss58_bytes]
+    return encoded_address
+
+
 def indexed_values_to_dataframe(
     prefix: Union[str, int],
     index: Union[list, torch.LongTensor],