add get_stake_operations_fee amd update stake methods

Author: Roman

bittensor/core/async_subtensor.py

@@ -117,7 +117,11 @@
     price_to_tick,
     LiquidityPosition,
 )
-from bittensor.utils.weight_utils import generate_weight_hash, convert_uids_and_weights
+from bittensor.utils.weight_utils import (
+    generate_weight_hash,
+    convert_uids_and_weights,
+    U16_MAX,
+)
 
 if TYPE_CHECKING:
     from async_substrate_interface.types import ScaleObj
@@ -2522,6 +2526,7 @@ async def get_stake(
 
         return Balance.from_rao(int(stake)).set_unit(netuid=netuid)
 
+    # TODO: remove unused parameters in SDK.v10
     async def get_stake_add_fee(
         self,
         amount: Balance,
@@ -2543,19 +2548,9 @@ async def get_stake_add_fee(
         Returns:
             The calculated stake fee as a Balance object
         """
-        result = await self.query_runtime_api(
-            runtime_api="StakeInfoRuntimeApi",
-            method="get_stake_fee",
-            params=[
-                None,
-                coldkey_ss58,
-                (hotkey_ss58, netuid),
-                coldkey_ss58,
-                amount.rao,
-            ],
-            block=block,
+        return await self.get_stake_operations_fee(
+            amount=amount, netuid=netuid, block=block
         )
-        return Balance.from_rao(result)
 
     async def get_subnet_info(
         self,
@@ -2669,6 +2664,7 @@ async def get_subnet_prices(
         prices.update({0: Balance.from_tao(1)})
         return prices
 
+    # TODO: remove unused parameters in SDK.v10
     async def get_unstake_fee(
         self,
         amount: Balance,
@@ -2690,20 +2686,11 @@ async def get_unstake_fee(
         Returns:
             The calculated stake fee as a Balance object
         """
-        result = await self.query_runtime_api(
-            runtime_api="StakeInfoRuntimeApi",
-            method="get_stake_fee",
-            params=[
-                None,
-                coldkey_ss58,
-                (hotkey_ss58, netuid),
-                coldkey_ss58,
-                amount.rao,
-            ],
-            block=block,
+        return await self.get_stake_operations_fee(
+            amount=amount, netuid=netuid, block=block
         )
-        return Balance.from_rao(result)
 
+    # TODO: remove unused parameters in SDK.v10
     async def get_stake_movement_fee(
         self,
         amount: Balance,
@@ -2731,19 +2718,9 @@ async def get_stake_movement_fee(
         Returns:
             The calculated stake fee as a Balance object
         """
-        result = await self.query_runtime_api(
-            runtime_api="StakeInfoRuntimeApi",
-            method="get_stake_fee",
-            params=[
-                (origin_hotkey_ss58, origin_netuid),
-                origin_coldkey_ss58,
-                (destination_hotkey_ss58, destination_netuid),
-                destination_coldkey_ss58,
-                amount.rao,
-            ],
-            block=block,
+        return await self.get_stake_operations_fee(
+            amount=amount, netuid=origin_netuid, block=block
         )
-        return Balance.from_rao(result)
 
     async def get_stake_for_coldkey_and_hotkey(
         self,
@@ -2864,6 +2841,39 @@ async def get_stake_for_hotkey(
 
     get_hotkey_stake = get_stake_for_hotkey
 
+    async def get_stake_operations_fee(
+        self,
+        amount: Balance,
+        netuid: int,
+        block: Optional[int] = None,
+        block_hash: Optional[str] = None,
+        reuse_block: bool = False,
+    ):
+        """Returns fee for any stake operation in specified subnet.
+
+        Args:
+            amount: Amount of stake to add in Alpha/TAO.
+            netuid: Netuid of subnet.
+            block: The block number at which to query the stake information. Do not specify if also specifying
+                block_hash or reuse_block.
+            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:
+            The calculated stake fee as a Balance object.
+        """
+        block_hash = await self.determine_block_hash(
+            block=block, block_hash=block_hash, reuse_block=reuse_block
+        )
+        result = await self.substrate.query(
+            module="Swap",
+            storage_function="FeeRate",
+            params=[netuid],
+            block_hash=block_hash,
+        )
+        return amount * (result.value / U16_MAX)
+
     async def get_subnet_burn_cost(
         self,
         block: Optional[int] = None,

bittensor/core/subtensor.py

@@ -120,7 +120,11 @@
     price_to_tick,
     LiquidityPosition,
 )
-from bittensor.utils.weight_utils import generate_weight_hash, convert_uids_and_weights
+from bittensor.utils.weight_utils import (
+    generate_weight_hash,
+    convert_uids_and_weights,
+    U16_MAX,
+)
 
 if TYPE_CHECKING:
     from bittensor_wallet import Wallet
@@ -1756,6 +1760,7 @@ def get_stake(
 
         return Balance.from_rao(int(stake)).set_unit(netuid=netuid)
 
+    # TODO: remove unused parameters in SDK.v10
     def get_stake_add_fee(
         self,
         amount: Balance,
@@ -1777,19 +1782,7 @@ def get_stake_add_fee(
         Returns:
             The calculated stake fee as a Balance object
         """
-        result = self.query_runtime_api(
-            runtime_api="StakeInfoRuntimeApi",
-            method="get_stake_fee",
-            params=[
-                None,
-                coldkey_ss58,
-                (hotkey_ss58, netuid),
-                coldkey_ss58,
-                amount.rao,
-            ],
-            block=block,
-        )
-        return Balance.from_rao(result)
+        return self.get_stake_operations_fee(amount=amount, netuid=netuid, block=block)
 
     def get_subnet_info(
         self, netuid: int, block: Optional[int] = None
@@ -1882,6 +1875,7 @@ def get_subnet_prices(
         prices.update({0: Balance.from_tao(1)})
         return prices
 
+    # TODO: remove unused parameters in SDK.v10
     def get_unstake_fee(
         self,
         amount: Balance,
@@ -1903,20 +1897,9 @@ def get_unstake_fee(
         Returns:
             The calculated stake fee as a Balance object
         """
-        result = self.query_runtime_api(
-            runtime_api="StakeInfoRuntimeApi",
-            method="get_stake_fee",
-            params=[
-                (hotkey_ss58, netuid),
-                coldkey_ss58,
-                None,
-                coldkey_ss58,
-                amount.rao,
-            ],
-            block=block,
-        )
-        return Balance.from_rao(result)
+        return self.get_stake_operations_fee(amount=amount, netuid=netuid, block=block)
 
+    # TODO: remove unused parameters in SDK.v10
     def get_stake_movement_fee(
         self,
         amount: Balance,
@@ -1944,19 +1927,9 @@ def get_stake_movement_fee(
         Returns:
             The calculated stake fee as a Balance object
         """
-        result = self.query_runtime_api(
-            runtime_api="StakeInfoRuntimeApi",
-            method="get_stake_fee",
-            params=[
-                (origin_hotkey_ss58, origin_netuid),
-                origin_coldkey_ss58,
-                (destination_hotkey_ss58, destination_netuid),
-                destination_coldkey_ss58,
-                amount.rao,
-            ],
-            block=block,
+        return self.get_stake_operations_fee(
+            amount=amount, netuid=origin_netuid, block=block
         )
-        return Balance.from_rao(result)
 
     def get_stake_for_coldkey_and_hotkey(
         self,
@@ -2044,6 +2017,31 @@ def get_stake_for_hotkey(
 
     get_hotkey_stake = get_stake_for_hotkey
 
+    def get_stake_operations_fee(
+        self,
+        amount: Balance,
+        netuid: int,
+        block: Optional[int] = None,
+    ):
+        """Returns fee for any stake operation in specified subnet.
+
+        Args:
+            amount: Amount of stake to add in Alpha/TAO.
+            netuid: Netuid of subnet.
+            block: Block number at which to perform the calculation.
+
+        Returns:
+            The calculated stake fee as a Balance object.
+        """
+        block_hash = self.determine_block_hash(block=block)
+        result = self.substrate.query(
+            module="Swap",
+            storage_function="FeeRate",
+            params=[netuid],
+            block_hash=block_hash,
+        )
+        return amount * (result.value / U16_MAX)
+
     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