Repository: reagento/dishka
Pull request #115 (2024-03-28 16:42:49)

Implement pass exception to generators #113
All discussion and info in #113 

Waiting for review, linter checks and tests passed.


Patch:
diff --git a/docs/provider/provide.rst b/docs/provider/provide.rst
index dfa1800..678c256 100644
--- a/docs/provider/provide.rst
+++ b/docs/provider/provide.rst
@@ -32,6 +32,20 @@ By default the result is cached within scope. You can disable it providing ``cac
             yield a
             a.close()
 
+Also, if an error occurs during process handling (inside the ``with`` block), it will be sent to the generator:
+
+.. code-block:: python
+
+    class MyProvider(Provider):
+        @provide(scope=Scope.REQUEST)
+        def get_a(self) -> Iterable[A]:
+            a = A()
+            exc = yield a
+            # exc will be None if an exception has not occurred
+            if exc:
+                print("Some exception while process handling: ", exc)
+            a.close()
+
 * Do not have any specific logic and just want to create class using its ``__init__``? Then add a provider attribute using ``provide`` as function passing that class.
 
 .. code-block:: python
diff --git a/src/dishka/async_container.py b/src/dishka/async_container.py
index 9e2af07..c832ba1 100644
--- a/src/dishka/async_container.py
+++ b/src/dishka/async_container.py
@@ -126,14 +126,14 @@ class AsyncContainer:
             e.add_path(self.registry.get_factory(key))
             raise
 
-    async def close(self):
+    async def close(self, exception: Exception | None = None):
         errors = []
         for exit_generator in self._exits[::-1]:
             try:
                 if exit_generator.type is FactoryType.ASYNC_GENERATOR:
-                    await anext(exit_generator.callable)
+                    await exit_generator.callable.asend(exception)
                 elif exit_generator.type is FactoryType.GENERATOR:
-                    next(exit_generator.callable)
+                    exit_generator.callable.send(exception)
             except StopIteration:  # noqa: PERF203
                 pass
             except StopAsyncIteration:
@@ -142,7 +142,7 @@ class AsyncContainer:
                 errors.append(err)
         if self.close_parent:
             try:
-                await self.parent_container.close()
+                await self.parent_container.close(exception)
             except Exception as err:  # noqa: BLE001
                 errors.append(err)
         if errors:
@@ -157,7 +157,7 @@ class AsyncContextWrapper:
         return self.container
 
     async def __aexit__(self, exc_type, exc_val, exc_tb):
-        await self.container.close()
+        await self.container.close(exception=exc_val)
 
 
 def make_async_container(
diff --git a/src/dishka/container.py b/src/dishka/container.py
index 2952673..000c1fa 100644
--- a/src/dishka/container.py
+++ b/src/dishka/container.py
@@ -123,19 +123,19 @@ class Container:
             e.add_path(self.registry.get_factory(key))
             raise
 
-    def close(self) -> None:
+    def close(self, exception: Exception | None = None) -> None:
         errors = []
         for exit_generator in self._exits[::-1]:
             try:
                 if exit_generator.type is FactoryType.GENERATOR:
-                    next(exit_generator.callable)
+                    exit_generator.callable.send(exception)
             except StopIteration:  # noqa: PERF203
                 pass
             except Exception as err:  # noqa: BLE001
                 errors.append(err)
         if self.close_parent:
             try:
-                self.parent_container.close()
+                self.parent_container.close(exception)
             except Exception as err:  # noqa: BLE001
                 errors.append(err)
 
@@ -153,7 +153,7 @@ class ContextWrapper:
         return self.container
 
     def __exit__(self, exc_type, exc_val, exc_tb):
-        self.container.close()
+        self.container.close(exception=exc_val)
 
 
 def make_container(


Test patch:
diff --git a/tests/unit/container/test_enter_exit.py b/tests/unit/container/test_enter_exit.py
index ce78b28..b28d1ab 100644
--- a/tests/unit/container/test_enter_exit.py
+++ b/tests/unit/container/test_enter_exit.py
@@ -1,3 +1,6 @@
+from collections.abc import AsyncGenerator, Generator
+from unittest.mock import Mock
+
 import pytest
 
 from dishka import (
@@ -9,6 +12,7 @@ from dishka import (
 )
 from ..sample_providers import (
     ClassA,
+    async_gen_a,
     sync_gen_a,
 )
 
@@ -90,6 +94,45 @@ def test_sync_enter(provide_scope, start_scope, expected_scope):
     assert a.closed
 
 
+@pytest.mark.parametrize(
+    ("provide_scope", "start_scope", "expected_scope"),
+    [
+        (Scope.SESSION, Scope.REQUEST, Scope.REQUEST),
+        (Scope.REQUEST, Scope.REQUEST, Scope.REQUEST),
+        (Scope.REQUEST, Scope.STEP, Scope.STEP),
+        (Scope.REQUEST, None, Scope.REQUEST),
+    ],
+)
+def test_error_in_sync_enter(provide_scope, start_scope, expected_scope):
+    exception = ValueError("Some exception")
+    mock = Mock()
+
+    class MyProvider(Provider):
+        a = provide(sync_gen_a, scope=provide_scope)
+
+        @provide(scope=provide_scope)
+        def get_int(self) -> Generator[None, int, None]:
+            caught_exc = yield 100
+            mock(caught_exc)
+
+    base_container = make_container(MyProvider())
+
+    try:
+        with base_container(scope=start_scope) as container:
+            assert container.registry.scope is expected_scope
+            a = container.get(ClassA)
+            assert a
+            assert a.dep == 100
+
+            raise exception
+    except ValueError as exc:
+        if exc != exception:
+            raise
+
+    mock.assert_called_with(exception)
+    assert a.closed
+
+
 @pytest.mark.parametrize(
     ("provide_scope", "start_scope", "expected_scope"),
     [
@@ -117,6 +160,94 @@ async def test_async_enter(provide_scope, start_scope, expected_scope):
     assert a.closed
 
 
+@pytest.mark.parametrize(
+    ("provide_scope", "start_scope", "expected_scope"),
+    [
+        (Scope.SESSION, Scope.REQUEST, Scope.REQUEST),
+        (Scope.REQUEST, Scope.REQUEST, Scope.REQUEST),
+        (Scope.REQUEST, Scope.STEP, Scope.STEP),
+        (Scope.REQUEST, None, Scope.REQUEST),
+    ],
+)
+@pytest.mark.asyncio
+async def test_error_in_async_enter(
+    provide_scope,
+    start_scope,
+    expected_scope,
+):
+    exception = ValueError("Some exception")
+    mock = Mock()
+
+    class MyProvider(Provider):
+        a = provide(sync_gen_a, scope=provide_scope)
+
+        @provide(scope=provide_scope)
+        def get_int(self) -> Generator[None, int, None]:
+            caught_exc = yield 100
+            mock(caught_exc)
+
+    base_container = make_async_container(MyProvider())
+
+    try:
+        async with base_container(scope=start_scope) as container:
+            assert container.registry.scope is expected_scope
+            a = await container.get(ClassA)
+            assert a
+            assert a.dep == 100
+
+            raise exception
+    except ValueError as exc:
+        if exc != exception:
+            raise
+
+    mock.assert_called_with(exception)
+    assert a.closed
+
+
+@pytest.mark.parametrize(
+    ("provide_scope", "start_scope", "expected_scope"),
+    [
+        (Scope.SESSION, Scope.REQUEST, Scope.REQUEST),
+        (Scope.REQUEST, Scope.REQUEST, Scope.REQUEST),
+        (Scope.REQUEST, Scope.STEP, Scope.STEP),
+        (Scope.REQUEST, None, Scope.REQUEST),
+    ],
+)
+@pytest.mark.asyncio
+async def test_error_in_async_enter_with_async_gen(
+    provide_scope,
+    start_scope,
+    expected_scope,
+):
+    exception = ValueError("Some exception")
+    mock = Mock()
+
+    class MyProvider(Provider):
+        a = provide(async_gen_a, scope=provide_scope)
+
+        @provide(scope=provide_scope)
+        async def get_int(self) -> AsyncGenerator[int, None]:
+            caught_exc = yield 100
+            mock(caught_exc)
+
+    base_container = make_async_container(MyProvider())
+
+    try:
+        async with base_container(scope=start_scope) as container:
+            assert container.registry.scope is expected_scope
+            a = await container.get(ClassA)
+            assert a
+            assert a.dep == 100
+
+            raise exception
+    except ValueError as exc:
+        if exc != exception:
+            raise
+
+    mock.assert_called_with(exception)
+    assert a.closed
+
+
 def test_sync_exit_one():
     class MyProvider(Provider):
         a = provide(sync_gen_a, scope=Scope.RUNTIME)