Adds explicit type for balance inputs

Author: ibraheem-abe

bittensor/core/async_subtensor.py

@@ -2795,7 +2795,7 @@ async def add_stake(
         wallet: "Wallet",
         hotkey_ss58: Optional[str] = None,
         netuid: Optional[int] = None,
-        amount: Optional[Union["Balance", float]] = None,
+        amount: Optional[Balance] = None,
         wait_for_inclusion: bool = True,
         wait_for_finalization: bool = False,
     ) -> bool:
@@ -2807,7 +2807,7 @@ async def add_stake(
         Args:
             wallet (bittensor_wallet.Wallet): The wallet to be used for staking.
             hotkey_ss58 (Optional[str]): The ``SS58`` address of the hotkey associated with the neuron.
-            amount (Union[Balance, float]): The amount of TAO to stake.
+            amount (Balance): The amount of TAO to stake.
             wait_for_inclusion (bool): Waits for the transaction to be included in a block.
             wait_for_finalization (bool): Waits for the transaction to be finalized on the blockchain.
 
@@ -2832,7 +2832,7 @@ async def add_stake_multiple(
         wallet: "Wallet",
         hotkey_ss58s: list[str],
         netuids: list[int],
-        amounts: Optional[list[Union["Balance", float]]] = None,
+        amounts: Optional[list[Balance]] = None,
         wait_for_inclusion: bool = True,
         wait_for_finalization: bool = False,
     ) -> bool:
@@ -2843,7 +2843,7 @@ async def add_stake_multiple(
         Args:
             wallet (bittensor_wallet.Wallet): The wallet used for staking.
             hotkey_ss58s (list[str]): List of ``SS58`` addresses of hotkeys to stake to.
-            amounts (list[Union[Balance, float]]): Corresponding amounts of TAO to stake for each hotkey.
+            amounts (list[Balance]): Corresponding amounts of TAO to stake for each hotkey.
             wait_for_inclusion (bool): Waits for the transaction to be included in a block.
             wait_for_finalization (bool): Waits for the transaction to be finalized on the blockchain.
 
@@ -2975,7 +2975,7 @@ async def move_stake(
         origin_netuid: int,
         destination_hotkey: str,
         destination_netuid: int,
-        amount: Union["Balance", float],
+        amount: Balance,
         wait_for_inclusion: bool = True,
         wait_for_finalization: bool = False,
     ) -> bool:
@@ -2988,15 +2988,13 @@ async def move_stake(
             origin_netuid (int): The netuid of the source subnet.
             destination_hotkey (str): The SS58 address of the destination hotkey.
             destination_netuid (int): The netuid of the destination subnet.
-            amount (Union[Balance, float]): Amount of stake to move.
+            amount (Balance): Amount of stake to move.
             wait_for_inclusion (bool): Waits for the transaction to be included in a block.
             wait_for_finalization (bool): Waits for the transaction to be finalized on the blockchain.
 
         Returns:
             success (bool): True if the stake movement was successful.
         """
-        if isinstance(amount, float):
-            amount = Balance.from_tao(amount)
 
         return await move_stake_extrinsic(
             subtensor=self,
@@ -3420,7 +3418,7 @@ async def transfer_stake(
         hotkey_ss58: str,
         origin_netuid: int,
         destination_netuid: int,
-        amount: Union["Balance", float],
+        amount: Balance,
         wait_for_inclusion: bool = True,
         wait_for_finalization: bool = False,
     ) -> bool:
@@ -3433,15 +3431,13 @@ async def transfer_stake(
             hotkey_ss58 (str): The hotkey SS58 address associated with the stake.
             origin_netuid (int): The source subnet UID.
             destination_netuid (int): The destination subnet UID.
-            amount (Union[Balance, float, int]): Amount to transfer.
+            amount (Balance): Amount to transfer.
             wait_for_inclusion (bool): If true, waits for inclusion before returning.
             wait_for_finalization (bool): If true, waits for finalization before returning.
 
         Returns:
             success (bool): True if the transfer was successful.
         """
-        if isinstance(amount, float):
-            amount = Balance.from_tao(amount)
 
         return await transfer_stake_extrinsic(
             subtensor=self,
@@ -3481,8 +3477,6 @@ async def transfer(
         Returns:
             `True` if the transferring was successful, otherwise `False`.
         """
-        if isinstance(amount, float):
-            amount = Balance.from_tao(amount)
 
         return await transfer_extrinsic(
             subtensor=self,
@@ -3500,7 +3494,7 @@ async def unstake(
         wallet: "Wallet",
         hotkey_ss58: Optional[str] = None,
         netuid: Optional[int] = None,
-        amount: Optional[Union["Balance", float]] = None,
+        amount: Optional[Balance] = None,
         wait_for_inclusion: bool = True,
         wait_for_finalization: bool = False,
     ) -> bool:
@@ -3512,7 +3506,7 @@ async def unstake(
             wallet (bittensor_wallet.wallet): The wallet associated with the neuron from which the stake is being
                 removed.
             hotkey_ss58 (Optional[str]): The ``SS58`` address of the hotkey account to unstake from.
-            amount (Union[Balance, float]): The amount of TAO to unstake. If not specified, unstakes all.
+            amount (Balance): The amount of TAO to unstake. If not specified, unstakes all.
             wait_for_inclusion (bool): Waits for the transaction to be included in a block.
             wait_for_finalization (bool): Waits for the transaction to be finalized on the blockchain.
 
@@ -3537,7 +3531,7 @@ async def unstake_multiple(
         wallet: "Wallet",
         hotkey_ss58s: list[str],
         netuids: list[int],
-        amounts: Optional[list[Union["Balance", float]]] = None,
+        amounts: Optional[list[Balance]] = None,
         wait_for_inclusion: bool = True,
         wait_for_finalization: bool = False,
     ) -> bool:

bittensor/core/extrinsics/asyncex/move_stake.py

@@ -1,4 +1,4 @@
-from typing import Union, Optional, TYPE_CHECKING
+from typing import TYPE_CHECKING
 
 from bittensor.utils.balance import Balance
 from bittensor.utils.btlogging import logging
@@ -7,21 +7,19 @@
     from bittensor_wallet import Wallet
     from bittensor.core.subtensor import Subtensor
 
+
 async def transfer_stake_extrinsic(
     subtensor: "Subtensor",
     wallet: "Wallet",
     destination_coldkey_ss58: str,
     hotkey_ss58: str,
     origin_netuid: int,
     destination_netuid: int,
-    amount: Optional[Union["Balance", float]] = None,
+    amount: Balance,
     wait_for_inclusion: bool = True,
     wait_for_finalization: bool = False,
 ) -> bool:
-    if not isinstance(amount, Balance):
-        amount = Balance.from_tao(amount)
     amount.set_unit(netuid=origin_netuid)
-
     # Verify ownership
     hotkey_owner = await subtensor.get_hotkey_owner(hotkey_ss58)
     if hotkey_owner != wallet.coldkeypub.ss58_address:
@@ -107,20 +105,18 @@ async def transfer_stake_extrinsic(
         logging.error(f":cross_mark: [red]Failed[/red]: {str(e)}")
         return False
 
+
 async def swap_stake_extrinsic(
     subtensor: "Subtensor",
     wallet: "Wallet",
     hotkey_ss58: str,
     origin_netuid: int,
     destination_netuid: int,
-    amount: Optional[Union["Balance", float]] = None,
+    amount: Balance,
     wait_for_inclusion: bool = True,
     wait_for_finalization: bool = False,
 ) -> bool:
-    if not isinstance(amount, Balance):
-        amount = Balance.from_tao(amount)
     amount.set_unit(netuid=origin_netuid)
-
     # Verify ownership
     hotkey_owner = await subtensor.get_hotkey_owner(hotkey_ss58)
     if hotkey_owner != wallet.coldkeypub.ss58_address:
@@ -205,21 +201,19 @@ async def swap_stake_extrinsic(
         logging.error(f":cross_mark: [red]Failed[/red]: {str(e)}")
         return False
 
+
 async def move_stake_extrinsic(
     subtensor: "Subtensor",
     wallet: "Wallet",
     origin_hotkey: str,
     origin_netuid: int,
     destination_hotkey: str,
     destination_netuid: int,
-    amount: Optional[Union["Balance", float]] = None,
+    amount: Balance,
     wait_for_inclusion: bool = True,
     wait_for_finalization: bool = False,
 ) -> bool:
-    if not isinstance(amount, Balance):
-        amount = Balance.from_tao(amount)
     amount.set_unit(netuid=origin_netuid)
-
     # Verify ownership of origin hotkey
     origin_owner = await subtensor.get_hotkey_owner(origin_hotkey)
     if origin_owner != wallet.coldkeypub.ss58_address:

bittensor/core/extrinsics/asyncex/staking.py

@@ -1,5 +1,5 @@
 import asyncio
-from typing import Optional, Sequence, TYPE_CHECKING, Union
+from typing import Optional, Sequence, TYPE_CHECKING
 
 from bittensor.core.errors import StakeError, NotRegisteredError
 from bittensor.utils import unlock_key
@@ -14,10 +14,10 @@
 async def add_stake_extrinsic(
     subtensor: "AsyncSubtensor",
     wallet: "Wallet",
-    old_balance: Optional["Balance"] = None,
+    old_balance: Optional[Balance] = None,
     hotkey_ss58: Optional[str] = None,
     netuid: Optional[int] = None,
-    amount: Optional[Union["Balance", float]] = None,
+    amount: Optional[Balance] = None,
     wait_for_inclusion: bool = True,
     wait_for_finalization: bool = False,
 ) -> bool:
@@ -72,10 +72,9 @@ async def add_stake_extrinsic(
     if amount is None:
         # Stake it all.
         staking_balance = Balance.from_tao(old_balance.tao)
-    elif not isinstance(amount, Balance):
-        staking_balance = Balance.from_tao(amount)
     else:
         staking_balance = amount
+    staking_balance.set_unit(netuid)
 
     # Leave existential balance to keep key alive.
     if staking_balance > old_balance - existential_deposit:
@@ -162,8 +161,8 @@ async def add_stake_multiple_extrinsic(
     wallet: "Wallet",
     hotkey_ss58s: list[str],
     netuids: list[int],
-    old_balance: Optional["Balance"] = None,
-    amounts: Optional[list["Balance"]] = None,
+    old_balance: Optional[Balance] = None,
+    amounts: Optional[list[Balance]] = None,
     wait_for_inclusion: bool = True,
     wait_for_finalization: bool = False,
 ) -> bool:
@@ -185,6 +184,7 @@ async def add_stake_multiple_extrinsic(
         success: `True` if extrinsic was finalized or included in the block. `True` if any wallet was staked. If we did
             not wait for finalization/inclusion, the response is `True`.
     """
+
     async def get_old_stakes() -> list[Balance]:
         old_stakes = []
         all_stakes = await subtensor.get_stake_for_coldkey(
@@ -224,8 +224,7 @@ async def get_old_stakes() -> list[Balance]:
         new_amounts = [None] * len(hotkey_ss58s)
     else:
         new_amounts = [
-            Balance.from_tao(amount) if not isinstance(amount, Balance) else amount
-            for amount in amounts
+            amount.set_unit(netuid=netuid) for amount, netuid in zip(amounts, netuids)
         ]
         if sum(amount.tao for amount in new_amounts) == 0:
             # Staking 0 tao
@@ -266,8 +265,7 @@ async def get_old_stakes() -> list[Balance]:
         # Reduce the amount to stake to each wallet to keep the balance above 1000 rao.
         percent_reduction = 1 - (1000 / total_staking_rao)
         new_amounts = [
-            Balance.from_tao(amount.tao * percent_reduction)
-            for amount in new_amounts
+            Balance.from_tao(amount.tao * percent_reduction) for amount in new_amounts
         ]
 
     successful_stakes = 0

bittensor/core/extrinsics/asyncex/unstaking.py

@@ -1,5 +1,5 @@
 import asyncio
-from typing import Union, Optional, TYPE_CHECKING
+from typing import Optional, TYPE_CHECKING
 
 from bittensor.core.errors import StakeError, NotRegisteredError
 from bittensor.utils import unlock_key
@@ -16,7 +16,7 @@ async def unstake_extrinsic(
     wallet: "Wallet",
     hotkey_ss58: Optional[str] = None,
     netuid: Optional[int] = None,
-    amount: Optional[Union[Balance, float]] = None,
+    amount: Optional[Balance] = None,
     wait_for_inclusion: bool = True,
     wait_for_finalization: bool = False,
 ) -> bool:
@@ -64,10 +64,9 @@ async def unstake_extrinsic(
     if amount is None:
         # Unstake it all.
         unstaking_balance = old_stake
-    elif not isinstance(amount, Balance):
-        unstaking_balance = Balance.from_tao(amount)
     else:
         unstaking_balance = amount
+        unstaking_balance.set_unit(netuid)
 
     # Check enough to unstake.
     stake_on_uid = old_stake
@@ -146,7 +145,7 @@ async def unstake_multiple_extrinsic(
     wallet: "Wallet",
     hotkey_ss58s: list[str],
     netuids: list[int],
-    amounts: Optional[list[Union[Balance, float]]] = None,
+    amounts: Optional[list[Balance]] = None,
     wait_for_inclusion: bool = True,
     wait_for_finalization: bool = False,
 ) -> bool:
@@ -212,11 +211,7 @@ async def get_old_stakes() -> list[Balance]:
         amounts = [None] * len(hotkey_ss58s)
     else:
         # Convert to Balance
-        amounts = [
-            Balance.from_tao(amount) if isinstance(amount, float) else amount
-            for amount in amounts
-        ]
-
+        amounts = [amount.set_unit(netuid) for amount, netuid in zip(amounts, netuids)]
         if sum(amount.tao for amount in amounts) == 0:
             # Staking 0 tao
             return True
@@ -244,9 +239,7 @@ async def get_old_stakes() -> list[Balance]:
             # Unstake it all.
             unstaking_balance = old_stake
         else:
-            unstaking_balance = (
-                amount if isinstance(amount, Balance) else Balance.from_tao(amount)
-            )
+            unstaking_balance = amount
 
         # Check enough to unstake.
         stake_on_uid = old_stake
@@ -267,7 +260,7 @@ async def get_old_stakes() -> list[Balance]:
                 call_function="remove_stake",
                 call_params={
                     "hotkey": hotkey_ss58,
-                    "amount_unstaked": amount.rao,
+                    "amount_unstaked": unstaking_balance.rao,
                     "netuid": netuid,
                 },
             )