diff --git a/celery_pubsub/pubsub.py b/celery_pubsub/pubsub.py index 2cd45c9..768a83e 100644 --- a/celery_pubsub/pubsub.py +++ b/celery_pubsub/pubsub.py @@ -33,7 +33,9 @@ P: TypeAlias = typing.Any # ParamSpec R: TypeAlias = typing.Any # Return type -task: typing.Callable[..., typing.Callable[[typing.Callable[[P], R]], Task[P, R]]] = celery.shared_task +task: typing.Callable[ + ..., typing.Callable[[typing.Callable[[P], R]], Task[P, R]] +] = celery.shared_task class PubSubManager: @@ -87,9 +89,12 @@ def _topic_to_re(topic: str) -> re.Pattern[str]: def subscribe_to(topic: str) -> typing.Callable[[typing.Callable[[P], R]], Task[P, R]]: def decorator(func: typing.Callable[[P], R]) -> Task[P, R]: - app_name, module_name = func.__module__.split(".", 1) - task_name = f"{app_name}.{module_name}.{func.__qualname__}" - task_instance = func if isinstance(func, Task) else task(name=task_name)(func) + if isinstance(func, Task): + task_instance: Task[P, R] = func + else: + app_name, module_name = func.__module__.split(".", 1) + task_name = f"{app_name}.{module_name}.{func.__qualname__}" + task_instance = task(name=task_name)(func) _pubsub_manager.subscribe(topic, task_instance) return task_instance