tw.ctl: move store-related methods + init to tw.store

This commit is contained in:
The MMGen Project 2025-05-28 11:40:40 +00:00
commit 6e084d4f71
Signed by: mmgen
GPG key ID: 3F8B1861E32B7DA2
3 changed files with 189 additions and 172 deletions

View file

@ -17,12 +17,10 @@ from ....util import msg, msg_r, rmsg, die, suf, fmt_list
class BitcoinTwCtl(TwCtl):
def init_empty(self):
self.data = {'coin': self.proto.coin, 'addresses': {}}
async def rpc_get_balance(self, addr, block='latest'):
raise NotImplementedError('not implemented')
# TODO: do check with check_import_mmid()
@write_mode
async def import_address(self, addr, *, label, rescan=False):
if (await self.rpc.walletinfo).get('descriptors'):

View file

@ -20,11 +20,9 @@
tw.ctl: Tracking wallet control class for the MMGen suite
"""
import json
from collections import namedtuple
from pathlib import Path
from ..util import msg, msg_r, suf, die
from ..util import msg, msg_r, ymsg, suf, die
from ..base_obj import AsyncInit
from ..objmethods import MMGenObject
from ..obj import TwComment, get_obj
@ -50,10 +48,7 @@ class TwCtl(MMGenObject, metaclass=AsyncInit):
caps = ('rescan', 'batch')
data_key = 'addresses'
use_tw_file = False
aggressive_sync = False
importing = False
tw_fn = 'tracking-wallet.json'
use_cached_balances = False
def __new__(cls, cfg, proto, *args, **kwargs):
@ -68,7 +63,6 @@ class TwCtl(MMGenObject, metaclass=AsyncInit):
mode = 'r',
token_addr = None,
no_rpc = False,
no_wallet_init = False,
rpc_ignore_wallet = False):
assert mode in ('r', 'w', 'i'), f"{mode!r}: wallet mode must be 'r', 'w' or 'i'"
@ -80,158 +74,10 @@ class TwCtl(MMGenObject, metaclass=AsyncInit):
self.proto = proto
self.mode = mode
self.desc = self.base_desc = f'{self.proto.name} tracking wallet'
self.cur_balances = {} # cache balances to prevent repeated lookups per program invocation
if cfg.cached_balances:
self.use_cached_balances = True
if not no_rpc:
self.rpc = await rpc_init(cfg, proto, ignore_wallet=rpc_ignore_wallet)
if self.use_tw_file:
if self.proto.coin == 'BTC':
self.tw_dir = Path(self.cfg.data_dir)
else:
self.tw_dir = Path(
self.cfg.data_dir_root,
'altcoins',
self.proto.coin.lower(),
('' if self.proto.network == 'mainnet' else self.proto.network)
)
self.tw_path = self.tw_dir / self.tw_fn
if no_wallet_init:
return
if self.use_tw_file:
self.init_from_wallet_file()
else:
self.init_empty()
if self.data['coin'] != self.proto.coin: # TODO remove?
die('WalletFileError',
f'Tracking wallet coin ({self.data["coin"]}) does not match current coin ({self.proto.coin})!')
self.conv_types(self.data[self.data_key])
def upgrade_wallet_maybe(self):
pass
def init_from_wallet_file(self):
from ..fileutil import check_or_create_dir, get_data_from_file
check_or_create_dir(self.tw_dir)
try:
self.orig_data = get_data_from_file(self.cfg, self.tw_path, quiet=True)
self.data = json.loads(self.orig_data)
except:
try:
self.tw_path.stat()
except:
self.orig_data = ''
self.init_empty()
self.force_write()
else:
die('WalletFileError', f'File ‘{self.tw_path}’ exists but does not contain valid JSON data')
else:
self.upgrade_wallet_maybe()
# ensure that wallet file is written when user exits via KeyboardInterrupt:
if self.mode == 'w':
import atexit
def del_twctl(twctl):
self.cfg._util.dmsg(f'Running exit handler del_twctl() for {twctl!r}')
del twctl
atexit.register(del_twctl, self)
def __del__(self):
"""
TwCtl instances opened in write or import mode must be explicitly destroyed with del
twuo.twctl and the like to ensure the instance is deleted and wallet is written before
global vars are destroyed by the interpreter at shutdown.
Not that this code can only be debugged by examining the program output, as exceptions
are ignored within __del__():
/usr/share/doc/python3.6-doc/html/reference/datamodel.html#object.__del__
Since no exceptions are raised, errors will not be caught by the test suite.
"""
if getattr(self, 'mode', None) == 'w': # mode attr might not exist in this state
self.write()
elif self.cfg.debug:
msg('read-only wallet, doing nothing')
def conv_types(self, ad):
for k, v in ad.items():
if k not in ('params', 'coin'):
v['mmid'] = TwMMGenID(self.proto, v['mmid'])
v['comment'] = TwComment(v['comment'])
@property
def data_root(self):
return self.data[self.data_key]
@property
def data_root_desc(self):
return self.data_key
def cache_balance(self, addr, bal, *, session_cache, data_root, force=False):
if force or addr not in session_cache:
session_cache[addr] = str(bal)
if addr in data_root:
data_root[addr]['balance'] = str(bal)
if self.aggressive_sync:
self.write()
def get_cached_balance(self, addr, session_cache, data_root):
if addr in session_cache:
return self.proto.coin_amt(session_cache[addr])
if not self.use_cached_balances:
return None
if addr in data_root and 'balance' in data_root[addr]:
return self.proto.coin_amt(data_root[addr]['balance'])
async def get_balance(self, addr, *, force_rpc=False, block='latest'):
ret = None if force_rpc else self.get_cached_balance(addr, self.cur_balances, self.data_root)
if ret is None:
ret = await self.rpc_get_balance(addr, block=block)
self.cache_balance(addr, ret, session_cache=self.cur_balances, data_root=self.data_root)
return ret
def force_write(self):
mode_save = self.mode
self.mode = 'w'
self.write()
self.mode = mode_save
@write_mode
def write_changed(self, data, quiet):
from ..fileutil import write_data_to_file
write_data_to_file(
self.cfg,
self.tw_path,
data,
desc = f'{self.base_desc} data',
ask_overwrite = False,
ignore_opt_outdir = True,
quiet = quiet,
check_data = True, # die if wallet has been altered by another program
cmp_data = self.orig_data)
self.orig_data = data
def write(self, *, quiet=True):
if not self.use_tw_file:
self.cfg._util.dmsg("'use_tw_file' is False, doing nothing")
return
self.cfg._util.dmsg(f'write(): checking if {self.desc} data has changed')
wdata = json.dumps(self.data)
if self.orig_data != wdata:
self.write_changed(wdata, quiet=quiet)
elif self.cfg.debug:
msg('Data is unchanged\n')
async def resolve_address(self, addrspec):
twmmid, coinaddr = (None, None)
@ -305,6 +151,19 @@ class TwCtl(MMGenObject, metaclass=AsyncInit):
async def remove_comment(self, mmaddr):
await self.set_comment(mmaddr, '')
def check_import_mmid(self, addr, old_mmid, new_mmid):
'returns True if mmid needs update, None otherwise'
if new_mmid != old_mmid:
if old_mmid.endswith(':' + addr):
ymsg(f'Warning: address {new_mmid} was previously imported as non-MMGen!')
return True
else:
fs = (
'attempting to import MMGen address {a!r} ({b}) as non-MMGen!'
if new_mmid.endswith(':' + addr) else
'imported MMGen ID {b!r} does not match tracking wallet MMGen ID {a!r}!')
die(2, fs.format(a=old_mmid, b=new_mmid))
async def import_address_common(self, data, *, batch=False, gather=False):
async def do_import(address, comment, message):

View file

@ -12,17 +12,86 @@
tw.store: Tracking wallet control class with store
"""
import json
from pathlib import Path
from ..base_obj import AsyncInit
from ..obj import TwComment
from ..util import msg, ymsg, die, cached_property
from ..addr import is_coin_addr, is_mmgen_id, CoinAddr
from .shared import TwLabel
from .shared import TwMMGenID, TwLabel
from .ctl import TwCtl, write_mode, label_addr_pair
class TwCtlWithStore(TwCtl):
class TwCtlWithStore(TwCtl, metaclass=AsyncInit):
caps = ('batch',)
data_key = 'addresses'
use_tw_file = True
tw_fn = 'tracking-wallet.json'
aggressive_sync = False
async def __init__(
self,
cfg,
proto,
*,
mode = 'r',
token_addr = None,
no_rpc = False,
no_wallet_init = False,
rpc_ignore_wallet = False):
await super().__init__(cfg, proto, mode=mode, no_rpc=no_rpc, rpc_ignore_wallet=rpc_ignore_wallet)
self.cur_balances = {} # cache balances to prevent repeated lookups per program invocation
if cfg.cached_balances:
self.use_cached_balances = True
self.tw_dir = Path(
self.cfg.data_dir_root,
'altcoins',
self.proto.coin.lower(),
('' if self.proto.network == 'mainnet' else self.proto.network)
)
self.tw_path = self.tw_dir / self.tw_fn
if no_wallet_init:
return
self.init_from_wallet_file()
if self.data['coin'] != self.proto.coin:
fs = 'Tracking wallet coin ({}) does not match current coin ({})!'
die('WalletFileError', fs.format(self.data['coin'], self.proto.coin))
self.conv_types(self.data[self.data_key])
def __del__(self):
"""
TwCtl instances opened in write or import mode must be explicitly destroyed with del
twuo.twctl and the like to ensure the instance is deleted and wallet is written before
global vars are destroyed by the interpreter at shutdown.
Not that this code can only be debugged by examining the program output, as exceptions
are ignored within __del__():
/usr/share/doc/python3.6-doc/html/reference/datamodel.html#object.__del__
Since no exceptions are raised, errors will not be caught by the test suite.
"""
if getattr(self, 'mode', None) == 'w': # mode attr might not exist in this state
self.write()
elif self.cfg.debug:
msg('read-only wallet, doing nothing')
def upgrade_wallet_maybe(self):
pass
def conv_types(self, ad):
for k, v in ad.items():
if k not in ('params', 'coin'):
v['mmid'] = TwMMGenID(self.proto, v['mmid'])
v['comment'] = TwComment(v['comment'])
def init_empty(self):
self.data = {
@ -31,6 +100,32 @@ class TwCtlWithStore(TwCtl):
'addresses': {},
}
def init_from_wallet_file(self):
from ..fileutil import check_or_create_dir, get_data_from_file
check_or_create_dir(self.tw_dir)
try:
self.orig_data = get_data_from_file(self.cfg, self.tw_path, quiet=True)
self.data = json.loads(self.orig_data)
except:
try:
self.tw_path.stat()
except:
self.orig_data = ''
self.init_empty()
self.force_write()
else:
die('WalletFileError', f'File ‘{self.tw_path}’ exists but does not contain valid JSON data')
else:
self.upgrade_wallet_maybe()
# ensure that wallet file is written when user exits via KeyboardInterrupt:
if self.mode == 'w':
import atexit
def del_twctl(twctl):
self.cfg._util.dmsg(f'Running exit handler del_twctl() for {twctl!r}')
del twctl
atexit.register(del_twctl, self)
@write_mode
async def batch_import_address(self, args_list):
return [await self.import_address(a, label=b, rescan=c) for a, b, c in args_list]
@ -42,17 +137,9 @@ class TwCtlWithStore(TwCtl):
async def import_address(self, addr, *, label, rescan=False):
r = self.data_root
if addr in r:
if r[addr]['mmid']:
if r[addr]['mmid'] != label.mmid:
fs = 'imported MMGen ID {!r} does not match tracking wallet MMGen ID {!r}!'
die(3, fs.format(label.mmid, r[addr]['mmid']))
elif label.mmid:
ymsg(f'Warning: MMGen ID {label.mmid!r} was missing in tracking wallet!')
if self.check_import_mmid(addr, r[addr]['mmid'], label.mmid):
r[addr]['mmid'] = label.mmid
if not 'comment' in r[addr]:
ymsg(f'Warning: Label for MMGen ID {label.mmid!r} was missing in tracking wallet!')
r[addr]['comment'] = label.comment
elif label.comment: # overwrite existing comment only if new comment not empty
if label.comment: # overwrite existing comment only if new comment not empty
r[addr]['comment'] = label.comment
else:
r[addr] = {'mmid': label.mmid, 'comment': label.comment}
@ -112,3 +199,76 @@ class TwCtlWithStore(TwCtl):
from decimal import Decimal
# TODO: for now, consider used addrs to be addrs with balance
return ({k for k, v in self.data['addresses'].items() if Decimal(v.get('balance', 0))})
@property
def data_root(self):
return self.data[self.data_key]
@property
def data_root_desc(self):
return self.data_key
def cache_balance(self, addr, bal, *, session_cache, data_root, force=False):
if force or addr not in session_cache:
session_cache[addr] = str(bal)
if addr in data_root:
data_root[addr]['balance'] = str(bal)
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])
if self.use_cached_balances:
return self.proto.coin_amt(
data_root[addr]['balance'] if addr in data_root and 'balance' in data_root[addr]
else '0')
async def get_balance(self, addr, *, force_rpc=False, block='latest'):
ret = None if force_rpc else self.get_cached_balance(addr, self.cur_balances, self.data_root)
if ret is None:
ret = await self.rpc_get_balance(addr, block=block)
if ret is not None:
self.cache_balance(addr, ret, session_cache=self.cur_balances, data_root=self.data_root)
return ret
def force_write(self):
mode_save = self.mode
self.mode = 'w'
self.write()
self.mode = mode_save
@write_mode
def write_changed(self, data, quiet):
from ..fileutil import write_data_to_file
write_data_to_file(
self.cfg,
self.tw_path,
data,
desc = f'{self.base_desc} data',
ask_overwrite = False,
ignore_opt_outdir = True,
quiet = quiet,
check_data = True, # die if wallet has been altered by another program
cmp_data = self.orig_data)
self.orig_data = data
def write(self, *, quiet=True):
self.cfg._util.dmsg(f'write(): checking if {self.desc} data has changed')
wdata = json.dumps(self.data)
if self.orig_data != wdata:
self.write_changed(wdata, quiet=quiet)
elif self.cfg.debug:
msg('Data is unchanged\n')