added support for rpm packaging and basic support for deb

lib/huey/registry.py

@@ -0,0 +1,77 @@
+import pickle
+
+from huey.exceptions import QueueException
+
+
+class TaskRegistry(object):
+    """
+    A simple Registry used to track subclasses of :class:`QueueTask` - the
+    purpose of this registry is to allow translation from queue messages to
+    task classes, and vice-versa.
+    """
+    _ignore = ['QueueTask', 'PeriodicQueueTask']
+
+    _registry = {}
+    _periodic_tasks = []
+
+    def task_to_string(self, task):
+        return '%s' % (task.__name__)
+
+    def register(self, task_class):
+        klass_str = self.task_to_string(task_class)
+        if klass_str in self._ignore:
+            return
+
+        if klass_str not in self._registry:
+            self._registry[klass_str] = task_class
+
+            # store an instance in a separate list of periodic tasks
+            if hasattr(task_class, 'validate_datetime'):
+                self._periodic_tasks.append(task_class())
+
+    def unregister(self, task_class):
+        klass_str = self.task_to_string(task_class)
+
+        if klass_str in self._registry:
+            del(self._registry[klass_str])
+
+            for task in self._periodic_tasks:
+                if isinstance(task, task_class):
+                    self._periodic_tasks.remove(task)
+
+    def __contains__(self, klass_str):
+        return klass_str in self._registry
+
+    def get_message_for_task(self, task):
+        """Convert a task object to a message for storage in the queue"""
+        return pickle.dumps((
+            task.task_id,
+            self.task_to_string(type(task)),
+            task.execute_time,
+            task.retries,
+            task.retry_delay,
+            task.get_data(),
+        ))
+
+    def get_task_class(self, klass_str):
+        klass = self._registry.get(klass_str)
+
+        if not klass:
+            raise QueueException('%s not found in TaskRegistry' % klass_str)
+
+        return klass
+
+    def get_task_for_message(self, msg):
+        """Convert a message from the queue into a task"""
+        # parse out the pieces from the enqueued message
+        raw = pickle.loads(msg)
+        task_id, klass_str, execute_time, retries, delay, data = raw
+
+        klass = self.get_task_class(klass_str)
+        return klass(data, task_id, execute_time, retries, delay)
+
+    def get_periodic_tasks(self):
+        return self._periodic_tasks
+
+
+registry = TaskRegistry()