| 1234567891011121314151617181920212223242526272829303132333435 |
- from functools import wraps
- from huey import RedisHuey
- def _task_wrapper(task_fn, pre_task=None, post_task=None):
- @wraps(task_fn)
- def inner(*args, **kwargs):
- if pre_task is not None:
- pre_task()
- result = task_fn(*args, **kwargs)
- if post_task is not None:
- post_task()
- return result
- return inner
- class RedisHueyExt(RedisHuey):
- def task(self, pre_task=None, post_task=None, *args, **kwargs):
- def decorator(fn):
- return (super(RedisHueyExt, self)
- .task(*args, **kwargs)(_task_wrapper(
- fn,
- pre_task=pre_task,
- post_task=post_task)))
- return decorator
- def periodic_task(self, pre_task=None, post_task=None, *args, **kwargs):
- def decorator(fn):
- return (super(RedisHueyExt, self)
- .periodic_task(*args, **kwargs)(_task_wrapper(
- fn,
- pre_task=pre_task,
- post_task=post_task)))
- return decorator
|
