proto.rune.tw.ctl: support zeroing balances manually

mmgen/proto/rune/tw/ctl.py

@@ -12,7 +12,24 @@
 proto.rune.tw.ctl: THORChain tracking wallet control class
 """
 
+from ....util import ymsg
 from ....tw.store import TwCtlWithStore
 
 class THORChainTwCtl(TwCtlWithStore):
+
 	use_cached_balances = True
+
+	async def rpc_get_balance(self, addr, block='latest'):
+		assert self.rpc.is_remote, 'tw.store.rpc_get_balance(): RPC is not remote!'
+		try:
+			return self.rpc.get_balance(addr, block=block)
+		except Exception as e:
+			ymsg(f'{type(e).__name__}: {e}')
+			ymsg(f'Unable to get balance for address ‘{addr}’')
+			if self.get_cached_balance(addr, self.cur_balances, self.data_root):
+				from ....ui import keypress_confirm
+				if keypress_confirm(self.cfg, 'Zero the balance for this address?'):
+					return self.proto.coin_amt('0')
+			else:
+				import asyncio
+				await asyncio.sleep(3)

mmgen/tw/store.py

@@ -17,7 +17,7 @@ from pathlib import Path
 
 from ..base_obj import AsyncInit
 from ..obj import TwComment
-from ..util import msg, ymsg, die, cached_property
+from ..util import msg, die, cached_property
 from ..addr import is_coin_addr, is_mmgen_id, CoinAddr
 
 from .shared import TwMMGenID, TwLabel
@@ -216,16 +216,6 @@ class TwCtlWithStore(TwCtl, metaclass=AsyncInit):
 				if self.aggressive_sync:
 					self.write()
 
-	async def rpc_get_balance(self, addr, block='latest'):
-		assert self.rpc.is_remote, 'tw.store.rpc_get_balance(): RPC is not remote!'
-		try:
-			return self.rpc.get_balance(addr, block=block)
-		except Exception as e:
-			ymsg(f'{type(e).__name__}: {e}')
-			ymsg(f'Unable to get balance for address ‘{addr}’')
-			import asyncio
-			await asyncio.sleep(3)
-
 	def get_cached_balance(self, addr, session_cache, data_root):
 		if addr in session_cache:
 			return self.proto.coin_amt(session_cache[addr])

mmgen/tw/view.py

@@ -743,6 +743,8 @@ class TwView(MMGenObject, metaclass=AsyncInit):
 				await parent.get_data()
 				parent.oneshot_msg = yellow(
 					f'{parent.proto.dcoin} balance for {parent.item_desc} #{idx} refreshed')
+				if res == 0:
+					return False # zeroing balance may mess up display
 
 		async def i_addr_delete(self, parent, idx):
 			if not parent.keypress_confirm(