From a46c3bbbee305b8c706a5bfc13776201855714df Mon Sep 17 00:00:00 2001
From: The MMGen Project <mmgen@tuta.io>
Date: Sat, 26 Apr 2025 10:38:55 +0000
Subject: [PATCH] amt: new `TokenAmt` class

---
 mmgen/amt.py          | 14 ++++++++++++++
 test/modtest_d/amt.py | 18 +++++++++++++++++-
 2 files changed, 31 insertions(+), 1 deletion(-)

diff --git a/mmgen/amt.py b/mmgen/amt.py
index e02d4d8a..b7cf5682 100755
--- a/mmgen/amt.py
+++ b/mmgen/amt.py
@@ -208,6 +208,20 @@ class ETHAmt(CoinAmt):
 class ETCAmt(ETHAmt):
 	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):
 	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 810b316f..064f2ba6 100755
--- a/test/modtest_d/amt.py
+++ b/test/modtest_d/amt.py
@@ -9,6 +9,7 @@ from decimal import Decimal
 from mmgen.protocol import init_proto
 from mmgen.tx.new import parse_fee_spec
 from mmgen.cfg import Config
+from mmgen.amt import TokenAmt
 
 from ..include.common import cfg, vmsg
 
@@ -38,7 +39,7 @@ def test_fee_spec(data):
 
 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)'):
 		return test_to_unit((
@@ -79,3 +80,18 @@ class unit_tests:
 			('eth', '3701w', '3701', 'wei'),
 			('eth', '3.07M', '3.07', 'Mwei'),
 			('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