Browse Source

amt: new `TokenAmt` class

The MMGen Project 7 months ago
parent
commit
a46c3bbbee
2 changed files with 31 additions and 1 deletions
  1. 14 0
      mmgen/amt.py
  2. 17 1
      test/modtest_d/amt.py

+ 14 - 0
mmgen/amt.py

@@ -208,6 +208,20 @@ class ETHAmt(CoinAmt):
 class ETCAmt(ETHAmt):
 class ETCAmt(ETHAmt):
 	coin = 'ETC'
 	coin = 'ETC'
 
 
+class TokenAmt(ETHAmt):
+	units = ('atomic',)
+
+	def __new__(cls, num, *, decimals, from_unit=None):
+		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)
+
+	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 CoinAmtChk(proto, num):
 def CoinAmtChk(proto, num):
 	assert type(num) is proto.coin_amt, f'CoinAmtChk: {type(num)} != {proto.coin_amt}'
 	assert type(num) is proto.coin_amt, f'CoinAmtChk: {type(num)} != {proto.coin_amt}'
 	return num
 	return num

+ 17 - 1
test/modtest_d/amt.py

@@ -9,6 +9,7 @@ from decimal import Decimal
 from mmgen.protocol import init_proto
 from mmgen.protocol import init_proto
 from mmgen.tx.new import parse_fee_spec
 from mmgen.tx.new import parse_fee_spec
 from mmgen.cfg import Config
 from mmgen.cfg import Config
+from mmgen.amt import TokenAmt
 
 
 from ..include.common import cfg, vmsg
 from ..include.common import cfg, vmsg
 
 
@@ -38,7 +39,7 @@ def test_fee_spec(data):
 
 
 class unit_tests:
 class unit_tests:
 
 
-	altcoin_deps = ('fee_spec_alt', 'to_unit_alt')
+	altcoin_deps = ('fee_spec_alt', 'to_unit_alt', 'token_amt')
 
 
 	def to_unit(self, name, ut, desc='CoinAmt.to_unit() (BTC)'):
 	def to_unit(self, name, ut, desc='CoinAmt.to_unit() (BTC)'):
 		return test_to_unit((
 		return test_to_unit((
@@ -79,3 +80,18 @@ class unit_tests:
 			('eth', '3701w', '3701', 'wei'),
 			('eth', '3701w', '3701', 'wei'),
 			('eth', '3.07M', '3.07', 'Mwei'),
 			('eth', '3.07M', '3.07', 'Mwei'),
 			('xmr', '3.07a', '3.07', 'atomic')))
 			('xmr', '3.07a', '3.07', 'atomic')))
+
+	def token_amt(self, name, ut, desc='TokenAmt (ETH)'):
+		for n, dec, unit, chk in (
+			(1234567,                    6,  'atomic', '1.234567'),
+			('1.234567',                 6,  None,     '1.234567'),
+			(1234567,                    22, 'atomic', '0.0000000000000001234567'),
+			('0.0000000000000001234567', 22, None,     '0.0000000000000001234567'),
+		):
+			amt = TokenAmt(n, decimals=dec, from_unit=unit)
+			amt_disp = amt.hl(color=False)
+			vmsg('  ' + amt_disp)
+			if unit == 'atomic':
+				assert amt.to_unit('atomic') == n
+			assert amt_disp == chk, f'{amt_disp} != {chk}'
+		return True