Merge pull request #676 from opentensor/feat/thewhaleking/wallet-balance-sorting

Adds wallet balance sorting

Author: thewhaleking

bittensor_cli/cli.py

@@ -39,6 +39,7 @@
     Constants,
     COLORS,
     HYPERPARAMS,
+    WalletOptions,
 )
 from bittensor_cli.src.bittensor import utils
 from bittensor_cli.src.bittensor.balances import Balance
@@ -92,6 +93,7 @@
     subnets,
     mechanisms as subnet_mechanisms,
 )
+from bittensor_cli.src.commands.wallets import SortByBalance
 from bittensor_cli.version import __version__, __version_as_int__
 
 try:
@@ -1915,7 +1917,7 @@ def wallet_ask(
         wallet_name: Optional[str],
         wallet_path: Optional[str],
         wallet_hotkey: Optional[str],
-        ask_for: Optional[list[Literal[WO.NAME, WO.PATH, WO.HOTKEY]]] = None,
+        ask_for: Optional[list[WalletOptions]] = None,
         validate: WV = WV.WALLET,
         return_wallet_and_hotkey: bool = False,
     ) -> Union[Wallet, tuple[Wallet, str]]:
@@ -3166,6 +3168,11 @@ def wallet_balance(
             "-a",
             help="Whether to display the balances for all the wallets.",
         ),
+        sort_by: Optional[wallets.SortByBalance] = typer.Option(
+            None,
+            "--sort",
+            help="When using `--all`, sorts the wallets by a given column",
+        ),
         network: Optional[list[str]] = Options.network,
         quiet: bool = Options.quiet,
         verbose: bool = Options.verbose,
@@ -3265,7 +3272,7 @@ def wallet_balance(
         subtensor = self.initialize_chain(network)
         return self._run_command(
             wallets.wallet_balance(
-                wallet, subtensor, all_balances, ss58_addresses, json_output
+                wallet, subtensor, all_balances, ss58_addresses, sort_by, json_output
             )
         )
 

bittensor_cli/src/commands/wallets.py

@@ -3,6 +3,7 @@
 import json
 import os
 from collections import defaultdict
+from enum import Enum
 from typing import Generator, Optional, Union
 
 import aiohttp
@@ -53,6 +54,27 @@
 )
 
 
+class SortByBalance(Enum):
+    name = "name"
+    free = "free"
+    staked = "staked"
+    total = "total"
+
+
+def _sort_by_balance_key(sort_by: SortByBalance):
+    """Get the sort key function based on the enum"""
+    if sort_by == SortByBalance.name:
+        return lambda row: row[0].lower()  # Case-insensitive alphabetical sort
+    elif sort_by == SortByBalance.free:
+        return lambda row: row[2]
+    elif sort_by == SortByBalance.staked:
+        return lambda row: row[3]
+    elif sort_by == SortByBalance.total:
+        return lambda row: row[4]
+    else:
+        raise ValueError("Invalid sort key")
+
+
 async def associate_hotkey(
     wallet: Wallet,
     subtensor: SubtensorInterface,
@@ -565,6 +587,7 @@ async def wallet_balance(
     subtensor: SubtensorInterface,
     all_balances: bool,
     ss58_addresses: Optional[str] = None,
+    sort_by: Optional[SortByBalance] = None,
     json_output: bool = False,
 ):
     """Retrieves the current balance of the specified wallet"""
@@ -644,14 +667,26 @@ async def wallet_balance(
         width=None,
         leading=True,
     )
-
-    for name, (coldkey, free, staked) in balances.items():
+    balance_rows = [
+        (name, coldkey, free, staked, free + staked)
+        for (name, (coldkey, free, staked)) in balances.items()
+    ]
+    sorted_balances = (
+        sorted(
+            balance_rows,
+            key=_sort_by_balance_key(sort_by),
+            reverse=(sort_by != SortByBalance.name),
+        )
+        if sort_by is not None
+        else balance_rows
+    )
+    for name, coldkey, free, staked, total in sorted_balances:
         table.add_row(
             name,
             coldkey,
             str(free),
             str(staked),
-            str(free + staked),
+            str(total),
         )
     table.add_row()
     table.add_row(