Add get_stake_weight methods

bittensor/core/async_subtensor.py

@@ -2924,6 +2924,7 @@ async def get_stake_operations_fee(
                 or reuse_block.
             reuse_block: Whether to reuse for this query the last-used block. Do not specify if also specifying block
                 or block_hash.
+
         Returns:
             The calculated stake fee as a Balance object.
         """
@@ -2938,6 +2939,38 @@ async def get_stake_operations_fee(
         )
         return amount * (result.value / U16_MAX)
 
+    async def get_stake_weight(
+        self,
+        netuid: int,
+        block: Optional[int] = None,
+        block_hash: Optional[str] = None,
+        reuse_block: bool = False,
+    ) -> list[float]:
+        """
+        Retrieves the stake weight for all hotkeys in a given subnet.
+
+        Arguments:
+            netuid: Netuid of subnet.
+            block: Block number at which to perform the calculation.
+            block_hash: The hash of the blockchain block number for the query. Do not specify if also specifying block
+                or reuse_block.
+            reuse_block: Whether to reuse for this query the last-used block. Do not specify if also specifying block
+                or block_hash.
+
+        Returns:
+            A list of stake weights for all hotkeys in the specified subnet.
+        """
+        block_hash = await self.determine_block_hash(
+            block=block, block_hash=block_hash, reuse_block=reuse_block
+        )
+        result = await self.substrate.query(
+            module="SubtensorModule",
+            storage_function="StakeWeight",
+            params=[netuid],
+            block_hash=block_hash,
+        )
+        return [u16_normalized_float(w) for w in result]
+
     async def get_subnet_burn_cost(
         self,
         block: Optional[int] = None,

bittensor/core/subtensor.py

@@ -2088,6 +2088,26 @@ def get_stake_operations_fee(
         )
         return amount * (result.value / U16_MAX)
 
+    def get_stake_weight(self, netuid: int, block: Optional[int] = None) -> list[float]:
+        """
+        Retrieves the stake weight for all hotkeys in a given subnet.
+
+        Arguments:
+            netuid: Netuid of subnet.
+            block: Block number at which to perform the calculation.
+
+        Returns:
+            A list of stake weights for all hotkeys in the specified subnet.
+        """
+        block_hash = self.determine_block_hash(block=block)
+        result = self.substrate.query(
+            module="SubtensorModule",
+            storage_function="StakeWeight",
+            params=[netuid],
+            block_hash=block_hash,
+        )
+        return [u16_normalized_float(w) for w in result]
+
     def get_subnet_burn_cost(self, block: Optional[int] = None) -> Optional[Balance]:
         """
         Retrieves the burn cost for registering a new subnet within the Bittensor network. This cost represents the

bittensor/core/subtensor_api/staking.py

@@ -20,6 +20,7 @@ def __init__(self, subtensor: Union["_Subtensor", "_AsyncSubtensor"]):
         self.get_stake_info_for_coldkey = subtensor.get_stake_info_for_coldkey
         self.get_stake_movement_fee = subtensor.get_stake_movement_fee
         self.get_stake_operations_fee = subtensor.get_stake_operations_fee
+        self.get_stake_weight = subtensor.get_stake_weight
         self.get_unstake_fee = subtensor.get_unstake_fee
         self.unstake = subtensor.unstake
         self.unstake_all = subtensor.unstake_all

bittensor/core/subtensor_api/utils.py

@@ -90,6 +90,7 @@ def add_legacy_methods(subtensor: "SubtensorApi"):
     )
     subtensor.get_stake_movement_fee = subtensor._subtensor.get_stake_movement_fee
     subtensor.get_stake_operations_fee = subtensor._subtensor.get_stake_operations_fee
+    subtensor.get_stake_weight = subtensor._subtensor.get_stake_weight
     subtensor.get_subnet_burn_cost = subtensor._subtensor.get_subnet_burn_cost
     subtensor.get_subnet_hyperparameters = (
         subtensor._subtensor.get_subnet_hyperparameters

tests/unit_tests/test_async_subtensor.py

@@ -4060,3 +4060,39 @@ async def test_get_stake_movement_fee(subtensor, mocker):
         amount=amount, netuid=netuid, block=None
     )
     assert result == mocked_get_stake_operations_fee.return_value
+
+
+@pytest.mark.asyncio
+async def test_get_stake_weight(subtensor, mocker):
+    """Verify that `get_stake_weight` method calls proper methods and returns the correct value."""
+    # Preps
+    netuid = mocker.Mock()
+    fake_weights = [0, 100, 15000]
+    expected_result = [0.0, 0.0015259021896696422, 0.22888532845044632]
+
+    mock_determine_block_hash = mocker.patch.object(
+        subtensor,
+        "determine_block_hash",
+    )
+    mocked_query = mocker.patch.object(
+        subtensor.substrate,
+        "query",
+        return_value=fake_weights,
+    )
+
+    # Call
+    result = await subtensor.get_stake_weight(netuid=netuid)
+
+    # Asserts
+    mock_determine_block_hash.assert_awaited_once_with(
+        block=None,
+        block_hash=None,
+        reuse_block=False,
+    )
+    mocked_query.assert_awaited_once_with(
+        module="SubtensorModule",
+        storage_function="StakeWeight",
+        params=[netuid],
+        block_hash=mock_determine_block_hash.return_value,
+    )
+    assert result == expected_result

tests/unit_tests/test_subtensor.py

@@ -4275,3 +4275,34 @@ def test_get_stake_movement_fee(subtensor, mocker):
         amount=amount, netuid=netuid, block=None
     )
     assert result == mocked_get_stake_operations_fee.return_value
+
+
+def test_get_stake_weight(subtensor, mocker):
+    """Verify that `get_stake_weight` method calls proper methods and returns the correct value."""
+    # Preps
+    netuid = mocker.Mock()
+    fake_weights = [0, 100, 15000]
+    expected_result = [0.0, 0.0015259021896696422, 0.22888532845044632]
+
+    mock_determine_block_hash = mocker.patch.object(
+        subtensor,
+        "determine_block_hash",
+    )
+    mocked_query = mocker.patch.object(
+        subtensor.substrate,
+        "query",
+        return_value=fake_weights,
+    )
+
+    # Call
+    result = subtensor.get_stake_weight(netuid=netuid)
+
+    # Asserts
+    mock_determine_block_hash.assert_called_once_with(block=None)
+    mocked_query.assert_called_once_with(
+        module="SubtensorModule",
+        storage_function="StakeWeight",
+        params=[netuid],
+        block_hash=mock_determine_block_hash.return_value,
+    )
+    assert result == expected_result