Balance update.

Author: thewhaleking

bittensor_cli/src/bittensor/balances.py

@@ -18,6 +18,7 @@
 # DEALINGS IN THE SOFTWARE.
 
 from typing import Union
+from bittensor_cli.src import UNITS
 
 
 class Balance:
@@ -279,3 +280,21 @@ def from_rao(amount: int):
         :return: A Balance object representing the given amount.
         """
         return Balance(amount)
+
+    @staticmethod
+    def get_unit(netuid: int):
+        units = UNITS
+        base = len(units)
+        if netuid < base:
+            return units[netuid]
+        else:
+            result = ""
+            while netuid > 0:
+                result = units[netuid % base] + result
+                netuid //= base
+            return result
+
+    def set_unit(self, netuid: int):
+        self.unit = Balance.get_unit(netuid)
+        self.rao_unit = Balance.get_unit(netuid)
+        return self