Merge pull request #2993 from opentensor/fix/thewhaleking/transfers

Transfers improvements

Author: thewhaleking

bittensor/core/async_subtensor.py

@@ -98,6 +98,7 @@
     torch,
     u16_normalized_float,
     u64_normalized_float,
+    get_transfer_fn_params,
 )
 from bittensor.core.extrinsics.asyncex.liquidity import (
     add_liquidity_extrinsic,
@@ -3134,7 +3135,7 @@ async def get_total_subnets(
         return getattr(result, "value", None)
 
     async def get_transfer_fee(
-        self, wallet: "Wallet", dest: str, value: Balance
+        self, wallet: "Wallet", dest: str, value: Balance, keep_alive: bool = True
     ) -> Balance:
         """
         Calculates the transaction fee for transferring tokens from a wallet to a specified destination address. This
@@ -3146,6 +3147,8 @@ async def get_transfer_fee(
             dest: The ``SS58`` address of the destination account.
             value: The amount of tokens to be transferred, specified as a Balance object, or in Tao (float) or Rao
                 (int) units.
+            keep_alive: Whether the transfer fee should be calculated based on keeping the wallet alive (existential
+                deposit) or not.
 
         Returns:
             bittensor.utils.balance.Balance: The estimated transaction fee for the transfer, represented as a Balance
@@ -3155,12 +3158,15 @@ async def get_transfer_fee(
         wallet has sufficient funds to cover both the transfer amount and the associated costs. This function provides
         a crucial tool for managing financial operations within the Bittensor network.
         """
-        value = check_and_convert_to_balance(value)
+        if value is not None:
+            value = check_and_convert_to_balance(value)
+        call_params: dict[str, Union[int, str, bool]]
+        call_function, call_params = get_transfer_fn_params(value, dest, keep_alive)
 
         call = await self.substrate.compose_call(
             call_module="Balances",
-            call_function="transfer_keep_alive",
-            call_params={"dest": dest, "value": value.rao},
+            call_function=call_function,
+            call_params=call_params,
         )
 
         try:
@@ -5455,7 +5461,7 @@ async def transfer(
         self,
         wallet: "Wallet",
         dest: str,
-        amount: Balance,
+        amount: Optional[Balance],
         transfer_all: bool = False,
         wait_for_inclusion: bool = True,
         wait_for_finalization: bool = False,
@@ -5468,7 +5474,7 @@ async def transfer(
         Arguments:
             wallet: Source wallet for the transfer.
             dest: Destination address for the transfer.
-            amount: Number of tokens to transfer.
+            amount: Number of tokens to transfer. `None` is transferring all.
             transfer_all: Flag to transfer all tokens. Default is `False`.
             wait_for_inclusion: Waits for the transaction to be included in a block. Defaults to `True`.
             wait_for_finalization: Waits for the transaction to be finalized on the blockchain. Defaults to `False`.
@@ -5479,7 +5485,8 @@ async def transfer(
         Returns:
             `True` if the transferring was successful, otherwise `False`.
         """
-        amount = check_and_convert_to_balance(amount)
+        if amount is not None:
+            amount = check_and_convert_to_balance(amount)
         return await transfer_extrinsic(
             subtensor=self,
             wallet=wallet,

bittensor/core/extrinsics/asyncex/transfer.py

@@ -6,6 +6,7 @@
     get_explorer_url_for_network,
     is_valid_bittensor_address_or_public_key,
     unlock_key,
+    get_transfer_fn_params,
 )
 from bittensor.utils.balance import Balance
 from bittensor.utils.btlogging import logging
@@ -19,10 +20,11 @@ async def _do_transfer(
     subtensor: "AsyncSubtensor",
     wallet: "Wallet",
     destination: str,
-    amount: "Balance",
+    amount: Optional[Balance],
     wait_for_inclusion: bool = True,
     wait_for_finalization: bool = False,
     period: Optional[int] = None,
+    keep_alive: bool = True,
 ) -> tuple[bool, str, str]:
     """
     Makes transfer from wallet to destination public key address.
@@ -39,14 +41,17 @@ async def _do_transfer(
         period (Optional[int]): The number of blocks during which the transaction will remain valid after it's submitted.
             If the transaction is not included in a block within that number of blocks, it will expire and be rejected.
             You can think of it as an expiration date for the transaction.
+        keep_alive (bool): If `True`, will keep the existential deposit in the account.
 
     Returns:
         success, block hash, formatted error message
     """
+    call_function, call_params = get_transfer_fn_params(amount, destination, keep_alive)
+
     call = await subtensor.substrate.compose_call(
         call_module="Balances",
-        call_function="transfer_keep_alive",
-        call_params={"dest": destination, "value": amount.rao},
+        call_function=call_function,
+        call_params=call_params,
     )
 
     success, message = await subtensor.sign_and_send_extrinsic(
@@ -73,7 +78,7 @@ async def transfer_extrinsic(
     subtensor: "AsyncSubtensor",
     wallet: "Wallet",
     dest: str,
-    amount: "Balance",
+    amount: Optional[Balance],
     transfer_all: bool = False,
     wait_for_inclusion: bool = True,
     wait_for_finalization: bool = False,
@@ -86,7 +91,8 @@ async def transfer_extrinsic(
         subtensor (bittensor.core.async_subtensor.AsyncSubtensor): initialized AsyncSubtensor object used for transfer
         wallet (bittensor_wallet.Wallet): Bittensor wallet object to make transfer from.
         dest (str): Destination public key address (ss58_address or ed25519) of recipient.
-        amount (bittensor.utils.balance.Balance): Amount to stake as Bittensor balance.
+        amount (Optional[bittensor.utils.balance.Balance]): Amount to stake as Bittensor balance. `None` if
+            transferring all.
         transfer_all (bool): Whether to transfer all funds from this wallet to the destination address.
         wait_for_inclusion (bool): If set, waits for the extrinsic to enter a block before returning `True`, or returns
             `False` if the extrinsic fails to enter the block within the timeout.
@@ -102,6 +108,11 @@ async def transfer_extrinsic(
             finalization / inclusion, the response is `True`, regardless of its inclusion.
     """
     destination = dest
+
+    if amount is None and not transfer_all:
+        logging.error("If not transferring all, `amount` must be specified.")
+        return False
+
     # Validate destination address.
     if not is_valid_bittensor_address_or_public_key(destination):
         logging.error(
@@ -128,7 +139,7 @@ async def transfer_extrinsic(
     )
 
     fee = await subtensor.get_transfer_fee(
-        wallet=wallet, dest=destination, value=amount
+        wallet=wallet, dest=destination, value=amount, keep_alive=keep_alive
     )
 
     if not keep_alive:
@@ -137,12 +148,10 @@ async def transfer_extrinsic(
 
     # Check if we have enough balance.
     if transfer_all is True:
-        amount = account_balance - fee - existential_deposit
-        if amount < Balance(0):
+        if (account_balance - fee) < existential_deposit:
             logging.error("Not enough balance to transfer")
             return False
-
-    if account_balance < (amount + fee + existential_deposit):
+    elif account_balance < (amount + fee + existential_deposit):
         logging.error(":cross_mark: [red]Not enough balance[/red]")
         logging.error(f"\t\tBalance:\t[blue]{account_balance}[/blue]")
         logging.error(f"\t\tAmount:\t[blue]{amount}[/blue]")
@@ -155,6 +164,7 @@ async def transfer_extrinsic(
         wallet=wallet,
         destination=destination,
         amount=amount,
+        keep_alive=keep_alive,
         wait_for_finalization=wait_for_finalization,
         wait_for_inclusion=wait_for_inclusion,
         period=period,

bittensor/core/extrinsics/transfer.py

@@ -5,6 +5,7 @@
     is_valid_bittensor_address_or_public_key,
     unlock_key,
     get_explorer_url_for_network,
+    get_transfer_fn_params,
 )
 from bittensor.utils.balance import Balance
 from bittensor.utils.btlogging import logging
@@ -18,10 +19,11 @@ def _do_transfer(
     subtensor: "Subtensor",
     wallet: "Wallet",
     destination: str,
-    amount: Balance,
+    amount: Optional[Balance],
     wait_for_inclusion: bool = True,
     wait_for_finalization: bool = False,
     period: Optional[int] = None,
+    keep_alive: bool = True,
 ) -> tuple[bool, str, str]:
     """
     Makes transfer from wallet to destination public key address.
@@ -30,22 +32,25 @@ def _do_transfer(
         subtensor (bittensor.core.subtensor.Subtensor): the Subtensor object used for transfer
         wallet (bittensor_wallet.Wallet): Bittensor wallet object to make transfer from.
         destination (str): Destination public key address (ss58_address or ed25519) of recipient.
-        amount (bittensor.utils.balance.Balance): Amount to stake as Bittensor balance.
+        amount (bittensor.utils.balance.Balance): Amount to stake as Bittensor balance. `None` if transferring all.
         wait_for_inclusion (bool): If set, waits for the extrinsic to enter a block before returning `True`, or returns
             `False` if the extrinsic fails to enter the block within the timeout.
         wait_for_finalization (bool): If set, waits for the extrinsic to be finalized on the chain before returning
             `True`, or returns `False` if the extrinsic fails to be finalized within the timeout.
         period (Optional[int]): The number of blocks during which the transaction will remain valid after it's submitted.
             If the transaction is not included in a block within that number of blocks, it will expire and be rejected.
             You can think of it as an expiration date for the transaction.
+        keep_alive (bool): If `True`, will keep the existential deposit in the account.
 
     Returns:
         success, block hash, formatted error message
     """
+    call_function, call_params = get_transfer_fn_params(amount, destination, keep_alive)
+
     call = subtensor.substrate.compose_call(
         call_module="Balances",
-        call_function="transfer_keep_alive",
-        call_params={"dest": destination, "value": amount.rao},
+        call_function=call_function,
+        call_params=call_params,
     )
 
     success, message = subtensor.sign_and_send_extrinsic(
@@ -72,7 +77,7 @@ def transfer_extrinsic(
     subtensor: "Subtensor",
     wallet: "Wallet",
     dest: str,
-    amount: Balance,
+    amount: Optional[Balance],
     transfer_all: bool = False,
     wait_for_inclusion: bool = True,
     wait_for_finalization: bool = False,
@@ -85,7 +90,7 @@ def transfer_extrinsic(
         subtensor (bittensor.core.subtensor.Subtensor): the Subtensor object used for transfer
         wallet (bittensor_wallet.Wallet): Bittensor wallet object to make transfer from.
         dest (str): Destination public key address (ss58_address or ed25519) of recipient.
-        amount (bittensor.utils.balance.Balance): Amount to stake as Bittensor balance.
+        amount (bittensor.utils.balance.Balance): Amount to stake as Bittensor balance. `None` if transferring all.
         transfer_all (bool): Whether to transfer all funds from this wallet to the destination address.
         wait_for_inclusion (bool): If set, waits for the extrinsic to enter a block before returning `True`, or returns
             `False` if the extrinsic fails to enter the block within the timeout.
@@ -100,6 +105,10 @@ def transfer_extrinsic(
         success (bool): Flag is `True` if extrinsic was finalized or included in the block. If we did not wait for
             finalization / inclusion, the response is `True`, regardless of its inclusion.
     """
+    if amount is None and not transfer_all:
+        logging.error("If not transferring all, `amount` must be specified.")
+        return False
+
     # Validate destination address.
     if not is_valid_bittensor_address_or_public_key(dest):
         logging.error(
@@ -127,16 +136,16 @@ def transfer_extrinsic(
     else:
         existential_deposit = subtensor.get_existential_deposit(block=block)
 
-    fee = subtensor.get_transfer_fee(wallet=wallet, dest=dest, value=amount)
+    fee = subtensor.get_transfer_fee(
+        wallet=wallet, dest=dest, value=amount, keep_alive=keep_alive
+    )
 
     # Check if we have enough balance.
     if transfer_all is True:
-        amount = account_balance - fee - existential_deposit
-        if amount < Balance(0):
+        if (account_balance - fee) < existential_deposit:
             logging.error("Not enough balance to transfer")
             return False
-
-    if account_balance < (amount + fee + existential_deposit):
+    elif account_balance < (amount + fee + existential_deposit):
         logging.error(":cross_mark: [red]Not enough balance[/red]")
         logging.error(f"\t\tBalance:\t[blue]{account_balance}[/blue]")
         logging.error(f"\t\tAmount:\t[blue]{amount}[/blue]")
@@ -149,6 +158,7 @@ def transfer_extrinsic(
         wallet=wallet,
         destination=dest,
         amount=amount,
+        keep_alive=keep_alive,
         wait_for_finalization=wait_for_finalization,
         wait_for_inclusion=wait_for_inclusion,
         period=period,

bittensor/core/subtensor.py

@@ -107,6 +107,7 @@
     u16_normalized_float,
     u64_normalized_float,
     deprecated_message,
+    get_transfer_fn_params,
 )
 from bittensor.utils.balance import (
     Balance,
@@ -2243,7 +2244,13 @@ def get_total_subnets(self, block: Optional[int] = None) -> Optional[int]:
         )
         return getattr(result, "value", None)
 
-    def get_transfer_fee(self, wallet: "Wallet", dest: str, value: Balance) -> Balance:
+    def get_transfer_fee(
+        self,
+        wallet: "Wallet",
+        dest: str,
+        value: Optional[Balance],
+        keep_alive: bool = True,
+    ) -> Balance:
         """
         Calculates the transaction fee for transferring tokens from a wallet to a specified destination address. This
             function simulates the transfer to estimate the associated cost, taking into account the current network
@@ -2254,6 +2261,8 @@ def get_transfer_fee(self, wallet: "Wallet", dest: str, value: Balance) -> Balan
             dest (str): The ``SS58`` address of the destination account.
             value (Union[bittensor.utils.balance.Balance, float, int]): The amount of tokens to be transferred,
                 specified as a Balance object, or in Tao (float) or Rao (int) units.
+            keep_alive: Whether the transfer fee should be calculated based on keeping the wallet alive (existential
+                deposit) or not.
 
         Returns:
             bittensor.utils.balance.Balance: The estimated transaction fee for the transfer, represented as a Balance
@@ -2263,11 +2272,15 @@ def get_transfer_fee(self, wallet: "Wallet", dest: str, value: Balance) -> Balan
             has sufficient funds to cover both the transfer amount and the associated costs. This function provides a
             crucial tool for managing financial operations within the Bittensor network.
         """
-        value = check_and_convert_to_balance(value)
+        if value is not None:
+            value = check_and_convert_to_balance(value)
+        call_params: dict[str, Union[int, str, bool]]
+        call_function, call_params = get_transfer_fn_params(value, dest, keep_alive)
+
         call = self.substrate.compose_call(
             call_module="Balances",
-            call_function="transfer_keep_alive",
-            call_params={"dest": dest, "value": value.rao},
+            call_function=call_function,
+            call_params=call_params,
         )
 
         try:
@@ -4266,7 +4279,7 @@ def transfer(
         self,
         wallet: "Wallet",
         dest: str,
-        amount: Balance,
+        amount: Optional[Balance],
         wait_for_inclusion: bool = True,
         wait_for_finalization: bool = False,
         transfer_all: bool = False,
@@ -4292,7 +4305,8 @@ def transfer(
         Returns:
             `True` if the transferring was successful, otherwise `False`.
         """
-        amount = check_and_convert_to_balance(amount)
+        if amount is not None:
+            amount = check_and_convert_to_balance(amount)
         return transfer_extrinsic(
             subtensor=self,
             wallet=wallet,

bittensor/utils/__init__.py

@@ -22,6 +22,7 @@
 
 if TYPE_CHECKING:
     from bittensor_wallet import Wallet
+    from bittensor.utils.balance import Balance
 
 BT_DOCS_LINK = "https://docs.bittensor.com"
 
@@ -445,3 +446,34 @@ def deprecated_message(message: str) -> None:
     """Shows a deprecation warning message with the given message."""
     warnings.simplefilter("default", DeprecationWarning)
     warnings.warn(message=message, category=DeprecationWarning, stacklevel=2)
+
+
+def get_transfer_fn_params(
+    amount: Optional["Balance"], destination: str, keep_alive: bool
+) -> tuple[str, dict[str, Union[str, int, bool]]]:
+    """
+    Helper function to get the transfer call function and call params, depending on the value and keep_alive flag
+        provided
+
+    Args:
+        amount: the amount of Tao to transfer. `None` if transferring all.
+        destination: the destination SS58 of the transfer
+        keep_alive: whether to enforce a retention of the existential deposit in the account after transfer.
+
+    Returns:
+        tuple[call function, call params]
+    """
+    call_params = {"dest": destination}
+    if amount is None:
+        call_function = "transfer_all"
+        if keep_alive:
+            call_params["keep_alive"] = True
+        else:
+            call_params["keep_alive"] = False
+    else:
+        call_params["value"] = amount.rao
+        if keep_alive:
+            call_function = "transfer_keep_alive"
+        else:
+            call_function = "transfer_allow_death"
+    return call_function, call_params