diff --git a/django_rq/decorators.py b/django_rq/decorators.py index 6e2a8995..4543c95d 100644 --- a/django_rq/decorators.py +++ b/django_rq/decorators.py @@ -1,15 +1,42 @@ from rq.decorators import job as _rq_job -from typing import TYPE_CHECKING, Union +from typing import Any, Callable, Optional, overload, Protocol, TYPE_CHECKING, TypeVar, Union from django.conf import settings from .queues import get_queue if TYPE_CHECKING: + from redis import Redis from rq import Queue + from typing_extensions import ParamSpec + P = ParamSpec('P') + R = TypeVar('R', covariant=True) -def job(func_or_queue, connection=None, *args, **kwargs): + class _JobFn(Protocol[P, R]): + def delay(self, *args: P.args, **kwargs: P.kwargs) -> R: ... + def enqueue(self, *args: P.args, **kwargs: P.kwargs) -> R: ... + def __call__(self, *args: P.args, **kwargs: P.kwargs) -> R: ... + + +@overload +def job(func_or_queue: 'Callable[P, R]') -> '_JobFn[P, R]': ... + +@overload +def job( + func_or_queue: Union['Queue', str], + connection: Optional['Redis'] = None, + *args: Any, + **kwargs: Any, +) -> Callable[['Callable[P, R]'], '_JobFn[P, R]']: ... + + +def job( + func_or_queue: Union['Callable[P, R]', 'Queue', str], + connection: Optional['Redis'] = None, + *args: Any, + **kwargs: Any, +) -> Union['_JobFn[P, R]', Callable[['Callable[P, R]'], '_JobFn[P, R]']]: """ The same as RQ's job decorator, but it automatically works out the ``connection`` argument from RQ_QUEUES.