From 205fec178592c197fcc9fe2b04c96fc7e2a5d338 Mon Sep 17 00:00:00 2001
From: The MMGen Project <mmgen@tuta.io>
Date: Wed, 6 Mar 2024 11:05:22 +0000
Subject: [PATCH] TwGetBalance: offload spendable warning to btc subclass

---
 mmgen/proto/btc/tw/bal.py | 21 +++++++++++++++++++++
 mmgen/tw/bal.py           | 12 ++----------
 2 files changed, 23 insertions(+), 10 deletions(-)

diff --git a/mmgen/proto/btc/tw/bal.py b/mmgen/proto/btc/tw/bal.py
index 0405ca90..73f08a47 100755
--- a/mmgen/proto/btc/tw/bal.py
+++ b/mmgen/proto/btc/tw/bal.py
@@ -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
diff --git a/mmgen/tw/bal.py b/mmgen/tw/bal.py
index c0224cc4..c9d84656 100755
--- a/mmgen/tw/bal.py
+++ b/mmgen/tw/bal.py
@@ -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())