Merge pull request opentensor#701 from opentensor/update/btcli-stake-claim-args

Update/btcli stake claim args

Author: ibraheem-abe

bittensor_cli/cli.py

@@ -7085,6 +7085,10 @@ def view_dashboard(
 
     def stake_set_claim_type(
         self,
+        claim_type: Optional[str] = typer.Argument(
+            None,
+            help="Claim type: 'keep' or 'swap'. If not provided, you'll be prompted to choose.",
+        ),
         wallet_name: Optional[str] = Options.wallet_name,
         wallet_path: Optional[str] = Options.wallet_path,
         wallet_hotkey: Optional[str] = Options.wallet_hotkey,
@@ -7105,13 +7109,26 @@ def stake_set_claim_type(
 
         USAGE:
 
-        [green]$[/green] btcli root set-claim-type
+        [green]$[/green] btcli stake claim
+        [green]$[/green] btcli stake claim keep
+        [green]$[/green] btcli stake claim swap
 
         With specific wallet:
 
-        [green]$[/green] btcli root set-claim-type --wallet-name my_wallet
+        [green]$[/green] btcli stake claim swap --wallet-name my_wallet
         """
         self.verbosity_handler(quiet, verbose, json_output)
+
+        if claim_type is not None:
+            claim_type_normalized = claim_type.capitalize()
+            if claim_type_normalized not in ["Keep", "Swap"]:
+                err_console.print(
+                    f":cross_mark: [red]Invalid claim type '{claim_type}'. Must be 'keep' or 'swap'.[/red]"
+                )
+                raise typer.Exit()
+        else:
+            claim_type_normalized = None
+
         wallet = self.wallet_ask(
             wallet_name,
             wallet_path,
@@ -7122,6 +7139,7 @@ def stake_set_claim_type(
             claim_stake.set_claim_type(
                 wallet=wallet,
                 subtensor=self.initialize_chain(network),
+                claim_type=claim_type_normalized,
                 prompt=prompt,
                 json_output=json_output,
             )

bittensor_cli/src/commands/stake/claim.py

@@ -25,6 +25,7 @@
 async def set_claim_type(
     wallet: Wallet,
     subtensor: "SubtensorInterface",
+    claim_type: Optional[str] = None,
     prompt: bool = True,
     json_output: bool = False,
 ) -> tuple[bool, str, Optional[str]]:
@@ -38,6 +39,7 @@ async def set_claim_type(
     Args:
         wallet: Bittensor wallet object
         subtensor: SubtensorInterface object
+        claim_type: Optional claim type ("Keep" or "Swap"). If None, user will be prompted.
         prompt: Whether to prompt for user confirmation
         json_output: Whether to output JSON
 
@@ -76,8 +78,13 @@ async def set_claim_type(
         wallet.coldkeypub.ss58_address, f"[yellow]{current_type}[/yellow]"
     )
     console.print(claim_table)
-    new_type = Prompt.ask(
-        "Select new root claim type", choices=["Swap", "Keep"], default=current_type
+
+    new_type = (
+        claim_type
+        if claim_type
+        else Prompt.ask(
+            "Select new root claim type", choices=["Swap", "Keep"], default=current_type
+        )
     )
     if new_type == current_type:
         msg = f"Root claim type is already set to '{current_type}'. No change needed."