Added exception propogation for dependencies.

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

Author: s3rius

.flake8

@@ -88,6 +88,8 @@ ignore =
   N802,
   ; Do not perform function calls in argument defaults.
   B008,
+  ; Found except `BaseException`
+  WPS424,
 
   ; all init files
   __init__.py:

taskiq_dependencies/ctx.py

@@ -1,6 +1,7 @@
 import asyncio
 import inspect
 from copy import copy
+from logging import getLogger
 from typing import TYPE_CHECKING, Any, Dict, Generator, List, Optional
 
 from taskiq_dependencies.utils import ParamInfo
@@ -9,6 +10,9 @@
     from taskiq_dependencies.graph import DependencyGraph  # pragma: no cover
 
 
+logger = getLogger("taskiq.dependencies.ctx")
+
+
 class BaseResolveContext:
     """Base resolver context."""
 
@@ -116,18 +120,34 @@ def __enter__(self) -> "SyncResolveContext":
         return self
 
     def __exit__(self, *args: Any) -> None:
-        self.close()
+        self.close(*args)
 
-    def close(self) -> None:
+    def close(self, *args: Any) -> None:  # noqa: C901
         """
         Close all opened dependencies.
 
         This function runs teardown of all dependencies.
+
+        :param args: exception info if any.
         """
+        exception_found = False
+        if args[1] is not None:
+            exception_found = True
         for ctx in self.sub_contexts:
-            ctx.close()
+            ctx.close(*args)
         for dep in reversed(self.opened_dependencies):
             if inspect.isgenerator(dep):
+                if exception_found:
+                    try:
+                        dep.throw(*args)
+                    except BaseException as exc:
+                        logger.warning(
+                            "Exception found on dependency teardown %s",
+                            exc,
+                            exc_info=True,
+                        )
+                        continue
+                    continue
                 for _ in dep:  # noqa: WPS328
                     pass  # noqa: WPS420
 
@@ -201,21 +221,48 @@ async def __aenter__(self) -> "AsyncResolveContext":
         return self
 
     async def __aexit__(self, *args: Any) -> None:
-        await self.close()
+        await self.close(*args)
 
-    async def close(self) -> None:  # noqa: C901
+    async def close(self, *args: Any) -> None:  # noqa: C901
         """
         Close all opened dependencies.
 
         This function runs teardown of all dependencies.
+
+        :param args: exception info if any.
         """
+        exception_found = False
+        if args[1] is not None:
+            exception_found = True
         for ctx in self.sub_contexts:
-            await ctx.close()  # type: ignore
+            await ctx.close(*args)  # type: ignore
         for dep in reversed(self.opened_dependencies):
             if inspect.isgenerator(dep):
+                if exception_found:
+                    try:
+                        dep.throw(*args)
+                    except BaseException as exc:
+                        logger.warning(
+                            "Exception found on dependency teardown %s",
+                            exc,
+                            exc_info=True,
+                        )
+                        continue
+                    continue
                 for _ in dep:  # noqa: WPS328
                     pass  # noqa: WPS420
             elif inspect.isasyncgen(dep):
+                if exception_found:
+                    try:
+                        await dep.athrow(*args)
+                    except BaseException as exc:
+                        logger.warning(
+                            "Exception found on dependency teardown %s",
+                            exc,
+                            exc_info=True,
+                        )
+                        continue
+                    continue
                 async for _ in dep:  # noqa: WPS328
                     pass  # noqa: WPS420
 

tests/test_graph.py

@@ -310,3 +310,93 @@ def target(class_val: str = Depends(TeClass("tval"))) -> None:
 
     info: str = kwargs["class_val"]
     assert info == "tval"
+
+
+def test_exception_generators() -> None:
+
+    errors_found = 0
+
+    def my_generator() -> Generator[int, None, None]:
+        nonlocal errors_found
+        try:
+            yield 1
+        except ValueError:
+            errors_found += 1
+
+    def target(_: int = Depends(my_generator)) -> None:
+        raise ValueError()
+
+    with pytest.raises(ValueError):
+        with DependencyGraph(target=target).sync_ctx() as g:
+            target(**g.resolve_kwargs())
+
+    assert errors_found == 1
+
+
+@pytest.mark.anyio
+async def test_async_exception_generators() -> None:
+
+    errors_found = 0
+
+    async def my_generator() -> AsyncGenerator[int, None]:
+        nonlocal errors_found
+        try:
+            yield 1
+        except ValueError:
+            errors_found += 1
+
+    def target(_: int = Depends(my_generator)) -> None:
+        raise ValueError()
+
+    with pytest.raises(ValueError):
+        async with DependencyGraph(target=target).async_ctx() as g:
+            target(**(await g.resolve_kwargs()))
+
+    assert errors_found == 1
+
+
+@pytest.mark.anyio
+async def test_async_exception_generators_multiple() -> None:
+
+    errors_found = 0
+
+    async def my_generator() -> AsyncGenerator[int, None]:
+        nonlocal errors_found
+        try:
+            yield 1
+        except ValueError:
+            errors_found += 1
+
+    def target(
+        _a: int = Depends(my_generator, use_cache=False),
+        _b: int = Depends(my_generator, use_cache=False),
+        _c: int = Depends(my_generator, use_cache=False),
+    ) -> None:
+        raise ValueError()
+
+    with pytest.raises(ValueError):
+        async with DependencyGraph(target=target).async_ctx() as g:
+            target(**(await g.resolve_kwargs()))
+
+    assert errors_found == 3
+
+
+@pytest.mark.anyio
+async def test_async_exception_generators_no_propogation() -> None:
+
+    errors_found = 0
+
+    async def my_generator() -> AsyncGenerator[int, None]:
+        nonlocal errors_found
+        try:
+            yield 1
+        except ValueError:
+            errors_found += 1
+            raise Exception()
+
+    def target(_: int = Depends(my_generator)) -> None:
+        raise ValueError()
+
+    with pytest.raises(ValueError):
+        async with DependencyGraph(target=target).async_ctx() as g:
+            target(**(await g.resolve_kwargs()))