feat: custom reciever (#123)

Co-authored-by: Anton <[email protected]>

Author: Sobes76rus

taskiq/cli/worker/args.py

@@ -24,6 +24,7 @@ class WorkerArgs:
     reload: bool = False
     no_gitignore: bool = False
     max_async_tasks: int = 100
+    receiver: str = "taskiq.receiver:Receiver"
 
     @classmethod
     def from_cli(  # noqa: WPS213
@@ -45,6 +46,15 @@ def from_cli(  # noqa: WPS213
                 "'module.module:variable' format."
             ),
         )
+        parser.add_argument(
+            "--receiver",
+            default="taskiq.receiver:Receiver",
+            help=(
+                "Where to search for receiver. "
+                "This string must be specified in "
+                "'module.module:variable' format."
+            ),
+        )
         parser.add_argument(
             "--tasks-pattern",
             "-tp",

taskiq/cli/worker/run.py

@@ -2,7 +2,7 @@
 import logging
 import signal
 from concurrent.futures import ThreadPoolExecutor
-from typing import Any
+from typing import Any, Type
 
 from taskiq.abc.broker import AsyncBroker
 from taskiq.cli.utils import import_object, import_tasks
@@ -51,7 +51,21 @@ async def shutdown_broker(broker: AsyncBroker, timeout: float) -> None:
         )
 
 
-def start_listen(args: WorkerArgs) -> None:  # noqa: WPS213
+def get_receiver_type(args: WorkerArgs) -> Type[Receiver]:
+    """
+    Import Receiver from args.
+
+    :param args: CLI arguments.
+    :raises ValueError: if receiver is not a Receiver type.
+    :return: Receiver type.
+    """
+    receiver_type = import_object(args.receiver)
+    if not (isinstance(receiver_type, type) and issubclass(receiver_type, Receiver)):
+        raise ValueError("Unknown receiver type. Please use Receiver class.")
+    return receiver_type
+
+
+def start_listen(args: WorkerArgs) -> None:  # noqa: WPS210, WPS213
     """
     This function starts actual listening process.
 
@@ -63,6 +77,7 @@ def start_listen(args: WorkerArgs) -> None:  # noqa: WPS213
 
     :param args: CLI arguments.
     :raises ValueError: if broker is not an AsyncBroker instance.
+    :raises ValueError: if receiver is not a Receiver type.
     """
     if uvloop is not None:
         logger.debug("UVLOOP found. Installing policy.")
@@ -77,6 +92,8 @@ def start_listen(args: WorkerArgs) -> None:  # noqa: WPS213
     if not isinstance(broker, AsyncBroker):
         raise ValueError("Unknown broker type. Please use AsyncBroker instance.")
 
+    receiver_type = get_receiver_type(args)
+
     # Here how we manage interruptions.
     # We have to remember shutting_down state,
     # because KeyboardInterrupt can be send multiple
@@ -105,10 +122,11 @@ def interrupt_handler(signum: int, _frame: Any) -> None:
     signal.signal(signal.SIGTERM, interrupt_handler)
 
     loop = asyncio.get_event_loop()
+
     try:
         logger.debug("Initialize receiver.")
         with ThreadPoolExecutor(args.max_threadpool_threads) as pool:
-            receiver = Receiver(
+            receiver = receiver_type(
                 broker=broker,
                 executor=pool,
                 validate_params=not args.no_parse,