Added dependencies replacemenets.

Signed-off-by: Pavel Kirilin <win10@list.ru>

Author: s3rius

taskiq_dependencies/ctx.py

@@ -20,12 +20,14 @@ def __init__(
         self,
         graph: "DependencyGraph",
         initial_cache: Optional[Dict[Any, Any]] = None,
+        replaced_deps: Optional[Dict[Any, Any]] = None,
         exception_propagation: bool = True,
     ) -> None:
         self.graph = graph
         self.opened_dependencies: List[Any] = []
         self.sub_contexts: "List[Any]" = []
         self.initial_cache = initial_cache or {}
+        self.replaced_funcs = replaced_deps or {}
         self.propagate_excs = exception_propagation
 
     def traverse_deps(  # noqa: C901, WPS210
@@ -56,7 +58,7 @@ def traverse_deps(  # noqa: C901, WPS210
             # later.
             if not dep.use_cache:
                 continue
-            # If somehow we have dependency with unknwon function.
+            # If somehow we have dependency with unknown function.
             if dep.dependency is None:
                 continue
             # If dependency is already calculated.
@@ -89,7 +91,13 @@ def traverse_deps(  # noqa: C901, WPS210
                         continue
                     if subdep.kwargs:
                         resolved_kwargs.update(subdep.kwargs)
-                    kwargs[subdep.param_name] = yield subdep.dependency(
+                    # We try to grab possible replacement for
+                    # function if any. Otherwise, original subdependency is returned.
+                    target_dependency = self.replaced_funcs.get(
+                        subdep.dependency,
+                        subdep.dependency,
+                    )
+                    kwargs[subdep.param_name] = yield target_dependency(
                         **resolved_kwargs,
                     )
 
@@ -103,7 +111,13 @@ def traverse_deps(  # noqa: C901, WPS210
             ):
                 user_kwargs = dep.kwargs
                 user_kwargs.update(kwargs)
-                cache[dep.dependency] = yield dep.dependency(**user_kwargs)
+                # From dict of replaced functions,
+                # we grab possible replacement or original function.
+                target_dependency = self.replaced_funcs.get(
+                    dep.dependency,
+                    dep.dependency,
+                )
+                cache[dep.dependency] = yield target_dependency(**user_kwargs)
         return kwargs
 
 
@@ -169,7 +183,12 @@ def resolver(self, executed_func: Any, initial_cache: Dict[Any, Any]) -> Any:
         :return: dict with resolved kwargs.
         """
         if getattr(executed_func, "dep_graph", False):
-            ctx = SyncResolveContext(executed_func, initial_cache)
+            ctx = SyncResolveContext(
+                graph=executed_func,
+                initial_cache=initial_cache,
+                replaced_deps=self.replaced_funcs,
+                exception_propagation=self.propagate_excs,
+            )
             self.sub_contexts.append(ctx)
             sub_result = ctx.resolve_kwargs()
         elif inspect.isgenerator(executed_func):
@@ -286,7 +305,12 @@ async def resolver(self, executed_func: Any, initial_cache: Dict[Any, Any]) -> A
         :return: dict with resolved kwargs.
         """
         if getattr(executed_func, "dep_graph", False):
-            ctx = AsyncResolveContext(executed_func, initial_cache)  # type: ignore
+            ctx = AsyncResolveContext(
+                graph=executed_func,
+                initial_cache=initial_cache,
+                replaced_deps=self.replaced_funcs,
+                exception_propagation=self.propagate_excs,
+            )  # type: ignore
             self.sub_contexts.append(ctx)
             sub_result = await ctx.resolve_kwargs()
         elif inspect.isgenerator(executed_func):

taskiq_dependencies/graph.py

@@ -41,6 +41,7 @@ def is_empty(self) -> bool:
     def async_ctx(
         self,
         initial_cache: Optional[Dict[Any, Any]] = None,
+        replaced_deps: Optional[Dict[Any, Any]] = None,
         exception_propagation: bool = True,
     ) -> AsyncResolveContext:
         """
@@ -51,17 +52,20 @@ def async_ctx(
         :param initial_cache: initial cache dict.
         :param exception_propagation: If true, all found errors within
             context will be propagated to dependencies.
+        :param replaced_deps: dict with dependencies to replace.
         :return: new resolver context.
         """
         return AsyncResolveContext(
-            self,
-            initial_cache,
-            exception_propagation,
+            graph=self,
+            initial_cache=initial_cache,
+            exception_propagation=exception_propagation,
+            replaced_deps=replaced_deps,
         )
 
     def sync_ctx(
         self,
         initial_cache: Optional[Dict[Any, Any]] = None,
+        replaced_deps: Optional[Dict[Any, Any]] = None,
         exception_propagation: bool = True,
     ) -> SyncResolveContext:
         """
@@ -72,12 +76,14 @@ def sync_ctx(
         :param initial_cache: initial cache dict.
         :param exception_propagation: If true, all found errors within
             context will be propagated to dependencies.
+        :param replaced_deps: dict with dependencies to replace.
         :return: new resolver context.
         """
         return SyncResolveContext(
-            self,
-            initial_cache,
-            exception_propagation,
+            graph=self,
+            initial_cache=initial_cache,
+            exception_propagation=exception_propagation,
+            replaced_deps=replaced_deps,
         )
 
     def _build_graph(self) -> None:  # noqa: C901, WPS210

tests/test_graph.py

@@ -1,4 +1,5 @@
 import asyncio
+import re
 import uuid
 from typing import Any, AsyncGenerator, Generator, Generic, Tuple, TypeVar
 
@@ -36,8 +37,9 @@ def testfunc(a: int = Depends(dep1)) -> int:
         return a
 
     with DependencyGraph(testfunc).sync_ctx({}) as sctx:
-        with pytest.raises(RuntimeError):
-            assert sctx.resolve_kwargs() == {"a": 1}
+        with pytest.warns(match=re.compile(".*was never awaited.*")):
+            with pytest.raises(RuntimeError):
+                assert sctx.resolve_kwargs() == {"a": 1}
 
     async with DependencyGraph(testfunc).async_ctx({}) as actx:
         assert await actx.resolve_kwargs() == {"a": 1}
@@ -611,3 +613,96 @@ def target(
         assert dep_obj.dependency == GenericClass[Tuple[str, int]]
         assert dep_obj.signature.name == "class_val"
         assert dep_obj.signature.annotation == GenericClass[Tuple[str, int]]
+
+
+@pytest.mark.anyio
+async def test_replaced_dep_simple() -> None:
+    def replaced() -> int:
+        return 321
+
+    def dep() -> int:
+        return 123
+
+    def target(val: int = Depends(dep)) -> None:
+        return None
+
+    graph = DependencyGraph(target=target)
+    async with graph.async_ctx(replaced_deps={dep: replaced}) as ctx:
+        kwargs = await ctx.resolve_kwargs()
+    assert kwargs["val"] == 321
+
+
+@pytest.mark.anyio
+async def test_replaced_dep_generators() -> None:
+    call_count = 0
+
+    def replaced() -> Generator[int, None, None]:
+        nonlocal call_count
+        yield 321
+        call_count += 1
+
+    def dep() -> int:
+        return 123
+
+    def target(val: int = Depends(dep)) -> None:
+        return None
+
+    graph = DependencyGraph(target=target)
+    async with graph.async_ctx(replaced_deps={dep: replaced}) as ctx:
+        kwargs = await ctx.resolve_kwargs()
+    assert kwargs["val"] == 321
+    assert call_count == 1
+
+
+@pytest.mark.anyio
+async def test_replaced_dep_exception_propogation() -> None:
+    exc_count = 0
+
+    def replaced() -> Generator[int, None, None]:
+        nonlocal exc_count
+        try:
+            yield 321
+        except ValueError:
+            exc_count += 1
+
+    def dep() -> int:
+        return 123
+
+    def target(val: int = Depends(dep)) -> None:
+        raise ValueError("lol")
+
+    graph = DependencyGraph(target=target)
+    with pytest.raises(ValueError):
+        async with graph.async_ctx(
+            replaced_deps={dep: replaced},
+            exception_propagation=True,
+        ) as ctx:
+            kwargs = await ctx.resolve_kwargs()
+            assert kwargs["val"] == 321
+            target(**kwargs)
+    assert exc_count == 1
+
+
+@pytest.mark.anyio
+async def test_replaced_dep_subdependencies() -> None:
+    def subdep() -> int:
+        return 321
+
+    def replaced(ret_val: int = Depends(subdep)) -> int:
+        return ret_val
+
+    def dep() -> int:
+        return 123
+
+    def target(val: int = Depends(dep)) -> None:
+        raise ValueError("lol")
+
+    graph = DependencyGraph(target=target)
+    with pytest.raises(ValueError):
+        async with graph.async_ctx(
+            replaced_deps={dep: replaced},
+            exception_propagation=True,
+        ) as ctx:
+            kwargs = await ctx.resolve_kwargs()
+            assert kwargs["val"] == 321
+            target(**kwargs)