battleman.py: Fix broken indexes after resuming

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
+    ### Shelf data
-    battle_pickle = None
+    session = None
-    bp_io_failed = False
+    SESSION_FILE = os.path.expanduser('~/.config/4etools/battleman/battleman-restore')
-    BP_FILE = os.path.expanduser('~/.config/4etools/battleman/battle.pickle')
     ###
 
     # Make sure config directory exists
-    if not os.path.exists(os.path.dirname(BP_FILE)):
+    if not os.path.exists(os.path.dirname(SESSION_FILE)):
-        os.makedirs(os.path.dirname(BP_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 = session['battle']
-                btl = pickle.load(f)
+            Combatant.next_index = session['combatant_next_index']
-                battle_pickle = pickle.dumps(btl)
         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.session = session
-        self.BP_FILE = BP_FILE
+        self.session_io_failed = False
         self.doc_header = 'Available commands (type help <command> for more help)'
         self.prompt = '\n> '
 
 
     def postloop(self):
-        # Re-pickle and write if changed after every query. It's cheap
+        # Just dumbly dump to the shelf DB
-        # and we only have to run at user-speed anyway
+        try:
-        old_bp = self.battle_pickle
+            self.session['battle'] = self.btl
-        self.battle_pickle = pickle.dumps(self.btl)
+            self.session['combatant_next_index'] = Combatant.next_index
-
+        except Exception:
-        if old_bp != self.battle_pickle:
+            if not self.session_io_failed:
-            try:
+                print("Warning: can't write to the session database. Resuming later will fail.")
-                with open(self.BP_FILE, 'w') as f:
+                self.session_io_failed = True
-                    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
 
 
     # This allows us to do partial command completion without <tab>,
@@ -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