Improved the JSON output of stake add to include error messages.

Author: thewhaleking

bittensor_cli/cli.py

@@ -3316,13 +3316,13 @@ def stake_add(
                 "You have specified hotkeys to include and also the `--all-hotkeys` flag. The flag"
                 "should only be used standalone (to use all hotkeys) or with `--exclude-hotkeys`."
             )
-            raise typer.Exit()
+            return
 
         if include_hotkeys and exclude_hotkeys:
             print_error(
                 "You have specified options for both including and excluding hotkeys. Select one or the other."
             )
-            raise typer.Exit()
+            return
 
         if not wallet_hotkey and not all_hotkeys and not include_hotkeys:
             if not wallet_name:

bittensor_cli/src/commands/stake/add.py

@@ -72,7 +72,7 @@ async def safe_stake_extrinsic(
         hotkey_ss58_: str,
         price_limit: Balance,
         status=None,
-    ) -> bool:
+    ) -> tuple[bool, str]:
         err_out = partial(print_error, status=status)
         failure_prelude = (
             f":cross_mark: [red]Failed[/red] to stake {amount_} on Netuid {netuid_}"
@@ -104,25 +104,24 @@ async def safe_stake_extrinsic(
             )
         except SubstrateRequestException as e:
             if "Custom error: 8" in str(e):
-                print_error(
-                    f"\n{failure_prelude}: Price exceeded tolerance limit. "
+                err_msg = (
+                    f"{failure_prelude}: Price exceeded tolerance limit. "
                     f"Transaction rejected because partial staking is disabled. "
                     f"Either increase price tolerance or enable partial staking.",
-                    status=status,
                 )
-                return False
+                print_error("\n" + err_msg, status=status)
             else:
-                err_out(f"\n{failure_prelude} with error: {format_error_message(e)}")
-            return False
+                err_msg = f"{failure_prelude} with error: {format_error_message(e)}"
+                err_out("\n" + err_msg)
+            return False, err_msg
         if not await response.is_success:
-            err_out(
-                f"\n{failure_prelude} with error: {format_error_message(await response.error_message)}"
-            )
-            return False
+            err_msg = f"{failure_prelude} with error: {format_error_message(await response.error_message)}"
+            err_out("\n" + err_msg)
+            return False, err_msg
         else:
             if json_output:
                 # the rest of this checking is not necessary if using json_output
-                return True
+                return True, ""
             block_hash = await subtensor.substrate.get_chain_head()
             new_balance, new_stake = await asyncio.gather(
                 subtensor.get_balance(wallet.coldkeypub.ss58_address, block_hash),
@@ -160,11 +159,11 @@ async def safe_stake_extrinsic(
                 f":arrow_right: "
                 f"[{COLOR_PALETTE['STAKE']['STAKE_AMOUNT']}]{new_stake}\n"
             )
-            return True
+            return True, ""
 
     async def stake_extrinsic(
         netuid_i, amount_, current, staking_address_ss58, status=None
-    ) -> bool:
+    ) -> tuple[bool, str]:
         err_out = partial(print_error, status=status)
         current_balance, next_nonce, call = await asyncio.gather(
             subtensor.get_balance(wallet.coldkeypub.ss58_address),
@@ -190,18 +189,18 @@ async def stake_extrinsic(
                 extrinsic, wait_for_inclusion=True, wait_for_finalization=False
             )
         except SubstrateRequestException as e:
-            err_out(f"\n{failure_prelude} with error: {format_error_message(e)}")
-            return False
+            err_msg = f"{failure_prelude} with error: {format_error_message(e)}"
+            err_out("\n" + err_msg)
+            return False, err_msg
         else:
             if not await response.is_success:
-                err_out(
-                    f"\n{failure_prelude} with error: {format_error_message(await response.error_message)}"
-                )
-                return False
+                err_msg = f"{failure_prelude} with error: {format_error_message(await response.error_message)}"
+                err_out("\n" + err_msg)
+                return False, err_msg
             else:
                 if json_output:
                     # the rest of this is not necessary if using json_output
-                    return True
+                    return True, ""
                 new_block_hash = await subtensor.substrate.get_chain_head()
                 new_balance, new_stake = await asyncio.gather(
                     subtensor.get_balance(
@@ -230,7 +229,7 @@ async def stake_extrinsic(
                     f":arrow_right: "
                     f"[{COLOR_PALETTE['STAKE']['STAKE_AMOUNT']}]{new_stake}\n"
                 )
-                return True
+                return True, ""
 
     netuids = (
         netuids if netuids is not None else await subtensor.get_all_subnet_netuids()
@@ -417,13 +416,17 @@ async def stake_extrinsic(
             for _, staking_address in hotkeys_to_stake_to
         }
     successes = defaultdict(dict)
+    error_messages = defaultdict(dict)
     with console.status(f"\n:satellite: Staking on netuid(s): {netuids} ..."):
         # We can gather them all at once but balance reporting will be in race-condition.
         for (ni, staking_address), coroutine in stake_coroutines.items():
-            success = await coroutine
+            success, er_msg = await coroutine
             successes[ni][staking_address] = success
+            error_messages[ni][staking_address] = er_msg
     if json_output:
-        json_console.print(json.dumps({"staking_success": successes}))
+        json_console.print(
+            json.dumps({"staking_success": successes, "error_messages": error_messages})
+        )
 
 
 # Helper functions

tests/e2e_tests/test_staking_sudo.py

@@ -343,12 +343,22 @@ def test_staking(local_chain, wallet_setup):
             "--no-prompt",
             "--era",
             "144",
+            "--json-output",
         ],
     )
-    assert "✅ Finalized" in add_stake_multiple.stdout, add_stake_multiple.stderr
+    add_stake_multiple_output = json.loads(add_stake_multiple.stdout)
     for netuid_ in multiple_netuids:
-        assert f"Stake added to netuid: {netuid_}" in add_stake_multiple.stdout, (
-            add_stake_multiple.stderr
+        assert (
+            add_stake_multiple_output["staking_success"][str(netuid_)][
+                wallet_alice.hotkey.ss58_address
+            ]
+            is True
+        )
+        assert (
+            add_stake_multiple_output["error_messages"][str(netuid_)][
+                wallet_alice.hotkey.ss58_address
+            ]
+            == ""
         )
 
     # Fetch the hyperparameters of the subnet