Added dependencies replacement.

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

Author: s3rius

taskiq_dependencies/ctx.py

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

taskiq_dependencies/graph.py

@@ -20,6 +20,7 @@ class DependencyGraph:
     def __init__(
         self,
         target: Callable[..., Any],
+        replaced_deps: Optional[Dict[Any, Any]] = None,
     ) -> None:
         self.target = target
         # Ordinary dependencies with cache.
@@ -28,6 +29,7 @@ def __init__(
         # Can be considered as sub graphs.
         self.subgraphs: Dict[Any, DependencyGraph] = {}
         self.ordered_deps: List[Dependency] = []
+        self.replaced_deps = replaced_deps
         self._build_graph()
 
     def is_empty(self) -> bool:
@@ -52,14 +54,16 @@ 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.
+        :param replaced_deps: Dependencies to replace during runtime.
         :return: new resolver context.
         """
+        graph = self
+        if replaced_deps:
+            graph = DependencyGraph(self.target, replaced_deps)
         return AsyncResolveContext(
-            graph=self,
-            initial_cache=initial_cache,
-            exception_propagation=exception_propagation,
-            replaced_deps=replaced_deps,
+            graph,
+            initial_cache,
+            exception_propagation,
         )
 
     def sync_ctx(
@@ -76,14 +80,16 @@ 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.
+        :param replaced_deps: Dependencies to replace during runtime.
         :return: new resolver context.
         """
+        graph = self
+        if replaced_deps:
+            graph = DependencyGraph(self.target, replaced_deps)
         return SyncResolveContext(
-            graph=self,
-            initial_cache=initial_cache,
-            exception_propagation=exception_propagation,
-            replaced_deps=replaced_deps,
+            graph,
+            initial_cache,
+            exception_propagation,
         )
 
     def _build_graph(self) -> None:  # noqa: C901, WPS210
@@ -108,6 +114,8 @@ def _build_graph(self) -> None:  # noqa: C901, WPS210
                 continue
             if dep.dependency is None:
                 continue
+            if self.replaced_deps and dep.dependency in self.replaced_deps:
+                dep.dependency = self.replaced_deps[dep.dependency]
             # Get signature and type hints.
             origin = getattr(dep.dependency, "__origin__", None)
             if origin is None:

tests/test_graph.py

@@ -695,14 +695,12 @@ def dep() -> int:
         return 123
 
     def target(val: int = Depends(dep)) -> None:
-        raise ValueError("lol")
+        """Stub function."""
 
     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)
+    async with graph.async_ctx(
+        replaced_deps={dep: replaced},
+        exception_propagation=True,
+    ) as ctx:
+        kwargs = await ctx.resolve_kwargs()
+        assert kwargs["val"] == 321