diff --git a/bittensor/core/async_subtensor.py b/bittensor/core/async_subtensor.py index c648f8ca8e..33bd99f76d 100644 --- a/bittensor/core/async_subtensor.py +++ b/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, diff --git a/bittensor/core/subtensor.py b/bittensor/core/subtensor.py index e6d862a081..fd04dfbc93 100644 --- a/bittensor/core/subtensor.py +++ b/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 diff --git a/bittensor/core/subtensor_api/staking.py b/bittensor/core/subtensor_api/staking.py index a99e201610..979d8a2632 100644 --- a/bittensor/core/subtensor_api/staking.py +++ b/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 diff --git a/bittensor/core/subtensor_api/utils.py b/bittensor/core/subtensor_api/utils.py index 3f31d51b90..adef4f31b1 100644 --- a/bittensor/core/subtensor_api/utils.py +++ b/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 diff --git a/tests/unit_tests/test_async_subtensor.py b/tests/unit_tests/test_async_subtensor.py index 24dee30ef7..93573438d2 100644 --- a/tests/unit_tests/test_async_subtensor.py +++ b/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 diff --git a/tests/unit_tests/test_subtensor.py b/tests/unit_tests/test_subtensor.py index 6ee02cea5a..217caeb9a9 100644 --- a/tests/unit_tests/test_subtensor.py +++ b/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