mmgen-tool getbalance: reimplement

mmgen/proto/btc/tw/bal.py

@@ -17,34 +17,37 @@ from ....tw.shared import get_tw_label
 
 class BitcoinTwGetBalance(TwGetBalance):
 
-	fs = '{w:13} {u:<16} {p:<16} {c}'
+	start_labels = ('TOTAL','Non-MMGen','Non-wallet')
+	conf_cols = {
+		'unconfirmed': 'Unconfirmed',
+		'lt_minconf':  '<{minconf} confs',
+		'ge_minconf':  '>={minconf} confs',
+	}
 
 	async def create_data(self):
-		# 0: unconfirmed, 1: below minconf, 2: confirmed, 3: spendable (privkey in wallet)
 		lbl_id = ('account','label')['label_api' in self.rpc.caps]
 		amt0 = self.proto.coin_amt('0')
 		for d in await self.rpc.call('listunspent',0):
-			lbl = get_tw_label(self.proto,d[lbl_id])
-			if lbl:
-				if lbl.mmid.type == 'mmgen':
-					key = lbl.mmid.obj.sid
-					if key not in self.data:
-						self.data[key] = [amt0] * 4
+			tw_lbl = get_tw_label(self.proto,d[lbl_id])
+			if tw_lbl:
+				if tw_lbl.mmid.type == 'mmgen':
+					label = tw_lbl.mmid.obj.sid
+					if label not in self.data:
+						self.data[label] = self.balance_info()
 				else:
-					key = 'Non-MMGen'
+					label = 'Non-MMGen'
 			else:
-				lbl,key = None,'Non-wallet'
+				label = 'Non-wallet'
 
 			amt = self.proto.coin_amt(d['amount'])
 
 			if not d['confirmations']:
-				self.data['TOTAL'][0] += amt
-				self.data[key][0] += amt
+				self.data['TOTAL']['unconfirmed'] += amt
+				self.data[label]['unconfirmed'] += amt
 
-			conf_level = (1,2)[d['confirmations'] >= self.minconf]
-
-			self.data['TOTAL'][conf_level] += amt
-			self.data[key][conf_level] += amt
+			col_key = ('lt_minconf','ge_minconf')[d['confirmations'] >= self.minconf]
+			self.data['TOTAL'][col_key] += amt
+			self.data[label][col_key] += amt
 
 			if d['spendable']:
-				self.data[key][3] += amt
+				self.data[label]['spendable'] += amt

mmgen/proto/eth/tw/bal.py

@@ -25,29 +25,31 @@ from ....tw.bal import TwGetBalance
 
 class EthereumTwGetBalance(TwGetBalance):
 
-	fs = '{w:13} {c}\n' # TODO - for now, just suppress display of meaningless data
+	start_labels = ('TOTAL','Non-MMGen')
+	conf_cols = {
+		'ge_minconf': 'Balance',
+	}
 
 	async def __init__(self,proto,*args,**kwargs):
 		self.wallet = await TrackingWallet(proto,mode='w')
 		await super().__init__(proto,*args,**kwargs)
 
 	async def create_data(self):
-		data = self.wallet.mmid_ordered_dict
-		amt0 = self.proto.coin_amt('0')
-		for d in data:
+		in_data = self.wallet.mmid_ordered_dict
+		for d in in_data:
 			if d.type == 'mmgen':
-				key = d.obj.sid
-				if key not in self.data:
-					self.data[key] = [amt0] * 4
+				label = d.obj.sid
+				if label not in self.data:
+					self.data[label] = self.balance_info()
 			else:
-				key = 'Non-MMGen'
+				label = 'Non-MMGen'
 
-			conf_level = 2 # TODO
-			amt = await self.wallet.get_balance(data[d]['addr'])
+			amt = await self.wallet.get_balance(in_data[d]['addr'])
 
-			self.data['TOTAL'][conf_level] += amt
-			self.data[key][conf_level] += amt
+			self.data['TOTAL']['ge_minconf'] += amt
+			self.data[label]['ge_minconf'] += amt
 
 		del self.wallet
 
-class EthereumTokenTwGetBalance(EthereumTwGetBalance): pass
+class EthereumTokenTwGetBalance(EthereumTwGetBalance):
+	pass

mmgen/tool/rpc.py

@@ -42,7 +42,8 @@ class tool_cmd(tool_cmd_base):
 			pager:   'send output to pager' = False ):
 		"list confirmed/unconfirmed, spendable/unspendable balances in tracking wallet"
 		from ..tw.bal import TwGetBalance
-		return (await TwGetBalance(self.proto,minconf,quiet)).format()
+		from ..globalvars import g
+		return (await TwGetBalance(self.proto,minconf,quiet)).format(color=g.color)
 
 	async def twops(self,
 			obj,pager,reverse,detail,sort,age_fmt,interactive,

mmgen/tw/bal.py

@@ -20,9 +20,11 @@
 twbal: Tracking wallet getbalance class for the MMGen suite
 """
 
-from ..color import red,green
+from collections import namedtuple
+
 from ..base_obj import AsyncInit
 from ..objmethods import MMGenObject
+from ..obj import NonNegativeInt
 from ..rpc import rpc_init
 
 class TwGetBalance(MMGenObject,metaclass=AsyncInit):
@@ -32,39 +34,77 @@ class TwGetBalance(MMGenObject,metaclass=AsyncInit):
 
 	async def __init__(self,proto,minconf,quiet):
 
-		self.minconf = minconf
+		class BalanceInfo(dict):
+			def __init__(self):
+				amt0 = proto.coin_amt('0')
+				data = {
+					'unconfirmed': amt0,
+					'lt_minconf': amt0,
+					'ge_minconf': amt0,
+					'spendable': amt0,
+				}
+				return dict.__init__(self,**data)
+
+		self.minconf = NonNegativeInt(minconf)
+		self.balance_info = BalanceInfo
 		self.quiet = quiet
-		self.data = {k:[proto.coin_amt('0')] * 4 for k in ('TOTAL','Non-MMGen','Non-wallet')}
-		self.rpc = await rpc_init(proto)
 		self.proto = proto
+		self.data = {k:self.balance_info() for k in self.start_labels}
+		self.rpc = await rpc_init(proto)
+
+		if minconf < 2 and 'lt_minconf' in self.conf_cols:
+			del self.conf_cols['lt_minconf']
+
 		await self.create_data()
 
-	def format(self):
+	def format(self,color):
+
 		def gen_output():
-			if self.proto.chain_name != 'mainnet':
-				yield 'Chain: ' + green(self.proto.chain_name.upper())
 
 			if self.quiet:
-				yield str(self.data['TOTAL'][2] if self.data else 0)
+				yield str(self.data['TOTAL']['ge_minconf'] if self.data else 0)
 			else:
-				yield self.fs.format(
-					w = 'Wallet',
-					u = ' Unconfirmed',
-					p = f' <{self.minconf} confirms',
-					c = f' >={self.minconf} confirms' )
-
-				for key in sorted(self.data):
-					if not any(self.data[key]):
-						continue
-					yield self.fs.format(**dict(zip(
-						('w','u','p','c'),
-						[key+':'] + [a.fmt(color=True,suf=' '+self.proto.dcoin) for a in self.data[key]]
-						)))
-
-			for key,vals in list(self.data.items()):
-				if key == 'TOTAL':
+
+				def get_col_iwidth(colname):
+					return len(str(int(max(v[colname] for v in self.data.values())))) + iwidth_adj
+
+				def make_col(label,col):
+					return(self.data[label][col].fmt(fs=f'{iwidths[col]}.{add_w-1}',color=color))
+
+				if color:
+					from ..color import red,green,yellow
+				else:
+					from ..color import nocolor
+					red = green = yellow = nocolor
+
+				add_w = self.proto.coin_amt.max_prec + 1 # 1 = len('.')
+				iwidth_adj = 1 # so that min iwidth (1) + add_w + iwidth_adj >= len('Unconfirmed')
+				col1_w = max(len(l) for l in self.start_labels) + 1 # 1 = len(':')
+
+				iwidths = {colname: get_col_iwidth(colname) for colname in self.conf_cols}
+
+				net_desc = self.proto.coin + ' ' + self.proto.network.upper()
+				if net_desc != 'BTC MAINNET':
+					yield 'Network: {}'.format(green(net_desc))
+
+				yield '{lbl:{w}} {cols}'.format(
+					lbl = 'Wallet',
+					w = col1_w + iwidth_adj,
+					cols = ' '.join(v.format(minconf=self.minconf).ljust(iwidths[k]+add_w)
+						for k,v in self.conf_cols.items()) )
+
+				from ..addr import MMGenID
+				for label in sorted(self.data.keys()):
+					yield '{lbl} {cols}'.format(
+						lbl = yellow((label + ' ' + self.proto.coin).ljust(col1_w)) if label == 'TOTAL'
+							else MMGenID.hlc((label+':').ljust(col1_w),color=color),
+						cols = ' '.join(make_col(label,col) for col in self.conf_cols)
+					)
+
+			for k,v in self.data.items():
+				if k == 'TOTAL':
 					continue
-				if vals[3]:
-					yield red(f'Warning: this wallet contains PRIVATE KEYS for {key} outputs!')
+				if v['spendable']:
+					yield red(f'Warning: this wallet contains PRIVATE KEYS for {k} outputs!')
 
 		return '\n'.join(gen_output()).rstrip()

test/test_py_d/ts_ethdev.py

@@ -838,16 +838,16 @@ class TestSuiteEthdev(TestSuiteBase,TestSuiteShared):
 		assert re.search(ss,text),ss
 		return t
 
-	def bal_getbalance(self,idx,etc_adj=False,extra_args=[]):
+	def bal_getbalance(self,sid,idx,etc_adj=False,extra_args=[]):
 		bal1 = token_bals_getbalance[idx][0]
 		bal2 = token_bals_getbalance[idx][1]
 		bal1 = Decimal(bal1)
 		if etc_adj and self.proto.coin == 'ETC':
 			bal1 += self.bal_corr
 		t = self.spawn('mmgen-tool', self.eth_args + extra_args + ['getbalance'])
-		t.expect(r'\n[0-9A-F]{8}: .*\D'+str(bal1),regex=True)
-		t.expect(r'\nNon-MMGen: .*\D'+bal2,regex=True)
-		total = strip_ansi_escapes(t.expect_getend(r'\nTOTAL:\s+',regex=True)).split()[0]
+		t.expect(rf'{sid}:.*'+str(bal1),regex=True)
+		t.expect(r'Non-MMGen:.*'+bal2,regex=True)
+		total = strip_ansi_escapes(t.expect_getend(rf'TOTAL {self.proto.coin}')).split()[0]
 		assert Decimal(bal1) + Decimal(bal2) == Decimal(total)
 		return t
 
@@ -1122,7 +1122,7 @@ class TestSuiteEthdev(TestSuiteBase,TestSuiteShared):
 		return t
 
 	def bal1_getbalance(self):
-		return self.bal_getbalance('1',etc_adj=True)
+		return self.bal_getbalance(dfl_sid,'1',etc_adj=True)
 
 	def addrimport_token_burn_addr(self):
 		return self.addrimport_one_addr(addr=burn_addr,extra_args=['--token=mm1'])
@@ -1131,7 +1131,7 @@ class TestSuiteEthdev(TestSuiteBase,TestSuiteShared):
 		return self.token_bal(n='4')
 
 	def token_bal_getbalance(self):
-		return self.bal_getbalance('2',extra_args=['--token=mm1'])
+		return self.bal_getbalance(dfl_sid,'2',extra_args=['--token=mm1'])
 
 	def txcreate_noamt(self):
 		return self.txcreate(args=['98831F3A:E:12'],eth_fee_res=True)

test/test_py_d/ts_regtest.py

@@ -59,19 +59,19 @@ rt_data = {
 	'rtBals_gb': {
 		'btc': {
 			'0conf0': {
-				'mmgen': ('283.22339537','0','283.22339537'),
-				'nonmm': ('16.77647763','0','116.77629233'),
-				'total': ('299.999873','0','399.9996877'),
+				'mmgen': ('283.22339537','283.22339537'),
+				'nonmm': ('16.77647763','116.77629233'),
+				'total': ('299.999873','399.9996877'),
 			},
 			'0conf1': {
-				'mmgen': ('283.22339537','283.22339537','0'),
-				'nonmm': ('16.77647763','16.77647763','99.9998147'),
-				'total': ('299.999873','299.999873','99.9998147'),
+				'mmgen': ('283.22339537','0'),
+				'nonmm': ('16.77647763','99.9998147'),
+				'total': ('299.999873','99.9998147'),
 			},
 			'1conf1': {
-				'mmgen': ('0','0','283.22339537'),
-				'nonmm': ('0','0','116.77629233'),
-				'total': ('0','0','399.9996877'),
+				'mmgen': ('0','283.22339537'),
+				'nonmm': ('0','116.77629233'),
+				'total': ('0','399.9996877'),
 			},
 			'1conf2': {
 				'mmgen': ('0','283.22339537','0'),
@@ -81,19 +81,19 @@ rt_data = {
 		},
 		'bch': {
 			'0conf0': {
-				'mmgen': ('283.22339437','0','283.22339437'),
-				'nonmm': ('16.77647763','0','116.77637483'),
-				'total': ('299.999872','0','399.9997692'),
+				'mmgen': ('283.22339437','283.22339437'),
+				'nonmm': ('16.77647763','116.77637483'),
+				'total': ('299.999872','399.9997692'),
 			},
 			'0conf1': {
-				'mmgen': ('283.22339437','283.22339437','0'),
-				'nonmm': ('16.77647763','16.77647763','99.9998972'),
-				'total': ('299.999872','299.999872','99.9998972'),
+				'mmgen': ('283.22339437','0'),
+				'nonmm': ('16.77647763','99.9998972'),
+				'total': ('299.999872','99.9998972'),
 			},
 			'1conf1': {
-				'mmgen': ('0','0','283.22339437'),
-				'nonmm': ('0','0','116.77637483'),
-				'total': ('0','0','399.9997692'),
+				'mmgen': ('0','283.22339437'),
+				'nonmm': ('0','116.77637483'),
+				'total': ('0','399.9997692'),
 			},
 			'1conf2': {
 				'mmgen': ('0','283.22339437','0'),
@@ -103,19 +103,19 @@ rt_data = {
 		},
 		'ltc': {
 			'0conf0': {
-				'mmgen': ('283.21717237','0','283.21717237'),
-				'nonmm': ('16.77647763','0','5116.77036263'),
-				'total': ('299.99365','0','5399.987535'),
+				'mmgen': ('283.21717237','283.21717237'),
+				'nonmm': ('16.77647763','5116.77036263'),
+				'total': ('299.99365','5399.987535'),
 			},
 			'0conf1': {
-				'mmgen': ('283.21717237','283.21717237','0'),
-				'nonmm': ('16.77647763','16.77647763','5099.993885'),
-				'total': ('299.99365','299.99365','5099.993885'),
+				'mmgen': ('283.21717237','0'),
+				'nonmm': ('16.77647763','5099.993885'),
+				'total': ('299.99365','5099.993885'),
 			},
 			'1conf1': {
-				'mmgen': ('0','0','283.21717237'),
-				'nonmm': ('0','0','5116.77036263'),
-				'total': ('0','0','5399.987535'),
+				'mmgen': ('0','283.21717237'),
+				'nonmm': ('0','5116.77036263'),
+				'total': ('0','5399.987535'),
 			},
 			'1conf2': {
 				'mmgen': ('0','283.21717237','0'),
@@ -727,16 +727,21 @@ class TestSuiteRegtest(TestSuiteBase,TestSuiteShared):
 		return t
 
 	def bob_getbalance(self,bals,confs=1):
-		for i in (0,1,2):
+		for i in range(len(bals['mmgen'])):
 			assert Decimal(bals['mmgen'][i]) + Decimal(bals['nonmm'][i]) == Decimal(bals['total'][i])
+		sid = self._user_sid('bob')
 		t = self.spawn('mmgen-tool',['--bob','getbalance',f'minconf={confs}'])
 		t.expect('Wallet')
-		for k in ('mmgen','nonmm','total'):
-			ret = strip_ansi_escapes(t.expect_getend(r'\S+: ',regex=True))
+		for k,lbl in (
+			('mmgen',f'{sid}:'),
+			('nonmm','Non-MMGen:'),
+			('total',f'TOTAL {self.proto.coin}')
+		):
+			ret = strip_ansi_escapes(t.expect_getend(lbl + ' '))
 			cmp_or_die(
 				' '.join(bals[k]),
-				re.sub(rf'\s+{self.proto.coin}\s*',' ',ret).strip(),
-				desc=k,
+				' '.join(ret.split()),
+				desc = k,
 			)
 		return t