Improve extensibility

huey/api.py

@@ -81,10 +81,15 @@ def add_numbers(a, b):
         def nightly_report():
             generate_nightly_report()
     """
+    consumer_class = Consumer
     storage_class = None
     _deprecated_params = ('result_store', 'events', 'store_errors',
                           'global_registry')
 
+    @property
+    def task_wrapper_class(self):
+        return TaskWrapper
+
     def __init__(self, name='huey', results=True, store_none=False, utc=True,
                  immediate=False, serializer=None, compression=False,
                  use_zlib=False, immediate_use_memory=True, always_eager=None,
@@ -162,12 +167,12 @@ def immediate(self, value):
                 self.storage = self.create_storage()
 
     def create_consumer(self, **options):
-        return Consumer(self, **options)
+        return self.consumer_class(self, **options)
 
     def task(self, retries=0, retry_delay=0, priority=None, context=False,
              name=None, expires=None, **kwargs):
         def decorator(func):
-            return TaskWrapper(
+            return self.task_wrapper_class(
                 self,
                 func.func if isinstance(func, TaskWrapper) else func,
                 context=context,
@@ -186,7 +191,7 @@ def decorator(func):
             def method_validate(self, timestamp):
                 return validate_datetime(timestamp)
 
-            return TaskWrapper(
+            return self.task_wrapper_class(
                 self,
                 func.func if isinstance(func, TaskWrapper) else func,
                 context=context,

huey/consumer.py

@@ -21,7 +21,7 @@
 from huey.constants import WORKER_PROCESS
 from huey.constants import WORKER_THREAD
 from huey.constants import WORKER_TYPES
-from huey.exceptions import ConfigurationError
+from huey.exceptions import ConfigurationError, StopFlagException
 from huey.utils import time_clock
 
 
@@ -426,36 +426,47 @@ def stop(self, graceful=False):
         else:
             self._logger.info('Shutting down')
 
+    def run_one_iteration(self, health_check_ts):
+        """
+        Run one iteration of the consumer event loop.
+        """
+        timeout = self._stop_flag_timeout
+        try:
+            self.stop_flag.wait(timeout=timeout)
+        except KeyboardInterrupt:
+            self._logger.info('Received SIGINT')
+            self.stop(graceful=True)
+        except:
+            self._logger.exception('Error in consumer.')
+            self.stop()
+        else:
+            if self._received_signal:
+                self.stop(graceful=self._graceful)
+
+        if self.stop_flag.is_set():
+            raise StopFlagException
+
+        if self._health_check:
+            now = time_clock()
+            if now >= health_check_ts + self._health_check_interval:
+                health_check_ts = now
+                self.check_worker_health()
+
+        return health_check_ts
+
     def run(self):
         """
         Run the consumer.
         """
         self.start()
-        timeout = self._stop_flag_timeout
         health_check_ts = time_clock()
 
         while True:
             try:
-                self.stop_flag.wait(timeout=timeout)
-            except KeyboardInterrupt:
-                self._logger.info('Received SIGINT')
-                self.stop(graceful=True)
-            except:
-                self._logger.exception('Error in consumer.')
-                self.stop()
-            else:
-                if self._received_signal:
-                    self.stop(graceful=self._graceful)
-
-            if self.stop_flag.is_set():
+                health_check_ts = self.run_one_iteration(health_check_ts)
+            except StopFlagException:
                 break
 
-            if self._health_check:
-                now = time_clock()
-                if now >= health_check_ts + self._health_check_interval:
-                    health_check_ts = now
-                    self.check_worker_health()
-
         self.huey.notify_interrupted_tasks()
 
         if self._restart:

huey/exceptions.py

@@ -1,6 +1,7 @@
 class HueyException(Exception): pass
 class ConfigurationError(HueyException): pass
 class TaskLockedException(HueyException): pass
+class StopFlagException(HueyException): pass
 
 class CancelExecution(Exception):
     def __init__(self, retry=None, *args, **kwargs):