Added local task registry, improved shared_tasks API. (#203)

Author: s3rius

taskiq/abc/broker.py

@@ -11,6 +11,7 @@
     AsyncGenerator,
     Awaitable,
     Callable,
+    ClassVar,
     DefaultDict,
     Dict,
     List,
@@ -68,7 +69,7 @@ class AsyncBroker(ABC):
     in async mode.
     """
 
-    available_tasks: Dict[str, AsyncTaskiqDecoratedTask[Any, Any]] = {}
+    global_task_registry: ClassVar[Dict[str, AsyncTaskiqDecoratedTask[Any, Any]]] = {}
 
     def __init__(
         self,
@@ -98,6 +99,7 @@ def __init__(
         self.decorator_class = AsyncTaskiqDecoratedTask
         self.formatter: "TaskiqFormatter" = JSONFormatter()
         self.id_generator = task_id_generator
+        self.local_task_registry: Dict[str, AsyncTaskiqDecoratedTask[Any, Any]] = {}
         # Every event has a list of handlers.
         # Every handler is a function which takes state as a first argument.
         # And handler can be either sync or async.
@@ -112,6 +114,41 @@ def __init__(
         # True only if broker runs in scheduler process.
         self.is_scheduler_process: bool = False
 
+    def find_task(self, task_name: str) -> Optional[AsyncTaskiqDecoratedTask[Any, Any]]:
+        """
+        Returns task by name.
+
+        This method should be used to get task by name.
+        Instead of accessing `available_tasks` or `local_available_tasks` directly.
+
+        It searches task by name in dict of tasks that
+        were registered for this broker directly.
+        If it fails, it checks global dict of all available tasks.
+
+        :param task_name: name of a task.
+        :returns: found task or None.
+        """
+        return self.local_task_registry.get(
+            task_name,
+        ) or self.global_task_registry.get(
+            task_name,
+        )
+
+    def get_all_tasks(self) -> Dict[str, AsyncTaskiqDecoratedTask[Any, Any]]:
+        """
+        Method to fetch all tasks available in broker.
+
+        This method returns all tasks, globally and locally
+        available in broker. With local tasks having higher priority.
+
+        So, if you have two tasks with the same name,
+        one registered in global registry and one registered
+        in local registry, then local task will be returned.
+
+        :return: dict of all tasks. Keys are task names, values are tasks.
+        """
+        return {**self.global_task_registry, **self.local_task_registry}
+
     def add_dependency_context(self, new_ctx: Dict[Any, Any]) -> None:
         """
         Add first-level dependencies.
@@ -291,7 +328,7 @@ def inner(
                     ),
                 )
 
-                self.available_tasks[decorated_task.task_name] = decorated_task
+                self._register_task(decorated_task.task_name, decorated_task)
 
                 return decorated_task
 
@@ -416,3 +453,19 @@ def with_event_handlers(
         """
         self.event_handlers[event].extend(handlers)
         return self
+
+    def _register_task(
+        self,
+        task_name: str,
+        task: AsyncTaskiqDecoratedTask[Any, Any],
+    ) -> None:
+        """
+        Mehtod is used to register tasks.
+
+        By default we register tasks in local task registry.
+        But this behaviour can be changed in subclasses.
+
+        :param task_name: Name of a task.
+        :param task: Decorated task.
+        """
+        self.local_task_registry[task_name] = task

taskiq/brokers/inmemory_broker.py

@@ -120,7 +120,7 @@ async def kick(self, message: BrokerMessage) -> None:
 
         :raises TaskiqError: if someone wants to kick unknown task.
         """
-        target_task = self.available_tasks.get(message.task_name)
+        target_task = self.find_task(message.task_name)
         if target_task is None:
             raise TaskiqError("Unknown task.")
 

taskiq/brokers/shared_broker.py

@@ -1,4 +1,4 @@
-from typing import AsyncGenerator, Optional, TypeVar
+from typing import Any, AsyncGenerator, Optional, TypeVar
 
 from typing_extensions import ParamSpec
 
@@ -71,5 +71,13 @@ async def listen(self) -> AsyncGenerator[bytes, None]:  # type: ignore
         """
         raise TaskiqError("Shared broker cannot listen")
 
+    def _register_task(
+        self,
+        task_name: str,
+        task: AsyncTaskiqDecoratedTask[Any, Any],
+    ) -> None:
+        self.global_task_registry[task_name] = task
+
 
 async_shared_broker = AsyncSharedBroker()
+shared_task = async_shared_broker.task

taskiq/cli/worker/run.py

@@ -73,10 +73,9 @@ def start_listen(args: WorkerArgs, event: Event) -> None:  # noqa: WPS210, WPS21
     This function starts actual listening process.
 
     It imports broker and all tasks.
-    Since tasks registers themselves in a global set,
-    it's easy to just import module where you have decorated
-    function and they will be available in broker's `available_tasks`
-    field.
+    Since tasks auto registeres themselves in a broker,
+    we don't need to do anything else other than importing.
+
 
     :param args: CLI arguments.
     :param event: Event for notification.

taskiq/receiver/receiver.py

@@ -60,7 +60,7 @@ def __init__(  # noqa: WPS211
         self.task_hints: Dict[str, Dict[str, Any]] = {}
         self.dependency_graphs: Dict[str, DependencyGraph] = {}
         self.propagate_exceptions = propagate_exceptions
-        for task in self.broker.available_tasks.values():
+        for task in self.broker.get_all_tasks().values():
             self.task_signatures[task.task_name] = inspect.signature(task.original_func)
             self.task_hints[task.task_name] = get_type_hints(task.original_func)
             self.dependency_graphs[task.task_name] = DependencyGraph(task.original_func)
@@ -106,7 +106,8 @@ async def callback(  # noqa: C901, WPS213, WPS217
             )
             return
         logger.debug(f"Received message: {taskiq_msg}")
-        if taskiq_msg.task_name not in self.broker.available_tasks:
+        task = self.broker.find_task(taskiq_msg.task_name)
+        if task is None:
             logger.warning(
                 'task "%s" is not found. Maybe you forgot to import it?',
                 taskiq_msg.task_name,
@@ -135,7 +136,7 @@ async def callback(  # noqa: C901, WPS213, WPS217
             await maybe_awaitable(message.ack())
 
         result = await self.run_task(
-            target=self.broker.available_tasks[taskiq_msg.task_name].original_func,
+            target=task.original_func,
             message=taskiq_msg,
         )
 

taskiq/schedule_sources/label_based.py

@@ -27,7 +27,7 @@ async def get_schedules(self) -> List["ScheduledTask"]:
         :return: list of schedules.
         """
         schedules = []
-        for task_name, task in self.broker.available_tasks.items():
+        for task_name, task in self.broker.get_all_tasks().items():
             if task.broker != self.broker:
                 continue
             for schedule in task.labels.get("schedule", []):
@@ -61,7 +61,7 @@ def post_send(self, scheduled_task: ScheduledTask) -> None:
         if scheduled_task.cron or not scheduled_task.time:
             return  # it's scheduled task with cron label, do not remove this trigger.
 
-        for task_name, task in self.broker.available_tasks.items():
+        for task_name, task in self.broker.get_all_tasks().items():
             if task.broker != self.broker or scheduled_task.task_name != task_name:
                 continue
 

tests/conftest.py

@@ -27,7 +27,7 @@ def reset_broker() -> Generator[None, None, None]:
     broker variables to default state.
     """
     yield
-    AsyncBroker.available_tasks = {}
+    AsyncBroker.global_task_registry = {}
     AsyncBroker.is_worker_process = False
     AsyncBroker.is_scheduler_process = False
 

tox.ini

@@ -12,5 +12,4 @@ allowlist_externals = poetry
 commands_pre =
     poetry install
 commands =
-    pre-commit run --all-files
-    poetry run pytest -vv
+    poetry run pytest -vv -n auto