Added ability to add custom first-level dependencies (#70)

Author: s3rius

docs/guide/state-and-deps.md

@@ -171,3 +171,13 @@ By default taskiq has only two deendencies:
 
 - Context from `taskiq.context.Context`
 - TaskiqState from `taskiq.state.TaskiqState`
+
+
+### Adding first-level dependencies
+
+You can expand default list of available dependencies for you application.
+Taskiq have an ability to add new first-level dependencies using brokers.
+
+The AsyncBroker interface has a function called `add_dependency_context` and you can add
+more default dependencies to the taskiq. This may be useful for libraries if you want to
+add new dependencies to users.

taskiq/abc/broker.py

@@ -86,11 +86,23 @@ def __init__(
         # 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.
-        self.event_handlers: DefaultDict[  # noqa: WPS234
+        self.event_handlers: DefaultDict[
             TaskiqEvents,
             List[Callable[[TaskiqState], Optional[Awaitable[None]]]],
         ] = defaultdict(list)
         self.state = TaskiqState()
+        self.custom_dependency_context: Dict[Any, Any] = {}
+
+    def add_dependency_context(self, new_ctx: Dict[Any, Any]) -> None:
+        """
+        Add first-level dependencies.
+
+        Provided dict will be used to inject new dependencies
+        in all dependency graph contexts.
+
+        :param new_ctx: Additional context values for dependnecy injection.
+        """
+        self.custom_dependency_context.update(new_ctx)
 
     def add_middlewares(self, *middlewares: "TaskiqMiddleware") -> None:
         """

taskiq/cli/worker/receiver.py

@@ -159,12 +159,14 @@ async def run_task(  # noqa: C901, WPS210
         dep_ctx = None
         if dependency_graph:
             # Create a context for dependency resolving.
-            dep_ctx = dependency_graph.async_ctx(
+            broker_ctx = self.broker.custom_dependency_context
+            broker_ctx.update(
                 {
                     Context: Context(message, self.broker),
                     TaskiqState: self.broker.state,
                 },
             )
+            dep_ctx = dependency_graph.async_ctx(broker_ctx)
             # Resolve all function's dependencies.
             dep_kwargs = await dep_ctx.resolve_kwargs()
             for key, val in dep_kwargs.items():

tests/cli/worker/test_custom_contexts.py

@@ -1,6 +1,7 @@
 from typing import Any, Optional
 
 import pytest
+from taskiq_dependencies import Depends
 
 from taskiq.abc.broker import AsyncBroker
 from taskiq.abc.middleware import TaskiqMiddleware
@@ -200,3 +201,43 @@ async def test_callback_unknown_task() -> None:
     )
 
     await receiver.callback(broker_message)
+
+
+@pytest.mark.anyio
+async def test_custom_ctx() -> None:
+    """Tests that run_task can run sync tasks."""
+
+    class MyTestClass:
+        """Class to test injection."""
+
+        def __init__(self, val: int) -> None:
+            self.val = val
+
+    broker = InMemoryBroker()
+
+    # We register a task into broker,
+    # to build dependency graph on startup.
+    @broker.task
+    def test_func(tes_val: MyTestClass = Depends()) -> int:
+        return tes_val.val
+
+    # We add custom first-level dependency.
+    broker.add_dependency_context({MyTestClass: MyTestClass(11)})
+    # Create a receiver.
+    receiver = get_receiver(broker)
+
+    result = await receiver.run_task(
+        test_func,
+        TaskiqMessage(
+            task_id="",
+            task_name=test_func.task_name,
+            labels={},
+            args=[],
+            kwargs={},
+        ),
+    )
+
+    # Check that the value is equal
+    # to the one we supplied.
+    assert result.return_value == 11
+    assert not result.is_err