source file: /opt/devel/celery/celery/registry.py
file stats: 46 lines, 45 executed: 97.8% covered
   1. from celery import discovery
   2. from UserDict import UserDict
   3. 
   4. 
   5. class NotRegistered(Exception):
   6.     """The task is not registered."""
   7. 
   8. 
   9. class AlreadyRegistered(Exception):
  10.     """The task is already registered."""
  11. 
  12. 
  13. class TaskRegistry(UserDict):
  14.     """Site registry for tasks."""
  15. 
  16.     AlreadyRegistered = AlreadyRegistered
  17.     NotRegistered = NotRegistered
  18. 
  19.     def __init__(self):
  20.         self.data = {}
  21. 
  22.     def autodiscover(self):
  23.         discovery.autodiscover()
  24. 
  25.     def register(self, task, task_name=None):
  26.         is_class = False
  27.         if hasattr(task, "run"):
  28.             is_class = True
  29.             task_name = task.name
  30.         if task_name in self.data:
  31.             raise self.AlreadyRegistered(
  32.                     "Task with name %s is already registered." % task_name)
  33. 
  34.         if is_class:
  35.             self.data[task_name] = task() # instantiate Task class
  36.         else:
  37.             task.type = "regular"
  38.             self.data[task_name] = task
  39. 
  40.     def unregister(self, task_name):
  41.         if hasattr(task_name, "run"):
  42.             task_name = task_name.name
  43.         if task_name not in self.data:
  44.             raise self.NotRegistered(
  45.                     "Task with name %s is not registered." % task_name)
  46.         del self.data[task_name]
  47. 
  48.     def get_all(self):
  49.         """Get all task types."""
  50.         return self.data
  51. 
  52.     def filter_types(self, type):
  53.         """Return all tasks of a specific type."""
  54.         return dict([(task_name, task)
  55.                         for task_name, task in self.data.items()
  56.                             if task.type == type])
  57. 
  58.     def get_all_regular(self):
  59.         """Get all regular task types."""
  60.         return self.filter_types(type="regular")
  61. 
  62.     def get_all_periodic(self):
  63.         """Get all periodic task types."""
  64.         return self.filter_types(type="periodic")
  65. 
  66.     def get_task(self, task_name):
  67.         """Get task by name."""
  68.         return self.data[task_name]
  69. 
  70. """This is the global task registry."""
  71. tasks = TaskRegistry()