stake remove

Author: thewhaleking

bittensor_cli/cli.py

@@ -3192,6 +3192,7 @@ def stake_remove(
         ),
         quiet: bool = Options.quiet,
         verbose: bool = Options.verbose,
+        json_output: bool = Options.json_output,
     ):
         """
         Unstake TAO from one or more hotkeys and transfer them back to the user's coldkey wallet.
@@ -3223,7 +3224,7 @@ def stake_remove(
         • [blue]--tolerance[/blue]: Max allowed rate change (0.05 = 5%)
         • [blue]--partial[/blue]: Complete partial unstake if rates exceed tolerance
         """
-        self.verbosity_handler(quiet, verbose)
+        self.verbosity_handler(quiet, verbose, json_output)
         if not unstake_all and not unstake_all_alpha:
             safe_staking = self.ask_safe_staking(safe_staking)
             if safe_staking:
@@ -3235,7 +3236,8 @@ def stake_remove(
             [hotkey_ss58_address, include_hotkeys, exclude_hotkeys, all_hotkeys]
         ):
             print_error(
-                "Interactive mode cannot be used with hotkey selection options like --include-hotkeys, --exclude-hotkeys, --all-hotkeys, or --hotkey."
+                "Interactive mode cannot be used with hotkey selection options like "
+                "--include-hotkeys, --exclude-hotkeys, --all-hotkeys, or --hotkey."
             )
             raise typer.Exit()
 
@@ -3372,6 +3374,7 @@ def stake_remove(
                     include_hotkeys=include_hotkeys,
                     exclude_hotkeys=exclude_hotkeys,
                     prompt=prompt,
+                    json_output=json_output,
                 )
             )
         elif (
@@ -3426,6 +3429,7 @@ def stake_remove(
                 safe_staking=safe_staking,
                 rate_tolerance=rate_tolerance,
                 allow_partial_stake=allow_partial_stake,
+                json_output=json_output,
             )
         )
 

bittensor_cli/src/commands/stake/remove.py

@@ -1,4 +1,5 @@
 import asyncio
+import json
 from functools import partial
 
 from typing import TYPE_CHECKING, Optional
@@ -20,6 +21,7 @@
     format_error_message,
     group_subnets,
     unlock_key,
+    json_console,
 )
 
 if TYPE_CHECKING:
@@ -41,6 +43,7 @@ async def unstake(
     safe_staking: bool,
     rate_tolerance: float,
     allow_partial_stake: bool,
+    json_output: bool,
 ):
     """Unstake from hotkey(s)."""
     unstake_all_from_hk = False
@@ -241,8 +244,11 @@ async def unstake(
                 base_unstake_op["price_with_tolerance"] = price_with_tolerance
                 base_table_row.extend(
                     [
-                        f"{rate_with_tolerance:.4f} {Balance.get_unit(0)}/{Balance.get_unit(netuid)}",  # Rate with tolerance
-                        f"[{'dark_sea_green3' if allow_partial_stake else 'red'}]{allow_partial_stake}[/{'dark_sea_green3' if allow_partial_stake else 'red'}]",  # Partial unstake
+                        # Rate with tolerance
+                        f"{rate_with_tolerance:.4f} {Balance.get_unit(0)}/{Balance.get_unit(netuid)}",
+                        # Partial unstake
+                        f"[{'dark_sea_green3' if allow_partial_stake else 'red'}]"
+                        f"{allow_partial_stake}[/{'dark_sea_green3' if allow_partial_stake else 'red'}]",
                     ]
                 )
 
@@ -273,45 +279,45 @@ async def unstake(
     if not unlock_key(wallet).success:
         return False
 
+    successes = []
     with console.status("\n:satellite: Performing unstaking operations...") as status:
-        if safe_staking:
-            for op in unstake_operations:
-                if op["netuid"] == 0:
-                    await _unstake_extrinsic(
-                        wallet=wallet,
-                        subtensor=subtensor,
-                        netuid=op["netuid"],
-                        amount=op["amount_to_unstake"],
-                        current_stake=op["current_stake_balance"],
-                        hotkey_ss58=op["hotkey_ss58"],
-                        status=status,
-                    )
-                else:
-                    await _safe_unstake_extrinsic(
-                        wallet=wallet,
-                        subtensor=subtensor,
-                        netuid=op["netuid"],
-                        amount=op["amount_to_unstake"],
-                        current_stake=op["current_stake_balance"],
-                        hotkey_ss58=op["hotkey_ss58"],
-                        price_limit=op["price_with_tolerance"],
-                        allow_partial_stake=allow_partial_stake,
-                        status=status,
-                    )
-        else:
-            for op in unstake_operations:
-                await _unstake_extrinsic(
-                    wallet=wallet,
-                    subtensor=subtensor,
-                    netuid=op["netuid"],
-                    amount=op["amount_to_unstake"],
-                    current_stake=op["current_stake_balance"],
-                    hotkey_ss58=op["hotkey_ss58"],
-                    status=status,
-                )
+        for op in unstake_operations:
+            common_args = {
+                "wallet": wallet,
+                "subtensor": subtensor,
+                "netuid": op["netuid"],
+                "amount": op["amount_to_unstake"],
+                "current_stake": op["current_stake_balance"],
+                "hotkey_ss58": op["hotkey_ss58"],
+                "status": status,
+            }
+
+            if safe_staking and op["netuid"] != 0:
+                func = _safe_unstake_extrinsic
+                specific_args = {
+                    "price_limit": op["price_with_tolerance"],
+                    "allow_partial_stake": allow_partial_stake,
+                }
+            else:
+                func = _unstake_extrinsic
+                specific_args = {}
+
+            suc = await func(**common_args, **specific_args)
+
+            successes.append(
+                {
+                    "netuid": op["netuid"],
+                    "hotkey_ss58": op["hotkey_ss58"],
+                    "unstake_amount": op["amount_to_unstake"].tao,
+                    "success": suc,
+                }
+            )
+
     console.print(
         f"[{COLOR_PALETTE['STAKE']['STAKE_AMOUNT']}]Unstaking operations completed."
     )
+    if json_output:
+        json_console.print(json.dumps(successes))
 
 
 async def unstake_all(
@@ -323,6 +329,7 @@ async def unstake_all(
     include_hotkeys: list[str] = [],
     exclude_hotkeys: list[str] = [],
     prompt: bool = True,
+    json_output: bool = False,
 ) -> bool:
     """Unstakes all stakes from all hotkeys in all subnets."""
 
@@ -448,11 +455,16 @@ async def unstake_all(
                 slippage_pct,
             )
     console.print(table)
-    message = ""
     if max_slippage > 5:
-        message += f"[{COLOR_PALETTE['STAKE']['SLIPPAGE_TEXT']}]-------------------------------------------------------------------------------------------------------------------\n"
-        message += f"[bold]WARNING:[/bold] The slippage on one of your operations is high: [{COLOR_PALETTE['STAKE']['SLIPPAGE_PERCENT']}]{max_slippage:.4f}%[/{COLOR_PALETTE['STAKE']['SLIPPAGE_PERCENT']}], this may result in a loss of funds.\n"
-        message += "-------------------------------------------------------------------------------------------------------------------\n"
+        message = (
+            f"[{COLOR_PALETTE['STAKE']['SLIPPAGE_TEXT']}]--------------------------------------------------------------"
+            f"-----------------------------------------------------\n"
+            f"[bold]WARNING:[/bold] The slippage on one of your operations is high: "
+            f"[{COLOR_PALETTE['STAKE']['SLIPPAGE_PERCENT']}]{max_slippage:.4f}%"
+            f"[/{COLOR_PALETTE['STAKE']['SLIPPAGE_PERCENT']}], this may result in a loss of funds.\n"
+            "----------------------------------------------------------------------------------------------------------"
+            "---------\n"
+        )
         console.print(message)
 
     console.print(
@@ -466,17 +478,19 @@ async def unstake_all(
 
     if not unlock_key(wallet).success:
         return False
-
+    successes = {}
     with console.status("Unstaking all stakes...") as status:
         for hotkey_ss58 in hotkey_ss58s:
-            await _unstake_all_extrinsic(
+            successes[hotkey_ss58] = await _unstake_all_extrinsic(
                 wallet=wallet,
                 subtensor=subtensor,
                 hotkey_ss58=hotkey_ss58,
                 hotkey_name=hotkey_names.get(hotkey_ss58, hotkey_ss58),
                 unstake_all_alpha=unstake_all_alpha,
                 status=status,
             )
+    if json_output:
+        return json_console.print(json.dumps({"success": successes}))
 
 
 # Extrinsics
@@ -488,7 +502,7 @@ async def _unstake_extrinsic(
     current_stake: Balance,
     hotkey_ss58: str,
     status=None,
-) -> None:
+) -> bool:
     """Execute a standard unstake extrinsic.
 
     Args:
@@ -510,15 +524,17 @@ async def _unstake_extrinsic(
             f"\n:satellite: Unstaking {amount} from {hotkey_ss58} on netuid: {netuid} ..."
         )
 
-    current_balance = await subtensor.get_balance(wallet.coldkeypub.ss58_address)
-    call = await subtensor.substrate.compose_call(
-        call_module="SubtensorModule",
-        call_function="remove_stake",
-        call_params={
-            "hotkey": hotkey_ss58,
-            "netuid": netuid,
-            "amount_unstaked": amount.rao,
-        },
+    current_balance, call = await asyncio.gather(
+        subtensor.get_balance(wallet.coldkeypub.ss58_address),
+        subtensor.substrate.compose_call(
+            call_module="SubtensorModule",
+            call_function="remove_stake",
+            call_params={
+                "hotkey": hotkey_ss58,
+                "netuid": netuid,
+                "amount_unstaked": amount.rao,
+            },
+        ),
     )
     extrinsic = await subtensor.substrate.create_signed_extrinsic(
         call=call, keypair=wallet.coldkey
@@ -528,15 +544,12 @@ async def _unstake_extrinsic(
         response = await subtensor.substrate.submit_extrinsic(
             extrinsic, wait_for_inclusion=True, wait_for_finalization=False
         )
-        await response.process_events()
-
         if not await response.is_success:
             err_out(
                 f"{failure_prelude} with error: "
                 f"{format_error_message(await response.error_message)}"
             )
-            return
-
+            return False
         # Fetch latest balance and stake
         block_hash = await subtensor.substrate.get_chain_head()
         new_balance, new_stake = await asyncio.gather(
@@ -557,9 +570,11 @@ async def _unstake_extrinsic(
             f"Subnet: [{COLOR_PALETTE['GENERAL']['SUBHEADING']}]{netuid}[/{COLOR_PALETTE['GENERAL']['SUBHEADING']}]"
             f" Stake:\n  [blue]{current_stake}[/blue] :arrow_right: [{COLOR_PALETTE['STAKE']['STAKE_AMOUNT']}]{new_stake}"
         )
+        return True
 
     except Exception as e:
         err_out(f"{failure_prelude} with error: {str(e)}")
+        return False
 
 
 async def _safe_unstake_extrinsic(
@@ -572,7 +587,7 @@ async def _safe_unstake_extrinsic(
     price_limit: Balance,
     allow_partial_stake: bool,
     status=None,
-) -> None:
+) -> bool:
     """Execute a safe unstake extrinsic with price limit.
 
     Args:
@@ -598,26 +613,27 @@ async def _safe_unstake_extrinsic(
 
     block_hash = await subtensor.substrate.get_chain_head()
 
-    current_balance, next_nonce, current_stake = await asyncio.gather(
+    current_balance, next_nonce, current_stake, call = await asyncio.gather(
         subtensor.get_balance(wallet.coldkeypub.ss58_address, block_hash),
         subtensor.substrate.get_account_next_index(wallet.coldkeypub.ss58_address),
         subtensor.get_stake(
             hotkey_ss58=hotkey_ss58,
             coldkey_ss58=wallet.coldkeypub.ss58_address,
             netuid=netuid,
+            block_hash=block_hash,
+        ),
+        subtensor.substrate.compose_call(
+            call_module="SubtensorModule",
+            call_function="remove_stake_limit",
+            call_params={
+                "hotkey": hotkey_ss58,
+                "netuid": netuid,
+                "amount_unstaked": amount.rao,
+                "limit_price": price_limit,
+                "allow_partial": allow_partial_stake,
+            },
+            block_hash=block_hash,
         ),
-    )
-
-    call = await subtensor.substrate.compose_call(
-        call_module="SubtensorModule",
-        call_function="remove_stake_limit",
-        call_params={
-            "hotkey": hotkey_ss58,
-            "netuid": netuid,
-            "amount_unstaked": amount.rao,
-            "limit_price": price_limit,
-            "allow_partial": allow_partial_stake,
-        },
     )
 
     extrinsic = await subtensor.substrate.create_signed_extrinsic(
@@ -636,17 +652,15 @@ async def _safe_unstake_extrinsic(
                 f"Either increase price tolerance or enable partial unstaking.",
                 status=status,
             )
-            return
         else:
             err_out(f"\n{failure_prelude} with error: {format_error_message(e)}")
-        return
+        return False
 
-    await response.process_events()
     if not await response.is_success:
         err_out(
             f"\n{failure_prelude} with error: {format_error_message(await response.error_message)}"
         )
-        return
+        return False
 
     block_hash = await subtensor.substrate.get_chain_head()
     new_balance, new_stake = await asyncio.gather(
@@ -677,6 +691,7 @@ async def _safe_unstake_extrinsic(
         f"Subnet: [{COLOR_PALETTE['GENERAL']['SUBHEADING']}]{netuid}[/{COLOR_PALETTE['GENERAL']['SUBHEADING']}] "
         f"Stake:\n  [blue]{current_stake}[/blue] :arrow_right: [{COLOR_PALETTE['STAKE']['STAKE_AMOUNT']}]{new_stake}"
     )
+    return True
 
 
 async def _unstake_all_extrinsic(