registry.py 1.9 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778
  1. """celery.registry"""
  2. import inspect
  3. from UserDict import UserDict
  4. from celery.exceptions import NotRegistered
  5. class TaskRegistry(UserDict):
  6. """Site registry for tasks."""
  7. NotRegistered = NotRegistered
  8. def __init__(self):
  9. self.data = {}
  10. def regular(self):
  11. """Get all regular task types."""
  12. return self.filter_types("regular")
  13. def periodic(self):
  14. """Get all periodic task types."""
  15. return self.filter_types("periodic")
  16. def register(self, task):
  17. """Register a task in the task registry.
  18. The task will be automatically instantiated if not already an
  19. instance.
  20. """
  21. task = inspect.isclass(task) and task() or task
  22. name = task.name
  23. self.data[name] = task
  24. def unregister(self, name):
  25. """Unregister task by name.
  26. :param name: name of the task to unregister, or a
  27. :class:`celery.task.base.Task` with a valid `name` attribute.
  28. :raises celery.exceptions.NotRegistered: if the task has not
  29. been registered.
  30. """
  31. try:
  32. # Might be a task class
  33. name = name.name
  34. except AttributeError:
  35. pass
  36. self.pop(name)
  37. def filter_types(self, type):
  38. """Return all tasks of a specific type."""
  39. return dict((task_name, task)
  40. for task_name, task in self.data.items()
  41. if task.type == type)
  42. def __getitem__(self, key):
  43. try:
  44. return UserDict.__getitem__(self, key)
  45. except KeyError, exc:
  46. raise self.NotRegistered(str(exc))
  47. def pop(self, key, *args):
  48. try:
  49. return UserDict.pop(self, key, *args)
  50. except KeyError, exc:
  51. raise self.NotRegistered(str(exc))
  52. """
  53. .. data:: tasks
  54. The global task registry.
  55. """
  56. tasks = TaskRegistry()