TwGetBalance: offload spendable warning to btc subclass

This commit is contained in:
The MMGen Project 2024-03-06 11:05:22 +00:00
commit 205fec1785
Signed by: mmgen
GPG key ID: 3F8B1861E32B7DA2
2 changed files with 23 additions and 10 deletions

View file

@ -14,9 +14,14 @@ proto.btc.tw.bal: Bitcoin base protocol tracking wallet balance class
from ....tw.bal import TwGetBalance
from ....tw.shared import get_tw_label
from ....rpc import rpc_init
class BitcoinTwGetBalance(TwGetBalance):
async def __init__(self, cfg, proto, minconf, quiet):
self.rpc = await rpc_init(cfg, proto)
await super().__init__(cfg, proto, minconf, quiet)
start_labels = ('TOTAL','Non-MMGen','Non-wallet')
conf_cols = {
'unconfirmed': 'Unconfirmed',
@ -50,3 +55,19 @@ class BitcoinTwGetBalance(TwGetBalance):
if d['spendable']:
self.data[label]['spendable'] += amt
def format(self, color):
def gen_spendable_warning():
for k,v in self.data.items():
if v['spendable']:
yield red(f'Warning: this wallet contains PRIVATE KEYS for {k} outputs!')
if color:
from ....color import red
else:
from ....color import nocolor as red
warning = '\n'.join(gen_spendable_warning())
return super().format(color) + ('\n' if warning else '') + warning

View file

@ -23,7 +23,6 @@ tw.bal: Tracking wallet getbalance class for the MMGen suite
from ..base_obj import AsyncInit
from ..objmethods import MMGenObject
from ..obj import NonNegativeInt
from ..rpc import rpc_init
class TwGetBalance(MMGenObject,metaclass=AsyncInit):
@ -48,7 +47,6 @@ class TwGetBalance(MMGenObject,metaclass=AsyncInit):
self.quiet = quiet
self.proto = proto
self.data = {k:self.balance_info() for k in self.start_labels}
self.rpc = await rpc_init(cfg,proto)
if minconf < 2 and 'lt_minconf' in self.conf_cols:
del self.conf_cols['lt_minconf']
@ -70,10 +68,10 @@ class TwGetBalance(MMGenObject,metaclass=AsyncInit):
return self.data[label][col].fmt( iwidth=iwidths[col], color=color )
if color:
from ..color import red,green,yellow
from ..color import green,yellow
else:
from ..color import nocolor
red = green = yellow = nocolor
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')
@ -99,10 +97,4 @@ class TwGetBalance(MMGenObject,metaclass=AsyncInit):
cols = ' '.join(make_col(label,col) for col in self.conf_cols)
).rstrip()
for k,v in self.data.items():
if k == 'TOTAL':
continue
if v['spendable']:
yield red(f'Warning: this wallet contains PRIVATE KEYS for {k} outputs!')
return '\n'.join(gen_output())