From b6b1bc3c516629e00f0e59a646444293a330b654 Mon Sep 17 00:00:00 2001 From: The MMGen Project Date: Tue, 29 Sep 2026 09:23:31 +0000 Subject: [PATCH] CoinAmt: support arithmetic operations with token amounts --- mmgen/amt.py | 67 ++++++++++++++++++++++++++++++------------- test/modtest_d/amt.py | 8 ++++++ test/modtest_d/obj.py | 12 ++++---- 3 files changed, 60 insertions(+), 27 deletions(-) diff --git a/mmgen/amt.py b/mmgen/amt.py index 7c3e0bd5..1ce98ec3 100755 --- a/mmgen/amt.py +++ b/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 diff --git a/test/modtest_d/amt.py b/test/modtest_d/amt.py index 7bba230a..847e7f69 100755 --- a/test/modtest_d/amt.py +++ b/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 diff --git a/test/modtest_d/obj.py b/test/modtest_d/obj.py index 96036113..858dace2 100755 --- a/test/modtest_d/obj.py +++ b/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),