Merge pull request #2764 from opentensor/fix/roman/balance-unit

Improve logic in Balance magic methods

Author: basfroman

bittensor/utils/balance.py

@@ -6,6 +6,35 @@
 from bittensor.core import settings
 
 
+def _check_currencies(self, other):
+    """Checks that Balance objects have the same netuids to perform arithmetic operations.
+
+    A warning is raised if the netuids differ.
+
+    Example:
+        >>> balance1 = Balance.from_rao(1000).set_unit(12)
+        >>> balance2 = Balance.from_rao(500).set_unit(12)
+        >>> balance1 + balance2  # No warning
+
+        >>> balance3 = Balance.from_rao(200).set_unit(15)
+        >>> balance1 + balance3  # Raises DeprecationWarning
+
+    In this example:
+        - `from_rao` creates a Balance instance from the amount in rao (smallest unit).
+        - `set_unit(12)` sets the unit to correspond to subnet 12 (i.e., Alpha from netuid 12).
+    """
+    if self.netuid != other.netuid:
+        warnings.simplefilter("default", DeprecationWarning)
+        warnings.warn(
+            "Balance objects must have the same netuid (Alpha currency) to perform arithmetic operations.\n"
+            f"First balance is `{self}`.  Second balance is `{other}`.\n\n"
+            "To create a Balance instance with the correct netuid, use:\n"
+            "Balance.from_rao(1000).set_unit(12)  # 1000 rao in subnet 12",
+            category=DeprecationWarning,
+            stacklevel=2,
+        )
+
+
 class Balance:
     """
     Represents the bittensor balance of the wallet, stored as rao (int).
@@ -23,6 +52,7 @@ class Balance:
     rao_unit: str = settings.RAO_SYMBOL  # This is the rao unit
     rao: int
     tao: float
+    netuid: int = 0
 
     def __init__(self, balance: Union[int, float]):
         """
@@ -78,7 +108,8 @@ def __eq__(self, other: Union[int, float, "Balance"]):
         if other is None:
             return False
 
-        if hasattr(other, "rao"):
+        if isinstance(other, Balance):
+            _check_currencies(self, other)
             return self.rao == other.rao
         else:
             try:
@@ -92,7 +123,8 @@ def __ne__(self, other: Union[int, float, "Balance"]):
         return not self == other
 
     def __gt__(self, other: Union[int, float, "Balance"]):
-        if hasattr(other, "rao"):
+        if isinstance(other, Balance):
+            _check_currencies(self, other)
             return self.rao > other.rao
         else:
             try:
@@ -103,7 +135,8 @@ def __gt__(self, other: Union[int, float, "Balance"]):
                 raise NotImplementedError("Unsupported type")
 
     def __lt__(self, other: Union[int, float, "Balance"]):
-        if hasattr(other, "rao"):
+        if isinstance(other, Balance):
+            _check_currencies(self, other)
             return self.rao < other.rao
         else:
             try:
@@ -115,111 +148,129 @@ def __lt__(self, other: Union[int, float, "Balance"]):
 
     def __le__(self, other: Union[int, float, "Balance"]):
         try:
+            if isinstance(other, Balance):
+                _check_currencies(self, other)
             return self < other or self == other
         except TypeError:
             raise NotImplementedError("Unsupported type")
 
     def __ge__(self, other: Union[int, float, "Balance"]):
         try:
+            if isinstance(other, Balance):
+                _check_currencies(self, other)
             return self > other or self == other
         except TypeError:
             raise NotImplementedError("Unsupported type")
 
     def __add__(self, other: Union[int, float, "Balance"]):
-        if hasattr(other, "rao"):
-            return Balance.from_rao(int(self.rao + other.rao))
+        if isinstance(other, Balance):
+            _check_currencies(self, other)
+            return Balance.from_rao(int(self.rao + other.rao)).set_unit(self.netuid)
         else:
             try:
                 # Attempt to cast to int from rao
-                return Balance.from_rao(int(self.rao + other))
+                return Balance.from_rao(int(self.rao + other)).set_unit(self.netuid)
             except (ValueError, TypeError):
                 raise NotImplementedError("Unsupported type")
 
     def __radd__(self, other: Union[int, float, "Balance"]):
         try:
+            if isinstance(other, Balance):
+                _check_currencies(self, other)
             return self + other
         except TypeError:
             raise NotImplementedError("Unsupported type")
 
     def __sub__(self, other: Union[int, float, "Balance"]):
         try:
+            if isinstance(other, Balance):
+                _check_currencies(self, other)
             return self + -other
         except TypeError:
             raise NotImplementedError("Unsupported type")
 
     def __rsub__(self, other: Union[int, float, "Balance"]):
         try:
+            if isinstance(other, Balance):
+                _check_currencies(self, other)
             return -self + other
         except TypeError:
             raise NotImplementedError("Unsupported type")
 
     def __mul__(self, other: Union[int, float, "Balance"]):
-        if hasattr(other, "rao"):
-            return Balance.from_rao(int(self.rao * other.rao))
+        if isinstance(other, Balance):
+            _check_currencies(self, other)
+            return Balance.from_rao(int(self.rao * other.rao)).set_unit(self.netuid)
         else:
             try:
                 # Attempt to cast to int from rao
-                return Balance.from_rao(int(self.rao * other))
+                return Balance.from_rao(int(self.rao * other)).set_unit(self.netuid)
             except (ValueError, TypeError):
                 raise NotImplementedError("Unsupported type")
 
     def __rmul__(self, other: Union[int, float, "Balance"]):
+        if isinstance(other, Balance):
+            _check_currencies(self, other)
         return self * other
 
     def __truediv__(self, other: Union[int, float, "Balance"]):
-        if hasattr(other, "rao"):
-            return Balance.from_rao(int(self.rao / other.rao))
+        if isinstance(other, Balance):
+            _check_currencies(self, other)
+            return Balance.from_rao(int(self.rao / other.rao)).set_unit(self.netuid)
         else:
             try:
                 # Attempt to cast to int from rao
-                return Balance.from_rao(int(self.rao / other))
+                return Balance.from_rao(int(self.rao / other)).set_unit(self.netuid)
             except (ValueError, TypeError):
                 raise NotImplementedError("Unsupported type")
 
     def __rtruediv__(self, other: Union[int, float, "Balance"]):
-        if hasattr(other, "rao"):
-            return Balance.from_rao(int(other.rao / self.rao))
+        if isinstance(other, Balance):
+            _check_currencies(self, other)
+            return Balance.from_rao(int(other.rao / self.rao)).set_unit(self.netuid)
         else:
             try:
                 # Attempt to cast to int from rao
-                return Balance.from_rao(int(other / self.rao))
+                return Balance.from_rao(int(other / self.rao)).set_unit(self.netuid)
             except (ValueError, TypeError):
                 raise NotImplementedError("Unsupported type")
 
     def __floordiv__(self, other: Union[int, float, "Balance"]):
-        if hasattr(other, "rao"):
-            return Balance.from_rao(int(self.tao // other.tao))
+        if isinstance(other, Balance):
+            _check_currencies(self, other)
+            return Balance.from_rao(int(self.tao // other.tao)).set_unit(self.netuid)
         else:
             try:
                 # Attempt to cast to int from rao
-                return Balance.from_rao(int(self.rao // other))
+                return Balance.from_rao(int(self.rao // other)).set_unit(self.netuid)
             except (ValueError, TypeError):
                 raise NotImplementedError("Unsupported type")
 
     def __rfloordiv__(self, other: Union[int, float, "Balance"]):
-        if hasattr(other, "rao"):
-            return Balance.from_rao(int(other.rao // self.rao))
+        if isinstance(other, Balance):
+            _check_currencies(self, other)
+            return Balance.from_rao(int(other.rao // self.rao)).set_unit(self.netuid)
         else:
             try:
                 # Attempt to cast to int from rao
-                return Balance.from_rao(int(other // self.rao))
+                return Balance.from_rao(int(other // self.rao)).set_unit(self.netuid)
             except (ValueError, TypeError):
                 raise NotImplementedError("Unsupported type")
 
     def __nonzero__(self) -> bool:
         return bool(self.rao)
 
     def __neg__(self):
-        return Balance.from_rao(-self.rao)
+        return Balance.from_rao(-self.rao).set_unit(self.netuid)
 
     def __pos__(self):
-        return Balance.from_rao(self.rao)
+        return Balance.from_rao(self.rao).set_unit(self.netuid)
 
     def __abs__(self):
-        return Balance.from_rao(abs(self.rao))
+        return Balance.from_rao(abs(self.rao)).set_unit(self.netuid)
 
     @staticmethod
-    def from_float(amount: float, netuid: int = 0):
+    def from_float(amount: float, netuid: int = 0) -> "Balance":
         """
         Given tao, return :func:`Balance` object with rao(``int``) and tao(``float``), where rao = int(tao*pow(10,9))
         Args:
@@ -233,7 +284,7 @@ def from_float(amount: float, netuid: int = 0):
         return Balance(rao_).set_unit(netuid)
 
     @staticmethod
-    def from_tao(amount: float, netuid: int = 0):
+    def from_tao(amount: float, netuid: int = 0) -> "Balance":
         """
         Given tao, return Balance object with rao(``int``) and tao(``float``), where rao = int(tao*pow(10,9))
 
@@ -248,7 +299,7 @@ def from_tao(amount: float, netuid: int = 0):
         return Balance(rao_).set_unit(netuid)
 
     @staticmethod
-    def from_rao(amount: int, netuid: int = 0):
+    def from_rao(amount: int, netuid: int = 0) -> "Balance":
         """
         Given rao, return Balance object with rao(``int``) and tao(``float``), where rao = int(tao*pow(10,9))
 
@@ -262,7 +313,7 @@ def from_rao(amount: int, netuid: int = 0):
         return Balance(amount).set_unit(netuid)
 
     @staticmethod
-    def get_unit(netuid: int):
+    def get_unit(netuid: int) -> str:
         base = len(units)
         if netuid < base:
             return units[netuid]
@@ -274,6 +325,7 @@ def get_unit(netuid: int):
             return result
 
     def set_unit(self, netuid: int):
+        self.netuid = netuid
         self.unit = Balance.get_unit(netuid)
         self.rao_unit = Balance.get_unit(netuid)
         return self
@@ -777,18 +829,18 @@ def fixed_to_float(
 ]
 
 
-def tao(amount: float) -> Balance:
+def tao(amount: float, netuid: int = 0) -> Balance:
     """
     Helper function to create a Balance object from a float (Tao)
     """
-    return Balance.from_tao(amount)
+    return Balance.from_tao(amount).set_unit(netuid)
 
 
-def rao(amount: int) -> Balance:
+def rao(amount: int, netuid: int = 0) -> Balance:
     """
     Helper function to create a Balance object from an int (Rao)
     """
-    return Balance.from_rao(amount)
+    return Balance.from_rao(amount).set_unit(netuid)
 
 
 def check_and_convert_to_balance(