Added dynamic dependency resolution for unknown tasks (#208)

Author: s3rius

taskiq/brokers/inmemory_broker.py

@@ -1,10 +1,7 @@
 import asyncio
-import inspect
 from collections import OrderedDict
 from concurrent.futures import ThreadPoolExecutor
-from typing import Any, AsyncGenerator, Set, TypeVar, get_type_hints
-
-from taskiq_dependencies import DependencyGraph
+from typing import Any, AsyncGenerator, Set, TypeVar
 
 from taskiq.abc.broker import AsyncBroker
 from taskiq.abc.result_backend import AsyncResultBackend, TaskiqResult
@@ -124,19 +121,6 @@ async def kick(self, message: BrokerMessage) -> None:
         if target_task is None:
             raise TaskiqError("Unknown task.")
 
-        if not self.receiver.dependency_graphs.get(target_task.task_name):
-            self.receiver.dependency_graphs[target_task.task_name] = DependencyGraph(
-                target_task.original_func,
-            )
-        if not self.receiver.task_signatures.get(target_task.task_name):
-            self.receiver.task_signatures[target_task.task_name] = inspect.signature(
-                target_task.original_func,
-            )
-        if not self.receiver.task_hints.get(target_task.task_name):
-            self.receiver.task_hints[target_task.task_name] = get_type_hints(
-                target_task.original_func,
-            )
-
         task = asyncio.create_task(self.receiver.callback(message=message.message))
         self._running_tasks.add(task)
         task.add_done_callback(self._running_tasks.discard)

taskiq/receiver/receiver.py

@@ -64,10 +64,9 @@ def __init__(  # noqa: WPS211
         self.dependency_graphs: Dict[str, DependencyGraph] = {}
         self.propagate_exceptions = propagate_exceptions
         self.on_exit = on_exit
+        self.known_tasks: Set[str] = set()
         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)
+            self._prepare_task(task.task_name, task.original_func)
         self.sem: "Optional[asyncio.Semaphore]" = None
         if max_async_tasks is not None and max_async_tasks > 0:
             self.sem = asyncio.Semaphore(max_async_tasks)
@@ -163,7 +162,7 @@ async def callback(  # noqa: C901, WPS213, WPS217
             if raise_err:
                 raise exc
 
-    async def run_task(  # noqa: C901, WPS210
+    async def run_task(  # noqa: C901, WPS210, WPS213
         self,
         target: Callable[..., Any],
         message: TaskiqMessage,
@@ -190,6 +189,8 @@ async def run_task(  # noqa: C901, WPS210
         returned = None
         found_exception: "Optional[BaseException]" = None
         signature = None
+        if message.task_name not in self.known_tasks:
+            self._prepare_task(message.task_name, target)
         if self.validate_params:
             signature = self.task_signatures.get(message.task_name)
         dependency_graph = self.dependency_graphs.get(message.task_name)
@@ -382,3 +383,23 @@ def task_cb(task: "asyncio.Task[Any]") -> None:
             # and this behaviour considered to be a Hisenbug.
             # https://textual.textualize.io/blog/2023/02/11/the-heisenbug-lurking-in-your-async-code/
             task.add_done_callback(task_cb)
+
+    def _prepare_task(self, name: str, handler: Callable[..., Any]) -> None:
+        """
+        Prepare task for execution.
+
+        This function gets function's signature,
+        type hints and builds dependency graph.
+
+        It's useful for dynamic dependency resolution,
+        because sometimes the receiver can get
+        funcion that is defined in runtime. We need
+        to be aware of that.
+
+        :param name: task name.
+        :param handler: task handler.
+        """
+        self.known_tasks.add(name)
+        self.task_signatures[name] = inspect.signature(handler)
+        self.task_hints[name] = get_type_hints(handler)
+        self.dependency_graphs[name] = DependencyGraph(handler)

tests/receiver/test_receiver.py

@@ -1,4 +1,5 @@
 import asyncio
+import random
 import time
 from concurrent.futures import ThreadPoolExecutor
 from typing import Any, AsyncGenerator, List, Optional, TypeVar
@@ -227,6 +228,41 @@ async def my_task() -> int:
     assert called_times == 1
 
 
+@pytest.mark.anyio
+async def test_callback_no_dep_info() -> None:
+    """Test that callback function works well."""
+    broker = InMemoryBroker()
+    expected = random.randint(1, 100)
+    ret_val = None
+
+    def dependency() -> int:
+        return expected
+
+    @broker.task
+    async def my_task(dep: int = Depends(dependency)) -> None:
+        nonlocal ret_val
+        ret_val = dep
+
+    receiver = get_receiver(broker)
+    receiver.known_tasks.remove(my_task.task_name)
+    receiver.dependency_graphs.pop(my_task.task_name, None)
+    receiver.task_signatures.pop(my_task.task_name, None)
+    receiver.task_hints.pop(my_task.task_name, None)
+
+    broker_message = broker.formatter.dumps(
+        TaskiqMessage(
+            task_id="task_id",
+            task_name=my_task.task_name,
+            labels={},
+            args=[],
+            kwargs={},
+        ),
+    )
+
+    await receiver.callback(broker_message.message)
+    assert ret_val == expected
+
+
 @pytest.mark.anyio
 async def test_callback_success_ackable() -> None:
     """Test that acking works."""