chore: clean code

Mulugruntz · Dec 4, 2023 · 637622c · 637622c
1 parent d8b703c
commit 637622c
Showing 1 changed file with 9 additions and 4 deletions.
diff --git 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