Adds support for ss58 addresses in wallet balance

bittensor_cli/cli.py

@@ -129,7 +129,9 @@ class Options:
         flag_value=False,
     )
     public_hex_key = typer.Option(None, help="The public key in hex format.")
-    ss58_address = typer.Option(None, help="The SS58 address of the coldkey.")
+    ss58_address = typer.Option(
+        None, "--ss58", "--ss58-address", help="The SS58 address of the coldkey."
+    )
     overwrite_coldkey = typer.Option(
         False,
         help="Overwrite the old coldkey with the newly generated coldkey.",
@@ -2026,6 +2028,7 @@ def wallet_balance(
         wallet_name: Optional[str] = Options.wallet_name,
         wallet_path: Optional[str] = Options.wallet_path,
         wallet_hotkey: Optional[str] = Options.wallet_hotkey,
+        ss58_addresses: Optional[list[str]] = Options.ss58_address,
         all_balances: Optional[bool] = typer.Option(
             False,
             "--all",
@@ -2039,6 +2042,8 @@ def wallet_balance(
         """
         Check the balance of the wallet. This command shows a detailed view of the wallet's coldkey balances, including free and staked balances.
 
+        You can also pass multiple ss58 addresses of coldkeys to check their balance (using --ss58).
+
         EXAMPLES:
 
         - To display the balance of a single wallet, use the command with the `--wallet-name` argument and provide the wallet name:
@@ -2052,17 +2057,41 @@ def wallet_balance(
         - To display the balances of all your wallets, use the `--all` argument:
 
             [green]$[/green] btcli w balance --all
+
+        - To display the balances of ss58 addresses, use the `--ss58` argument:
+
+            [green]$[/green] btcli w balance --ss58 <ss58_address> --ss58 <ss58_address>
+
         """
         self.verbosity_handler(quiet, verbose)
 
-        ask_for = [WO.PATH] if all_balances else [WO.NAME, WO.PATH]
-        validate = WV.NONE if all_balances else WV.WALLET
-        wallet = self.wallet_ask(
-            wallet_name, wallet_path, wallet_hotkey, ask_for=ask_for, validate=validate
-        )
+        if ss58_addresses:
+            valid_ss58s = [
+                ss58 for ss58 in set(ss58_addresses) if is_valid_ss58_address(ss58)
+            ]
+
+            invalid_ss58s = set(ss58_addresses) - set(valid_ss58s)
+            for invalid_ss58 in invalid_ss58s:
+                print_error(f"Incorrect ss58 address: {invalid_ss58}. Skipping.")
+
+            if valid_ss58s:
+                wallet = None
+                ss58_addresses = valid_ss58s
+            else:
+                raise typer.Exit()
+        else:
+            ask_for = [WO.PATH] if all_balances else [WO.NAME, WO.PATH]
+            validate = WV.NONE if all_balances else WV.WALLET
+            wallet = self.wallet_ask(
+                wallet_name,
+                wallet_path,
+                wallet_hotkey,
+                ask_for=ask_for,
+                validate=validate,
+            )
         subtensor = self.initialize_chain(network)
         return self._run_command(
-            wallets.wallet_balance(wallet, subtensor, all_balances)
+            wallets.wallet_balance(wallet, subtensor, all_balances, ss58_addresses)
         )
 
     def wallet_history(

bittensor_cli/src/commands/wallets.py

@@ -221,16 +221,25 @@ def _get_coldkey_ss58_addresses_for_path(path: str) -> tuple[list[str], list[str
 
 
 async def wallet_balance(
-    wallet: Wallet, subtensor: SubtensorInterface, all_balances: bool
+    wallet: Optional[Wallet],
+    subtensor: SubtensorInterface,
+    all_balances: bool,
+    ss58_addresses: Optional[str] = None,
 ):
     """Retrieves the current balance of the specified wallet"""
-    if not all_balances:
+    if ss58_addresses:
+        coldkeys = ss58_addresses
+        wallet_names = [f"Provided Address {i + 1}" for i in range(len(ss58_addresses))]
+
+    elif not all_balances:
         if not wallet.coldkeypub_file.exists_on_device():
             err_console.print("[bold red]No wallets found.[/bold red]")
             return
 
     with console.status("Retrieving balances", spinner="aesthetic") as status:
-        if all_balances:
+        if ss58_addresses:
+            print_verbose(f"Fetching data for ss58 address: {ss58_addresses}", status)
+        elif all_balances:
             print_verbose("Fetching data for all wallets", status)
             coldkeys, wallet_names = _get_coldkey_ss58_addresses_for_path(wallet.path)
         else: