CoinAmt: support arithmetic operations with token amounts

mmgen/amt.py

@@ -40,7 +40,9 @@ class CoinAmt(Decimal, Hilite, InitErrors): # abstract class
 	max_amt  = None   # coin supply if known, otherwise None
 	units    = ()     # defined unit names, e.g. ('satoshi',...)
 
-	def __new__(cls, num, *, from_unit=None, from_decimal=False):
+	def __new__(cls, num, *, from_unit=None, from_decimal=False, decimals=None):
+
+		assert decimals == None
 
 		if isinstance(num, CoinAmt):
 			raise TypeError(f'CoinAmt: {num} is instance of {cls.__name__}')
@@ -129,39 +131,42 @@ class CoinAmt(Decimal, Hilite, InitErrors): # abstract class
 		raise TypeError(
 			f'operand {other} is of incorrect type ({type(other).__name__} != {type(self).__name__})')
 
-	def __add__(self, other):
+	def __add__(self, other, decimals=None):
 		"""
 		we must allow other to be int(0) to use the sum() builtin
 		"""
-		if type(other) is type(self) or (other == 0 and isinstance(other, int)):
-			return type(self)(Decimal.__add__(self, other),  from_decimal=True)
+		if type(other) in (type(self), Decimal) or (other == 0 and isinstance(other, int)):
+			return type(self)(Decimal.__add__(self, other), from_decimal=True, decimals=decimals)
 		self.incorrect_type_error(other)
 
 	__radd__ = __add__
 
-	def __sub__(self, other):
-		if type(other) is type(self):
-			return type(self)(Decimal.__sub__(Decimal(self), Decimal(other)),  from_decimal=True)
+	def __sub__(self, other, decimals=None):
+		if type(other) in (type(self), Decimal):
+			return type(self)(Decimal.__sub__(self, other), from_decimal=True, decimals=decimals)
 		self.incorrect_type_error(other)
 
-	def __rsub__(self, other):
+	def __rsub__(self, other, decimals=None):
 		if type(other) is type(self):
-			return type(self)(Decimal.__rsub__(self, other),  from_decimal=True)
+			return type(self)(Decimal.__rsub__(self, other),  from_decimal=True, decimals=decimals)
 		self.incorrect_type_error(other)
 
-	def __mul__(self, other):
-		return type(self)('{:0.{p}f}'.format(
-			Decimal.__mul__(self, Decimal(other)), p=self.max_prec))
+	def __mul__(self, other, decimals=None):
+		return type(self)(
+			'{:0.{p}f}'.format(Decimal.__mul__(self, Decimal(other)), p=self.max_prec),
+			decimals = decimals)
 
 	__rmul__ = __mul__
 
-	def __truediv__(self, other):
-		return type(self)('{:0.{p}f}'.format(
-			Decimal.__truediv__(self, Decimal(other)), p=self.max_prec))
+	def __truediv__(self, other, decimals=None):
+		return type(self)(
+			'{:0.{p}f}'.format(Decimal.__truediv__(self, Decimal(other)), p=self.max_prec),
+			decimals = decimals)
 
-	def __rtruediv__(self, other):
-		return type(self)('{:0.{p}f}'.format(
-			Decimal.__rtruediv__(self, Decimal(other)), p=self.max_prec))
+	def __rtruediv__(self, other, decimals=None):
+		return type(self)(
+			'{:0.{p}f}'.format(Decimal.__rtruediv__(self, Decimal(other)), p=self.max_prec),
+			decimals = decimals)
 
 	def __neg__(self):
 		self.method_not_implemented()
@@ -228,17 +233,39 @@ class ETCAmt(ETHAmt):
 class TokenAmt(ETHAmt):
 	units = ('atomic',)
 
-	def __new__(cls, num, *, decimals, from_unit=None):
+	def __new__(cls, num, *, decimals, from_unit=None, from_decimal=False):
 		assert isinstance(decimals, int)
 		cls.max_prec = decimals
 		cls.atomic = Decimal(f'{10 ** -decimals:0.{decimals}f}')
-		return ETHAmt.__new__(cls, num=num, from_unit=from_unit)
+		return ETHAmt.__new__(cls, num=num, from_unit=from_unit, from_decimal=from_decimal)
 
 	def to_unit(self, unit):
 		if (u := getattr(self, unit)) == self.atomic:
 			return int(Decimal(self) // u)
 		raise ValueError('TokenAmt unit must be ‘atomic’')
 
+	def __add__(self, other):
+		return super().__add__(other, decimals=self.max_prec)
+
+	__radd__ = __add__
+
+	def __sub__(self, other):
+		return super().__sub__(other, decimals=self.max_prec)
+
+	def __rsub__(self, other):
+		return super().__rsub__(other, decimals=self.max_prec)
+
+	def __mul__(self, other):
+		return super().__mul__(other, decimals=self.max_prec)
+
+	__rmul__ = __mul__
+
+	def __truediv__(self, other):
+		return super().__truediv__(other, decimals=self.max_prec)
+
+	def __rtruediv__(self, other):
+		return super().__rtruediv__(other, decimals=self.max_prec)
+
 def CoinAmtChk(proto, num):
 	assert type(num) is proto.coin_amt, f'CoinAmtChk: {type(num)} != {proto.coin_amt}'
 	return num

test/modtest_d/amt.py

@@ -140,4 +140,12 @@ class unit_tests:
 			if unit == 'atomic':
 				assert amt.to_unit('atomic') == n
 			assert amt_disp == chk, f'{amt_disp} != {chk}'
+			assert amt + amt == amt * 2
+			assert amt - amt == 0
+			assert amt + 0 == amt
+			assert amt + Decimal(0) == amt
+			assert amt + Decimal('2.123') == Decimal(amt + Decimal('2.123'))
+			assert amt - Decimal(0) == Decimal(amt)
+			assert amt * 7 == Decimal(amt) * 7
+			assert amt / 7 == round(Decimal(amt) / 7, dec)
 		return True

test/modtest_d/obj.py

@@ -47,15 +47,16 @@ def coinamt_test(cls, aa, bb, ut):
 	do('a * b / a', a * b / a, A * B / A)
 	do('a * b / b', a * b / b, A * B / B)
 
+	do('a + B', a + B, B + a)
+	do('b - A', b - A, b - a)
+
 	vmsg(f'\nChecking {cls.__name__} error handling...')
 
 	bad_data = (
 		('negation',          'NotImplementedError', 'not implemented',    lambda: -a),
 		('modulus',           'NotImplementedError', 'not implemented',    lambda: b % a),
 		('floor division',    'NotImplementedError', 'not implemented',    lambda: b // a),
-		('negative result',   'ObjectInitError',     'cannot be negative', lambda: a - b),
-		('operand type',      'TypeError',           'incorrect type',     lambda: a + B),
-		('operand type',      'TypeError',           'incorrect type',     lambda: b - A))
+		('negative result',   'ObjectInitError',     'cannot be negative', lambda: a - b))
 
 	if cls.max_amt is not None:
 		bad_data += (
@@ -131,16 +132,15 @@ class unit_tests:
 		b = 2 / a
 		test_equal(type(b), coin_amt)
 		test_equal(b, coin_amt('1.62074554'))
+		test_equal(a + Decimal(1), Decimal(1) + a)
 
 		def bad1(): b = a + 1
 		def bad2(): b = a - 1
-		def bad3(): a + Decimal(1)
 		def bad4(): b = a + 0.0
 		def bad5(): b = a - 0.0
 
 		def bad1r(): b = 1 + a
 		def bad2r(): b = 3 - a
-		def bad3r(): Decimal(1) + a
 		def bad4r(): b = 0.0 + a
 		def bad5r(): b = 0.0 - a
 
@@ -159,13 +159,11 @@ class unit_tests:
 			(
 				('addition with int',      'TypeError',       'incorrect type',     bad1),
 				('subtraction with int',   'TypeError',       'incorrect type',     bad2),
-				('addition with Decimal',  'TypeError',       'incorrect type',     bad3),
 				('addition with float',    'TypeError',       'incorrect type',     bad4),
 				('subtraction with float', 'TypeError',       'incorrect type',     bad5),
 
 				('addition with int',      'TypeError',       'incorrect type',     bad1r),
 				('subtraction with int',   'TypeError',       'incorrect type',     bad2r),
-				('addition with Decimal',  'TypeError',       'incorrect type',     bad3r),
 				('addition with float',    'TypeError',       'incorrect type',     bad4r),
 				('subtraction with float', 'TypeError',       'incorrect type',     bad5r),