Add btcli wallet balance

bittensor/cli.py

@@ -85,7 +85,7 @@
             "overview": OverviewCommand,
             "transfer": TransferCommand,
             "inspect": InspectCommand,
-            # "balance": None,
+            "balance": WalletBalanceCommand,
             "create": WalletCreateCommand,
             "new_hotkey": NewHotkeyCommand,
             "new_coldkey": NewColdkeyCommand,

bittensor/commands/__init__.py

@@ -81,6 +81,7 @@
     RegenHotkeyCommand,
     UpdateWalletCommand,
     WalletCreateCommand,
+    WalletBalanceCommand,
 )
 from .transfer import TransferCommand
 from .inspect import InspectCommand

bittensor/commands/root.py

@@ -235,9 +235,14 @@ def run(cli):
         for matrix in weights:
             [uid, weights_data] = matrix
 
-            normalized_weights = np.array(weights_data)[:, 1] / max(
-                np.sum(weights_data, axis=0)[1], 1
-            )
+            if not len(weights_data):
+                uid_to_weights[uid] = {}
+                normalized_weights = []
+            else:
+                normalized_weights = np.array(weights_data)[:, 1] / max(
+                    np.sum(weights_data, axis=0)[1], 1
+                )
+
             for weight_data, normalized_weight in zip(weights_data, normalized_weights):
                 [netuid, _] = weight_data
                 netuids.add(netuid)

bittensor/commands/wallets.py

@@ -20,6 +20,7 @@
 import os
 import sys
 from rich.prompt import Prompt, Confirm
+from rich.table import Table
 from typing import Optional, List
 from . import defaults
 
@@ -527,3 +528,129 @@ def check_config(config: "bittensor.Config"):
                 "Enter wallet name", default=bittensor.defaults.wallet.name
             )
             config.wallet.name = str(wallet_name)
+
+
+def _get_coldkey_ss58_addresses_for_path(path: str) -> List[str]:
+    """Get all coldkey ss58 addresses from path."""
+
+    def list_coldkeypub_files(dir_path):
+        abspath = os.path.abspath(os.path.expanduser(dir_path))
+        coldkey_files = []
+
+        for file in os.listdir(abspath):
+            coldkey_path = os.path.join(abspath, file, "coldkeypub.txt")
+            if os.path.exists(coldkey_path):
+                coldkey_files.append(coldkey_path)
+            else:
+                bittensor.logging.warning(
+                    f"{coldkey_path} does not exist. Excluding..."
+                )
+        return coldkey_files
+
+    return [
+        bittensor.keyfile(file).keypair.ss58_address
+        for file in list_coldkeypub_files(path)
+    ]
+
+
+class WalletBalanceCommand:
+    @staticmethod
+    def run(cli):
+        """Check the balance of the wallet."""
+        wallet_names = os.listdir(os.path.expanduser(cli.config.wallet.path))
+        coldkeys = _get_coldkey_ss58_addresses_for_path(cli.config.wallet.path)
+        subtensor = bittensor.subtensor(config=cli.config)
+
+        free_balances = [
+            subtensor.get_balance(coldkeys[i]) for i in range(len(coldkeys))
+        ]
+        staked_balances = [
+            subtensor.get_total_stake_for_coldkey(coldkeys[i])
+            for i in range(len(coldkeys))
+        ]
+        total_free_balance = sum(free_balances)
+        total_staked_balance = sum(staked_balances)
+
+        balances = {
+            name: (coldkey, free, staked)
+            for name, coldkey, free, staked in sorted(
+                zip(wallet_names, coldkeys, free_balances, staked_balances)
+            )
+        }
+
+        table = Table(show_footer=False)
+        table.title = "[white]Wallet Coldkey Balances"
+        table.add_column(
+            "[white]Wallet Name",
+            header_style="overline white",
+            footer_style="overline white",
+            style="rgb(50,163,219)",
+            no_wrap=True,
+        )
+
+        table.add_column(
+            "[white]Coldkey Address",
+            header_style="overline white",
+            footer_style="overline white",
+            style="rgb(50,163,219)",
+            no_wrap=True,
+        )
+
+        for typestr in ["Free", "Staked", "Total"]:
+            table.add_column(
+                f"[white]{typestr} Balance",
+                header_style="overline white",
+                footer_style="overline white",
+                justify="right",
+                style="green",
+                no_wrap=True,
+            )
+
+        for name, (coldkey, free, staked) in balances.items():
+            table.add_row(
+                name,
+                coldkey,
+                str(free),
+                str(staked),
+                str(free + staked),
+            )
+        table.add_row()
+        table.add_row(
+            "Total Balance Across All Coldkeys",
+            "",
+            str(total_free_balance),
+            str(total_staked_balance),
+            str(total_free_balance + total_staked_balance),
+        )
+        table.show_footer = True
+
+        table.box = None
+        table.pad_edge = False
+        table.width = None
+        bittensor.__console__.print(table)
+
+    @staticmethod
+    def add_args(parser: argparse.ArgumentParser):
+        balance_parser = parser.add_parser(
+            "balance", help="""Checks the balance of the wallet."""
+        )
+        bittensor.wallet.add_args(balance_parser)
+        bittensor.subtensor.add_args(balance_parser)
+
+    @staticmethod
+    def check_config(config: "bittensor.config"):
+        if not config.is_set("wallet.path") and not config.no_prompt:
+            path = Prompt.ask("Enter wallets path", default=defaults.wallet.path)
+            config.wallet.path = str(path)
+
+        if not config.is_set("subtensor.network") and not config.no_prompt:
+            network = Prompt.ask(
+                "Enter network",
+                default=defaults.subtensor.network,
+                choices=bittensor.__networks__,
+            )
+            config.subtensor.network = str(network)
+            (
+                _,
+                config.subtensor.chain_endpoint,
+            ) = bittensor.subtensor.determine_chain_endpoint_and_network(str(network))