All imported tasks are now available for brokers.

Signed-off-by: Pavel Kirilin <[email protected]>

Author: s3rius

taskiq/abc/broker.py

@@ -276,6 +276,8 @@ class AsyncBroker(ABC):
     in async mode.
     """
 
+    available_tasks: Dict[str, AsyncTaskiqDecoratedTask[Any, Any]] = {}
+
     def __init__(
         self,
         result_backend: Optional[AsyncResultBackend[_T]] = None,
@@ -284,7 +286,6 @@ def __init__(
             result_backend = DummyResultBackend()
         self.result_backend = result_backend
         self.is_worker_process = False
-        self._related_tasks: Set[AsyncTaskiqDecoratedTask[..., Any]] = set()
 
     async def startup(self) -> None:
         """Do something when starting broker."""
@@ -395,10 +396,7 @@ def inner(
                     ),
                 )
 
-                # Adds this task to the list of tasks.
-                # This option is used by workers.
-                if self.is_worker_process:
-                    self._related_tasks.add(decorated_task)  # type: ignore
+                self.available_tasks[decorated_task.task_name] = decorated_task
 
                 return decorated_task
 

taskiq/brokers/inmemory_broker.py

@@ -6,6 +6,7 @@
 from taskiq.abc.broker import AsyncBroker
 from taskiq.abc.result_backend import AsyncResultBackend, TaskiqResult
 from taskiq.cli.async_task_runner import run_task
+from taskiq.exceptions import TaskiqError
 from taskiq.message import TaskiqMessage
 
 _ReturnType = TypeVar("_ReturnType")
@@ -87,18 +88,15 @@ def __init__(
         sync_tasks_pool_size: int = 4,
         logs_format: Optional[str] = None,
         max_stored_results: int = 100,
+        cast_types: bool = True,
     ) -> None:
         super().__init__(
             InmemoryResultBackend(
                 max_stored_results=max_stored_results,
             ),
         )
-        # We mock as if it's a worker process.
-        # So every task call will add tasks in
-        # _related_tasks attribute.
-        self.is_worker_process = True
-        self.tasks_mapping = None
         self.executor = ThreadPoolExecutor(max_workers=sync_tasks_pool_size)
+        self.cast_types = cast_types
         if logs_format is None:
             logs_format = "%(levelname)s %(message)s"
         self.logs_format = logs_format
@@ -110,13 +108,11 @@ async def kick(self, message: TaskiqMessage) -> None:
         This method just executes given task.
 
         :param message: incomming message.
-        :raises ValueError: if someone wants to kick unknown task.
+        :raises TaskiqError: if someone wants to kick unknown task.
         """
-        for task in self._related_tasks:
-            if task.task_name == message.task_name:
-                target_task = task
+        target_task = self.available_tasks.get(message.task_name)
         if target_task is None:
-            raise ValueError("Unknown task.")
+            raise TaskiqError("Unknown task.")
         result = await run_task(
             target=target_task.original_func,
             signature=inspect.signature(target_task.original_func),

taskiq/cli/async_task_runner.py

@@ -249,22 +249,19 @@ async def async_listen_messages(  # noqa: C901, WPS210, WPS213
         max_workers=cli_args.max_threadpool_threads,
     )
     logger.info("Listening started.")
-    task_registry: Dict[str, Callable[..., Any]] = {}
     task_signatures: Dict[str, inspect.Signature] = {}
-    for task in broker._related_tasks:  # noqa: WPS437
-        task_registry[task.task_name] = task.original_func
-        # If we need to parse parameters we remember all tasks signatures.
+    for task in broker.available_tasks.values():
         if not cli_args.no_parse:
             task_signatures[task.task_name] = inspect.signature(task.original_func)
     async for message in broker.listen():
         logger.debug(f"Received message: {message}")
-        if message.task_name not in task_registry:
+        if message.task_name not in broker.available_tasks:
             logger.warning(
                 'task "%s" is not found. Maybe you forgot to import it?',
                 message.task_name,
             )
             continue
-        func = task_registry[message.task_name]
+        func = broker.available_tasks[message.task_name]
         logger.debug(
             "Function for task %s is resolved. Executing...",
             message.task_name,

taskiq/cli/worker.py

@@ -92,9 +92,9 @@ def start_listen(args: TaskiqArgs) -> None:
     This function starts actual listening process.
 
     It imports broker and all tasks.
-    Since tasks registers themselfs in broker
+    Since tasks registers themselfs in a global set,
     it's easy to just import module where you have decorated
-    function and they will be available in broker's `_related_tasks`
+    function and they will be available in broker's `available_tasks`
     field.
 
     :param args: CLI arguments.