diff --git a/battleman.py b/battleman.py index 4811e3e..ae86be9 100755 --- a/battleman.py +++ b/battleman.py @@ -12,7 +12,7 @@ import sys sys.path.append('lib/') -import cPickle as pickle +import shelve import argparse import cmd import os.path @@ -25,15 +25,14 @@ import easyinput def main(): btl = battle.Battle() - ### This is the pickling jar - battle_pickle = None - bp_io_failed = False - BP_FILE = os.path.expanduser('~/.config/4etools/battleman/battle.pickle') + ### Shelf data + session = None + SESSION_FILE = os.path.expanduser('~/.config/4etools/battleman/battleman-restore') ### # Make sure config directory exists - if not os.path.exists(os.path.dirname(BP_FILE)): - os.makedirs(os.path.dirname(BP_FILE)) + if not os.path.exists(os.path.dirname(SESSION_FILE)): + os.makedirs(os.path.dirname(SESSION_FILE)) # Get command-line args settings = parse_args() @@ -42,14 +41,14 @@ def main(): print "Welcome to 4e Battle Manager.\n" # Resume battle if needed - # fixme: critical: indexes get reused when restoring! Maybe switch to shelve? + session = shelve.open(SESSION_FILE) + if settings.resume: try: - with open(BP_FILE, 'r') as f: - btl = pickle.load(f) - battle_pickle = pickle.dumps(btl) + btl = session['battle'] + Combatant.next_index = session['combatant_next_index'] except: - print "Error: Couldn't resume. Quitting to preserve our pickle." + print "Error: Couldn't resume. Quitting to preserve our old session." sys.exit(1) else: @@ -83,38 +82,35 @@ def main(): print btl - cmd_parser = CommandParser(btl, battle_pickle, BP_FILE) + cmd_parser = CommandParser(btl, session) cmd_parser.cmdloop() + session.close() + class CommandParser(cmd.Cmd): """Parse the commands from the command-line.""" - def __init__(self, btl, battle_pickle, BP_FILE): + def __init__(self, btl, session): cmd.Cmd.__init__(self) self.btl = btl - self.battle_pickle = battle_pickle - self.BP_FILE = BP_FILE + self.session = session + self.session_io_failed = False self.doc_header = 'Available commands (type help for more help)' self.prompt = '\n> ' def postloop(self): - # Re-pickle and write if changed after every query. It's cheap - # and we only have to run at user-speed anyway - old_bp = self.battle_pickle - self.battle_pickle = pickle.dumps(self.btl) - - if old_bp != self.battle_pickle: - try: - with open(self.BP_FILE, 'w') as f: - f.write(self.battle_pickle) - except Exception: - if not self.bp_io_failed: - print("Warning: can't write the battle pickle. Resuming later will fail.") - self.bp_io_failed = True + # Just dumbly dump to the shelf DB + try: + self.session['battle'] = self.btl + self.session['combatant_next_index'] = Combatant.next_index + except Exception: + if not self.session_io_failed: + print("Warning: can't write to the session database. Resuming later will fail.") + self.session_io_failed = True # This allows us to do partial command completion without , @@ -404,9 +400,9 @@ class CommandParser(cmd.Cmd): def do_sync(self, line): """sync - This function is still a stub""" + Save the current battle data to the session database. Really not necessary, but nice to have for peace of mind.""" - # Since the postloop pickles & writes, we don't actually need to + # Since the postloop does this, we don't actually need to # do a damn thing here - just let it be an 'empty' command. pass