diff --git a/bittensor/cli.py b/bittensor/cli.py index e4513f5ec0..a7ee579624 100644 --- a/bittensor/cli.py +++ b/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, diff --git a/bittensor/commands/__init__.py b/bittensor/commands/__init__.py index 86e41f91a8..03822f9bf5 100644 --- a/bittensor/commands/__init__.py +++ b/bittensor/commands/__init__.py @@ -81,6 +81,7 @@ RegenHotkeyCommand, UpdateWalletCommand, WalletCreateCommand, + WalletBalanceCommand, ) from .transfer import TransferCommand from .inspect import InspectCommand diff --git a/bittensor/commands/root.py b/bittensor/commands/root.py index f39f95d5c3..78c31fa3d7 100644 --- a/bittensor/commands/root.py +++ b/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) diff --git a/bittensor/commands/wallets.py b/bittensor/commands/wallets.py index 5417aa5496..9e892b5524 100644 --- a/bittensor/commands/wallets.py +++ b/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))