Use shelve instead of PersistentDict

celery/beat.py

@@ -1,12 +1,12 @@
-from celery import registry
-from datetime import datetime
 from UserDict import UserDict
-from celery.serialization import pickle
+from datetime import datetime
+from celery import registry
+import shelve
 import atexit
-import errno
 import time
 
-schedule = PersistentDict(save_at_exit=True)
+schedule = shelve.open(filename="celerybeat-schedule")
+atexit.register(schedule.close)
 
 
 class ScheduleEntry(object):
@@ -40,7 +40,7 @@ class ScheduleEntry(object):
         return True if datetime.now() > self.last_run_at else return False
 
 
-class Scheduler(object):
+class Scheduler(UserDict):
     """Scheduler for periodic tasks.
 
     :keyword registry: The task registry to use.
@@ -61,9 +61,12 @@ class Scheduler(object):
         """Run the scheduler.
 
         This runs :meth:`tick` every second in a never-exit loop."""
-        while True:
-            self.tick() 
-            time.sleep(1)
+        try:
+            while True:
+                self.tick() 
+                time.sleep(1)
+        finally:
+            self.schedule.close()
 
     def tick(self):
         """Run a tick, that is one iteration of the scheduler.
@@ -85,68 +88,3 @@ class Scheduler(object):
     @property
     def schedule(self):
         return self.data
-
-
-class PersistentDict(UserDict):
-    """Dictionary that can be stored to disk.
-
-    :param filename: Name of the file to save to.
-
-    :keyword initial_data: Initial dict to start with.
-    :keyword save_at_exit: Register an atexit handler to automatically save
-        the data when the program exits (not safe, but as an extra precaution)
-    :keyword encoding: The encoding to write the file with, default is
-        ``"zlib"``.
-
-    """
-    encoding = "zlib"
-    save_at_exit = False
-
-    def __init__(self, filename, initial_data=None, save_at_exit=None,
-            encoding=None):
-        self.data = initial_data or {}
-        self.filename = filename
-        self.encoding = encoding
-        if save_at_exit is not None:
-            self.save_at_exit = save_at_exit
-        self.reload()
-        self._saved = False
-        self.save_at_exit and self.register_atexit()
-
-    def reload(self):
-        """Reload data from disk."""
-        persisted_data = self._read_file(self.filename, self.encoding)
-        self.data = dict(self.data, **persisted_data)
-
-    def save(self):
-        """Save data to disk."""
-        self._saved = True
-        encoded = pickle.dump(fh, self.data).encode(self.encoding)
-
-        fh = open(self.filename, "w")
-        try:
-            fh.write(encoded)
-        finally:
-            fh.close()
-
-    def register_atexit(self):
-        """Register an atexit handler to save data to disk when the
-        program terminates."""
-        atexit.register(self._save_atexit)
-
-    def _save_atexit(self):
-        self._saved or self.save()
-
-    def _read_file(self, filename, encoding):
-        try:
-            fh = open(filename)
-        except IOError, exc:
-            if exc.errno == errno.ENOENT:
-                return
-            raise
-
-        try:
-            return pickle.loads(fh.read().decode(encoding))
-        finally:
-            fh.close()
-