Merge pull request #358 from opentensor/fix/stake-transfer-selection

Improves stake transfer, adds interactive selection of delegates

Author: ibraheem-abe

bittensor_cli/cli.py

@@ -3528,6 +3528,9 @@ def stake_transfer(
             "-a",
             help="Amount of stake to transfer",
         ),
+        stake_all: bool = typer.Option(
+            False, "--stake-all", "--all", help="Stake all", prompt=False
+        ),
         prompt: bool = Options.prompt,
         quiet: bool = Options.quiet,
         verbose: bool = Options.verbose,
@@ -3545,6 +3548,8 @@ def stake_transfer(
         - The destination subnet (--dest-netuid)
         - The destination wallet/address (--dest)
         - The amount to transfer (--amount)
+        - The origin wallet (--wallet-name)
+        - The origin hotkey wallet/address (--wallet-hotkey)
 
         If no arguments are provided, an interactive selection menu will be shown.
 
@@ -3553,14 +3558,37 @@ def stake_transfer(
         Transfer 100 TAO from subnet 1 to subnet 2:
         [green]$[/green] btcli stake transfer --origin-netuid 1 --dest-netuid 2 --dest wallet2 --amount 100
 
-        Using SS58 address:
+        Using Destination SS58 address:
         [green]$[/green] btcli stake transfer --origin-netuid 1 --dest-netuid 2 --dest 5FrLxJsyJ5x9n2rmxFwosFraxFCKcXZDngEP9H7qjkKgHLcK --amount 100
+
+        Using Origin hotkey SS58 address (useful when transferring stake from a delegate):
+        [green]$[/green] btcli stake transfer --wallet-hotkey 5FrLxJsyJ5x9n2rmxFwosFraxFCKcXZDngEP9H7qjkKgHLcK --wallet-name sample_wallet
+
+        Transfer all available stake from origin hotkey:
+        [green]$[/green] btcli stake transfer --all --origin-netuid 1 --dest-netuid 2
         """
         console.print(
             "[dim]This command transfers stake from one coldkey to another while keeping the same hotkey.[/dim]"
         )
         self.verbosity_handler(quiet, verbose)
 
+        if not dest_ss58:
+            dest_ss58 = Prompt.ask(
+                "Enter the [blue]destination wallet name[/blue] or [blue]coldkey SS58 address[/blue]"
+            )
+
+        if is_valid_ss58_address(dest_ss58):
+            dest_ss58 = dest_ss58
+        else:
+            dest_wallet = self.wallet_ask(
+                dest_ss58,
+                wallet_path,
+                None,
+                ask_for=[WO.NAME, WO.PATH],
+                validate=WV.WALLET,
+            )
+            dest_ss58 = dest_wallet.coldkeypub.ss58_address
+
         if not wallet_name:
             wallet_name = Prompt.ask(
                 "Enter the [blue]origin wallet name[/blue]",
@@ -3570,13 +3598,16 @@ def stake_transfer(
             wallet_name, wallet_path, wallet_hotkey, ask_for=[WO.NAME]
         )
 
+        interactive_selection = False
         if not wallet_hotkey:
             origin_hotkey = Prompt.ask(
-                "Enter the [blue]origin hotkey[/blue] name or "
-                "[blue]ss58 address[/blue] where the stake will be moved from",
-                default=self.config.get("wallet_hotkey") or defaults.wallet.hotkey,
+                "Enter the [blue]origin hotkey[/blue] name or ss58 address [bold](stake will be transferred FROM here)[/bold] "
+                "[dim](or press Enter to select from existing stakes)[/dim]"
             )
-            if is_valid_ss58_address(origin_hotkey):
+            if origin_hotkey == "":
+                interactive_selection = True
+
+            elif is_valid_ss58_address(origin_hotkey):
                 origin_hotkey = origin_hotkey
             else:
                 wallet = self.wallet_ask(
@@ -3600,33 +3631,11 @@ def stake_transfer(
                 )
                 origin_hotkey = wallet.hotkey.ss58_address
 
-        if not dest_ss58:
-            dest_ss58 = Prompt.ask(
-                "Enter the [blue]destination wallet name[/blue] or [blue]coldkey SS58 address[/blue]"
-            )
-
-        if is_valid_ss58_address(dest_ss58):
-            dest_ss58 = dest_ss58
-        else:
-            dest_wallet = self.wallet_ask(
-                dest_ss58,
-                wallet_path,
-                None,
-                ask_for=[WO.NAME, WO.PATH],
-                validate=WV.WALLET,
-            )
-            dest_ss58 = dest_wallet.coldkeypub.ss58_address
-
-        interactive_selection = False
-        if origin_netuid is None and dest_netuid is None and not amount:
-            interactive_selection = True
-        else:
+        if not interactive_selection:
             if origin_netuid is None:
                 origin_netuid = IntPrompt.ask(
                     "Enter the [blue]origin subnet[/blue] (netuid)"
                 )
-            if not amount:
-                amount = FloatPrompt.ask("Enter the [blue]amount[/blue] to transfer")
 
             if dest_netuid is None:
                 dest_netuid = IntPrompt.ask(
@@ -3643,6 +3652,7 @@ def stake_transfer(
                 dest_coldkey_ss58=dest_ss58,
                 amount=amount,
                 interactive_selection=interactive_selection,
+                stake_all=stake_all,
                 prompt=prompt,
             )
         )

bittensor_cli/src/commands/stake/move.py

@@ -207,7 +207,7 @@ def prompt_stake_amount(
             console.print("[red]Please enter a valid number or 'all'[/red]")
 
 
-async def stake_move_selection(
+async def stake_move_transfer_selection(
     subtensor: "SubtensorInterface",
     wallet: Wallet,
 ):
@@ -293,25 +293,24 @@ async def stake_move_selection(
         title_justify="center",
         width=len(origin_hotkey_ss58) + 20,
     )
-    table.add_column("Index", justify="right")
     table.add_column("Netuid", style="cyan")
     table.add_column("Stake Amount", style=COLOR_PALETTE["STAKE"]["STAKE_AMOUNT"])
 
     available_netuids = []
-    for idx, netuid in enumerate(origin_hotkey_info["netuids"]):
+    for netuid in origin_hotkey_info["netuids"]:
         stake = origin_hotkey_info["stakes"][netuid]
         if stake.tao > 0:
             available_netuids.append(netuid)
-            table.add_row(str(idx), str(netuid), str(stake))
+            table.add_row(str(netuid), str(stake))
 
     console.print("\n", table)
 
     # Select origin netuid
-    netuid_idx = Prompt.ask(
-        "\nEnter the index of the subnet you want to move stake from",
-        choices=[str(i) for i in range(len(available_netuids))],
+    origin_netuid = Prompt.ask(
+        "\nEnter the netuid you want to move stake from",
+        choices=[str(netuid) for netuid in available_netuids],
     )
-    origin_netuid = available_netuids[int(netuid_idx)]
+    origin_netuid = int(origin_netuid)
     origin_stake = origin_hotkey_info["stakes"][origin_netuid]
 
     # Ask for amount to move
@@ -334,104 +333,6 @@ async def stake_move_selection(
     }
 
 
-async def stake_transfer_selection(
-    wallet: Wallet, subtensor: "SubtensorInterface", origin_hotkey: str
-):
-    """Selection interface for transferring stakes."""
-    (
-        stakes,
-        all_netuids,
-        all_subnets,
-    ) = await asyncio.gather(
-        subtensor.get_stake_for_coldkey(coldkey_ss58=wallet.coldkeypub.ss58_address),
-        subtensor.get_all_subnet_netuids(),
-        subtensor.all_subnets(),
-    )
-    all_netuids = sorted(all_netuids)
-    all_subnets = {di.netuid: di for di in all_subnets}
-
-    available_stakes = {}
-    for stake in stakes:
-        if stake.stake.tao > 0 and stake.hotkey_ss58 == origin_hotkey:
-            available_stakes[stake.netuid] = {
-                "hotkey_ss58": stake.hotkey_ss58,
-                "stake": stake.stake,
-                "is_registered": stake.is_registered,
-            }
-
-    if not available_stakes:
-        console.print("[red]No stakes available to transfer.[/red]")
-        return None
-
-    table = Table(
-        title=(
-            f"\n[{COLOR_PALETTE['GENERAL']['HEADER']}]"
-            f"Available Stakes to Transfer\n"
-            f"for wallet hotkey:\n"
-            f"[{COLOR_PALETTE['GENERAL']['HOTKEY']}]{wallet.hotkey_str}: {wallet.hotkey.ss58_address}"
-            f"[/{COLOR_PALETTE['GENERAL']['HOTKEY']}]\n"
-        ),
-        show_edge=False,
-        header_style="bold white",
-        border_style="bright_black",
-        title_justify="center",
-        width=len(wallet.hotkey_str + wallet.hotkey.ss58_address) + 10,
-    )
-
-    table.add_column("Index", justify="right", style="cyan")
-    table.add_column("Netuid")
-    table.add_column("Name", style="cyan", justify="left")
-    table.add_column("Stake Amount", style=COLOR_PALETTE["STAKE"]["STAKE_AMOUNT"])
-    table.add_column("Registered", justify="center")
-
-    for idx, (netuid, stake_info) in enumerate(available_stakes.items()):
-        subnet_name_cell = (
-            f"[{COLOR_PALETTE['GENERAL']['SYMBOL']}]{all_subnets[netuid].symbol if netuid != 0 else 'τ'}[/{COLOR_PALETTE['GENERAL']['SYMBOL']}]"
-            f" {get_subnet_name(all_subnets[netuid])}"
-        )
-        table.add_row(
-            str(idx),
-            str(netuid),
-            subnet_name_cell,
-            str(stake_info["stake"]),
-            "[dark_sea_green3]YES"
-            if stake_info["is_registered"]
-            else f"[{COLOR_PALETTE['STAKE']['NOT_REGISTERED']}]NO",
-        )
-
-    console.print(table)
-
-    if not available_stakes:
-        console.print("[red]No stakes available to transfer.[/red]")
-        return None
-
-    # Prompt to select index of stake to transfer
-    selection = Prompt.ask(
-        "\nEnter the index of the stake you want to transfer",
-        choices=[str(i) for i in range(len(available_stakes))],
-    )
-    selected_netuid = list(available_stakes.keys())[int(selection)]
-    selected_stake = available_stakes[selected_netuid]
-
-    # Prompt for amount
-    stake_balance = selected_stake["stake"]
-    amount, _ = prompt_stake_amount(stake_balance, selected_netuid, "transfer")
-
-    # Prompt for destination subnet
-    destination_netuid = Prompt.ask(
-        "\nEnter the netuid of the subnet you want to move stake to"
-        + f" ([dim]{group_subnets(all_netuids)}[/dim])",
-        choices=[str(netuid) for netuid in all_netuids],
-        show_choices=False,
-    )
-
-    return {
-        "origin_netuid": selected_netuid,
-        "amount": amount.tao,
-        "destination_netuid": int(destination_netuid),
-    }
-
-
 async def stake_swap_selection(
     subtensor: "SubtensorInterface",
     wallet: Wallet,
@@ -540,7 +441,7 @@ async def move_stake(
     prompt: bool = True,
 ):
     if interactive_selection:
-        selection = await stake_move_selection(subtensor, wallet)
+        selection = await stake_move_transfer_selection(subtensor, wallet)
         origin_hotkey = selection["origin_hotkey"]
         origin_netuid = selection["origin_netuid"]
         amount = selection["amount"]
@@ -699,6 +600,7 @@ async def transfer_stake(
     dest_netuid: int,
     dest_coldkey_ss58: str,
     interactive_selection: bool = False,
+    stake_all: bool = False,
     prompt: bool = True,
 ) -> bool:
     """Transfers stake from one network to another.
@@ -717,12 +619,13 @@ async def transfer_stake(
     Returns:
         bool: True if transfer was successful, False otherwise.
     """
-    origin_hotkey = origin_hotkey or wallet.hotkey.ss58_address
     if interactive_selection:
-        selection = await stake_transfer_selection(wallet, subtensor, origin_hotkey)
+        selection = await stake_move_transfer_selection(subtensor, wallet)
         origin_netuid = selection["origin_netuid"]
         amount = selection["amount"]
         dest_netuid = selection["destination_netuid"]
+        stake_all = selection["stake_all"]
+        origin_hotkey = selection["origin_hotkey"]
 
     # Check if both subnets exist
     block_hash = await subtensor.substrate.get_chain_head()
@@ -750,7 +653,22 @@ async def transfer_stake(
             hotkey_ss58=origin_hotkey,
             netuid=dest_netuid,
         )
-    amount_to_transfer = Balance.from_tao(amount).set_unit(origin_netuid)
+
+    if current_stake.tao == 0:
+        err_console.print(
+            f"[red]No stake found for hotkey: {origin_hotkey} on netuid: {origin_netuid}[/red]"
+        )
+        return False
+
+    amount_to_transfer = None
+    if amount:
+        amount_to_transfer = Balance.from_tao(amount).set_unit(origin_netuid)
+    elif stake_all:
+        amount_to_transfer = current_stake
+    else:
+        amount_to_transfer, _ = prompt_stake_amount(
+            current_stake, origin_netuid, "transfer"
+        )
 
     # Check if enough stake to transfer
     if amount_to_transfer > current_stake: