如何编写正确将参数传递给芹菜任务的自定义装饰器?

时间:2013-10-11 15:36:56

标签: python decorator django-celery python-decorators djcelery

背景和问题

我正在使用Django 1.5.1和django-celery 3.0.17。我想编写一个自定义装饰器,以确保一次只运行一个函数实例,类似于this但没有重复的try / finally。

如何编写装饰器,以便可以使用参数调用像applydelay这样的Celery方法?

Others创造了这样的装饰者,取得了明显的成功。我错过了什么?

实现

我已经尝试将装饰器编写为函数和类,但是使用任一实现,当尝试用我的装饰器和@celery.task装饰函数时,参数不会传递给装饰函数,导致错误消息:

TypeError: foo() takes exactly 1 argument (0 given),其中foo是已修饰函数的名称。

功能实施

# util.py
from functools import wraps
from django.core.cache import get_cache


cache = get_cache('filesystem')


def cache_lock(lock_id, timeout=cache.get('TIMEOUT', 720)):
    def _decorator(func):
        try:
            timeout_secs = timeout.total_seconds()
        except AttributeError:
            # Timeout is None (forever) or number of seconds.
            timeout_secs = timeout

        acquire_lock = lambda: cache.add(lock_id, 'true', timeout_secs) if timeout_secs else cache.add(lock_id, 'true')
        release_lock = lambda: cache.delete(lock_id)

        @wraps(func)
        def _apply_lock(*args, **kwargs):
            if acquire_lock():
                try:
                    return func(*args, **kwargs)
                finally:
                    release_lock()
            else:
                return False

        return _apply_lock
    return _decorator

基于类的实现

# util.py
from functools import wraps
from django.core.cache import get_cache


cache = get_cache('filesystem')


class cache_lock(object):
    def __init__(self, lock_id, timeout=cache.get('TIMEOUT', 720)):
        self.lock_id = lock_id
        self.timeout = timeout

    def __call__(self, func):
        try:
            timeout_secs = self.timeout.total_seconds()
        except AttributeError:
            # Timeout is None (forever) or number of seconds.
            timeout_secs = self.timeout

        acquire_lock = lambda: cache.add(self.lock_id, 'true', timeout_secs) if timeout_secs else cache.add(self.lock_id, 'true')
        release_lock = lambda: cache.delete(self.lock_id)

        @wraps(func)
        def _apply_lock(*args, **kwargs):
            if acquire_lock():
                try:
                    return func(*args, **kwargs)
                finally:
                    release_lock()
            else:
                return False

        return _apply_lock

测试用例

对于这两种实现,第一种测试方法成功,第二种测试方法失败。

# tests.py
from datetime import timedelta
from celery import task  # using celery.task.task does not help
from django.test import TestCase
from django.test.utils import override_settings
from .util import cache_lock


class UtilTests(TestCase):
    def test_cache_lock_without_celery(self):
        @cache_lock('recursive', timedelta(seconds=1))
        def call_count(i):
            self.assertFalse(call_count(i + 1))
            return i + 1

        self.assertEqual(call_count(0), 1)  # succeeds

    celery_settings = {
        'CELERY_ALWAYS_EAGER': True,
        'CELERY_EAGER_PROPAGATES_EXCEPTIONS': True,
        'DEBUG': True,
    }

    @override_settings(**celery_settings)
    def test_cache_lock_with_celery(self):
        @task(name='test_cache_lock_with_celery')
        @cache_lock('recursive', timedelta(seconds=600))
        def call_count(i):
            self.assertFalse(call_count.apply(i + 1).result)
            return i + 1

        self.assertEqual(call_count.apply(0).result, 1)  # fails!

0 个答案:

没有答案