Repository: reagento/dishka
Pull request #452 (2025-04-28 12:58:31)

feat: allow to extend container context on calling injected func
Relates #450 

## Description
In this PR `wrap_injection` function accepts `provide_context` that allows to provide injected variables on injected function call
It could be useful for building integrations where there is no support for middleware and/or request hooks, as well as for more granular deeper scopes
it works only with `manage_scope=Tru`

## Notes (UPD)
1. I tried not create @ovlerload to minimize the code, but check it during runtime with the corresponding error message
2. [x] Shall I add documentation?
3. [x] There is a problem with function `wrap_injection` complexity - C901 and PLR0912 (needs to discuss)


Tishka17: Docs?

https://dishka.readthedocs.io/en/stable/integrations/adding_new.html

Patch:
diff --git a/docs/integrations/adding_new.rst b/docs/integrations/adding_new.rst
index 50dc43b..aa11f45 100644
--- a/docs/integrations/adding_new.rst
+++ b/docs/integrations/adding_new.rst
@@ -11,6 +11,8 @@ The main points are:
 2. Find a place to enter and exit request scope and how to pass the container to a handler. Usually, it is entered in a middleware and container is stored in some kind of request context. 
 
    Alternatively, you can use the ``wrap_injection`` function with ``manage_scope=True`` to automate entering and exiting the request scope without relying on middleware. When enabled, ``manage_scope`` ensures that the container passed to ``wrap_injection`` enters and exits the next scope.
+
+   When you use ``manage_scope=True``, you can pass ``provide_context`` function that allows to populate the container context with passed arguments. This is useful when you want to pass the function scoped arguments to other dependencies without relying on middleware.
 3. Configure a decorator. The main option here is to provide a way for retrieving container. Often, need to modify handler signature adding additional parameters. It is also available.
 4. Check if you can apply decorator automatically.
 
diff --git a/src/dishka/entities/with_parents.py b/src/dishka/entities/with_parents.py
index e4fdd3c..29ec257 100644
--- a/src/dishka/entities/with_parents.py
+++ b/src/dishka/entities/with_parents.py
@@ -193,7 +193,7 @@ class ParentsResolver:
 
 if TYPE_CHECKING:
     T = TypeVar("T")
-    WithParents: TypeAlias = T
+    WithParents: TypeAlias = T | T  # noqa: PYI016
 else:
     class WithParents:
         def __class_getitem__(cls, item: TypeHint) -> TypeHint:
diff --git a/src/dishka/integrations/base.py b/src/dishka/integrations/base.py
index 31f35bd..680d1ca 100644
--- a/src/dishka/integrations/base.py
+++ b/src/dishka/integrations/base.py
@@ -35,12 +35,19 @@ from dishka.container import Container
 from dishka.entities.component import DEFAULT_COMPONENT
 from dishka.entities.depends_marker import FromDishka
 from dishka.entities.key import DependencyKey, _FromComponent
-from dishka.integrations.exceptions import InvalidInjectedFuncTypeError
+from dishka.integrations.exceptions import (
+    ImproperProvideContextUsageError,
+    InvalidInjectedFuncTypeError,
+)
 
 T = TypeVar("T")
 P = ParamSpec("P")
 DependencyParser: TypeAlias = Callable[[Parameter, Any], DependencyKey | None]
 ContainerGetter: TypeAlias = Callable[[tuple[Any, ...], dict[str, Any]], T]
+ProvideContext: TypeAlias = Callable[
+    [tuple[Any, ...], dict[str, Any]],
+    dict[Any, Any],
+]
 DependsClass: TypeAlias = cast(
     type | Sequence[type],
     FromDishka | _FromComponent,
@@ -52,6 +59,7 @@ def _get_auto_injected_async_gen_scoped(
     additional_params: Sequence[Parameter],
     dependencies: dict[str, DependencyKey],
     func: Callable[P, AsyncIterator[T]],
+    provide_context: ProvideContext | None = None,
 ) -> Callable[P, AsyncIterator[T]]:
     async def auto_injected_generator(
         *args: P.args,
@@ -60,8 +68,11 @@ def _get_auto_injected_async_gen_scoped(
         for param in additional_params:
             kwargs.pop(param.name)
 
+        additional_context = (
+            {} if provide_context is None else provide_context(args, kwargs)
+        )
         container = container_getter(args, kwargs)
-        async with container() as container:
+        async with container(additional_context) as container:
             solved = {
                 name: await container.get(
                     dep.type_hint,
@@ -80,7 +91,11 @@ def _get_auto_injected_async_gen(
     additional_params: Sequence[Parameter],
     dependencies: dict[str, DependencyKey],
     func: Callable[P, AsyncIterator[T]],
+    provide_context: ProvideContext | None = None,
 ) -> Callable[P, AsyncIterator[T]]:
+    if provide_context is not None:
+        raise ImproperProvideContextUsageError
+
     async def auto_injected_generator(
         *args: P.args,
         **kwargs: P.kwargs,
@@ -107,12 +122,17 @@ def _get_auto_injected_async_func_scoped(
     additional_params: Sequence[Parameter],
     dependencies: dict[str, DependencyKey],
     func: Callable[P, Awaitable[T]],
+    provide_context: ProvideContext | None = None,
 ) -> Callable[P, Awaitable[T]]:
     async def auto_injected_func(*args: P.args, **kwargs: P.kwargs) -> T:
         container = container_getter(args, kwargs)
         for param in additional_params:
             kwargs.pop(param.name)
-        async with container() as container:
+
+        additional_context = (
+            {} if provide_context is None else provide_context(args, kwargs)
+        )
+        async with container(additional_context) as container:
             solved = {
                 name: await container.get(
                     dep.type_hint,
@@ -130,7 +150,11 @@ def _get_auto_injected_async_func(
     additional_params: Sequence[Parameter],
     dependencies: dict[str, DependencyKey],
     func: Callable[P, Awaitable[T]],
+    provide_context: ProvideContext | None = None,
 ) -> Callable[P, Awaitable[T]]:
+    if provide_context is not None:
+        raise ImproperProvideContextUsageError
+
     async def auto_injected_func(*args: P.args, **kwargs: P.kwargs) -> T:
         container = container_getter(args, kwargs)
         for param in additional_params:
@@ -152,6 +176,7 @@ def _get_auto_injected_sync_gen_scoped(
     additional_params: Sequence[Parameter],
     dependencies: dict[str, DependencyKey],
     func: Callable[P, Iterator[T]],
+    provide_context: ProvideContext | None = None,
 ) -> Callable[P, Iterator[T]]:
     def auto_injected_generator(
         *args: P.args,
@@ -160,8 +185,11 @@ def _get_auto_injected_sync_gen_scoped(
         for param in additional_params:
             kwargs.pop(param.name)
 
+        additional_context = (
+            {} if provide_context is None else provide_context(args, kwargs)
+        )
         container = container_getter(args, kwargs)
-        with container() as container:
+        with container(additional_context) as container:
             solved = {
                 name: container.get(dep.type_hint, component=dep.component)
                 for name, dep in dependencies.items()
@@ -176,7 +204,11 @@ def _get_auto_injected_sync_gen(
     additional_params: Sequence[Parameter],
     dependencies: dict[str, DependencyKey],
     func: Callable[P, Iterator[T]],
+    provide_context: ProvideContext | None = None,
 ) -> Callable[P, Iterator[T]]:
+    if provide_context is not None:
+        raise ImproperProvideContextUsageError
+
     def auto_injected_generator(
         *args: P.args,
         **kwargs: P.kwargs,
@@ -198,13 +230,17 @@ def _get_auto_injected_sync_func_scoped(
     additional_params: Sequence[Parameter],
     dependencies: dict[str, DependencyKey],
     func: Callable[P, T],
+    provide_context: ProvideContext | None = None,
 ) -> Callable[P, T]:
     def auto_injected_func(*args: P.args, **kwargs: P.kwargs) -> T:
         for param in additional_params:
             kwargs.pop(param.name)
 
+        additional_context = (
+            {} if provide_context is None else provide_context(args, kwargs)
+        )
         container = container_getter(args, kwargs)
-        with container() as container:
+        with container(additional_context) as container:
             solved = {
                 name: container.get(dep.type_hint, component=dep.component)
                 for name, dep in dependencies.items()
@@ -219,7 +255,11 @@ def _get_auto_injected_sync_func(
     additional_params: Sequence[Parameter],
     dependencies: dict[str, DependencyKey],
     func: Callable[P, T],
+    provide_context: ProvideContext | None = None,
 ) -> Callable[P, T]:
+    if provide_context is not None:
+        raise ImproperProvideContextUsageError
+
     def auto_injected_func(*args: P.args, **kwargs: P.kwargs) -> T:
         container = container_getter(args, kwargs)
         for param in additional_params:
@@ -239,6 +279,7 @@ def _get_auto_injected_sync_container_async_gen_scoped(
     additional_params: Sequence[Parameter],
     dependencies: dict[str, DependencyKey],
     func: Callable[P, Iterator[T]],
+    provide_context: ProvideContext | None = None,
 ) -> Callable[P, AsyncIterator[T]]:
     async def auto_injected_generator(
         *args: P.args,
@@ -247,8 +288,11 @@ def _get_auto_injected_sync_container_async_gen_scoped(
         for param in additional_params:
             kwargs.pop(param.name)
 
+        additional_context = (
+            {} if provide_context is None else provide_context(args, kwargs)
+        )
         container = container_getter(args, kwargs)
-        with container() as container:
+        with container(additional_context) as container:
             solved = {
                 name: container.get(dep.type_hint, component=dep.component)
                 for name, dep in dependencies.items()
@@ -264,7 +308,11 @@ def _get_auto_injected_sync_container_async_gen(
     additional_params: Sequence[Parameter],
     dependencies: dict[str, DependencyKey],
     func: Callable[P, Iterator[T]],
+    provide_context: ProvideContext | None = None,
 ) -> Callable[P, AsyncIterator[T]]:
+    if provide_context is not None:
+        raise ImproperProvideContextUsageError
+
     async def auto_injected_generator(
         *args: P.args,
         **kwargs: P.kwargs,
@@ -287,13 +335,17 @@ def _get_auto_injected_sync_container_async_func_scoped(
     additional_params: Sequence[Parameter],
     dependencies: dict[str, DependencyKey],
     func: Callable[P, T],
+    provide_context: ProvideContext | None = None,
 ) -> Callable[P, T]:
     async def auto_injected_func(*args: P.args, **kwargs: P.kwargs) -> T:
         for param in additional_params:
             kwargs.pop(param.name)
 
+        additional_context = (
+            {} if provide_context is None else provide_context(args, kwargs)
+        )
         container = container_getter(args, kwargs)
-        with container() as container:
+        with container(additional_context) as container:
             solved = {
                 name: container.get(dep.type_hint, component=dep.component)
                 for name, dep in dependencies.items()
@@ -308,7 +360,11 @@ def _get_auto_injected_sync_container_async_func(
     additional_params: Sequence[Parameter],
     dependencies: dict[str, DependencyKey],
     func: Callable[P, T],
+    provide_context: ProvideContext | None = None,
 ) -> Callable[P, T]:
+    if provide_context is not None:
+        raise ImproperProvideContextUsageError
+
     async def auto_injected_func(*args: P.args, **kwargs: P.kwargs) -> T:
         container = container_getter(args, kwargs)
         for param in additional_params:
@@ -449,6 +505,7 @@ def wrap_injection(
     additional_params: Sequence[Parameter] = (),
     parse_dependency: DependencyParser = default_parse_dependency,
     manage_scope: bool = False,
+    provide_context: ProvideContext | None = None,
 ) -> Callable[P, T]: ...
 
 
@@ -462,6 +519,7 @@ def wrap_injection(
     additional_params: Sequence[Parameter] = (),
     parse_dependency: DependencyParser = default_parse_dependency,
     manage_scope: bool = False,
+    provide_context: ProvideContext | None = None,
 ) -> Callable[P, T]: ...
 
 
@@ -474,6 +532,7 @@ def wrap_injection(
     additional_params: Sequence[Parameter] = (),
     parse_dependency: DependencyParser = default_parse_dependency,
     manage_scope: bool = False,
+    provide_context: ProvideContext | None = None,
 ) -> Callable[P, T]:
     hints = get_type_hints(func, include_extras=True)
     func_signature = signature(func)
@@ -541,6 +600,7 @@ def wrap_injection(
 
     auto_injected_func = get_auto_injected_func(
         func=func,
+        provide_context=provide_context,
         dependencies=dependencies,
         additional_params=additional_params,
         container_getter=container_getter,
diff --git a/src/dishka/integrations/exceptions.py b/src/dishka/integrations/exceptions.py
index ef6551b..81dee07 100644
--- a/src/dishka/integrations/exceptions.py
+++ b/src/dishka/integrations/exceptions.py
@@ -3,6 +3,9 @@ from dishka.exception_base import DishkaError
 
 class InvalidInjectedFuncTypeError(DishkaError):
     def __str__(self) -> str:
-        return (
-            "An async container cannot be used in a synchronous context."
-        )
+        return "An async container cannot be used in a synchronous context."
+
+
+class ImproperProvideContextUsageError(DishkaError):
+    def __str__(self) -> str:
+        return "provide_context can only be used with manage_scope=True."


Test patch:
diff --git a/tests/integrations/base/test_wrap_injection_provide_context.py b/tests/integrations/base/test_wrap_injection_provide_context.py
new file mode 100644
index 0000000..fd6cd7e
--- /dev/null
+++ b/tests/integrations/base/test_wrap_injection_provide_context.py
@@ -0,0 +1,134 @@
+from collections.abc import AsyncIterator, Callable, Iterable, Iterator
+from inspect import isasyncgenfunction, isgeneratorfunction
+from typing import Any, cast
+
+import pytest
+
+from dishka.async_container import AsyncContainer
+from dishka.container import Container
+from dishka.entities.depends_marker import FromDishka
+from dishka.integrations.base import wrap_injection
+from dishka.integrations.exceptions import ImproperProvideContextUsageError
+from tests.integrations.common import ContextDep, UserDep
+
+context = ContextDep("context")
+user = UserDep(f"user_id.from_context({context})")
+
+
+def generate_str(
+    user: UserDep,
+    context: ContextDep,
+    prefix: int | None = None,
+) -> str:
+    s = f"{user}, context={context}"
+    if prefix is not None:
+        s = f"{prefix}. {s}"
+    return s
+
+
+def sync_func(
+    context: ContextDep,
+    user: FromDishka[UserDep],
+) -> str:
+    return generate_str(user, context)
+
+
+def sync_gen(
+    context: ContextDep,
+    data: Iterable[int],
+    user: FromDishka[UserDep],
+) -> Iterator[str]:
+    for i in data:
+        yield generate_str(user, context, i)
+
+
+async def async_func(
+    context: ContextDep,
+    user: FromDishka[UserDep],
+) -> str:
+    return generate_str(user, context)
+
+
+async def async_gen(
+    context: ContextDep,
+    data: Iterable[int],
+    user: FromDishka[UserDep],
+) -> AsyncIterator[str]:
+    for i in data:
+        yield generate_str(user, context, i)
+
+
+def provide_context(
+    args: tuple[Any, ...],
+    _: dict[str, Any],
+) -> dict[Any, Any]:
+    context = cast(ContextDep, args[0])
+    return {ContextDep: context}
+
+
+@pytest.mark.parametrize("func", [sync_func, sync_gen])
+def test_sync_provide_context(func: Callable, container: Container) -> None:
+    wrapped_func = wrap_injection(
+        func=func,
+        container_getter=lambda *_: container,
+        is_async=False,
+        manage_scope=True,
+        provide_context=provide_context,
+    )
+
+    if isgeneratorfunction(func):
+        result = list(wrapped_func(context, [1, 2]))
+        assert result == [generate_str(user, context, i) for i in [1, 2]]
+    else:
+        result = wrapped_func(context)
+        assert result == generate_str(user, context)
+
+
+@pytest.mark.parametrize("func", [sync_func, sync_gen])
+def test_invalid_provide_context(func: Callable, container: Container) -> None:
+    with pytest.raises(ImproperProvideContextUsageError):
+        wrap_injection(
+            func=func,
+            container_getter=lambda *_: container,
+            is_async=False,
+            manage_scope=False,
+            provide_context=provide_context,
+        )
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("func", [async_func, async_gen])
+async def test_async_provide_context(
+    func: Callable,
+    async_container: AsyncContainer,
+) -> None:
+    wrapped_func = wrap_injection(
+        func=func,
+        container_getter=lambda *_: async_container,
+        is_async=True,
+        manage_scope=True,
+        provide_context=provide_context,
+    )
+
+    if isasyncgenfunction(func):
+        result = [item async for item in wrapped_func(context, [1, 2])]
+        assert result == [generate_str(user, context, i) for i in [1, 2]]
+    else:
+        result = await wrapped_func(context)
+        assert result == generate_str(user, context)
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("func", [async_func, async_gen])
+async def test_invalid_async_provide_context(
+    func: Callable,
+    async_container: AsyncContainer,
+) -> None:
+    with pytest.raises(ImproperProvideContextUsageError):
+        wrap_injection(
+            func=func,
+            container_getter=lambda *_: async_container,
+            is_async=True,
+            manage_scope=False,
+            provide_context=provide_context,
+        )
diff --git a/tests/integrations/common.py b/tests/integrations/common.py
index ee6b5fc..aeb2ccb 100644
--- a/tests/integrations/common.py
+++ b/tests/integrations/common.py
@@ -2,7 +2,11 @@ from collections.abc import Iterable
 from typing import NewType
 from unittest.mock import Mock
 
-from dishka import Provider, Scope, provide
+from dishka import Provider, Scope, from_context, provide
+from dishka.entities.depends_marker import FromDishka
+
+ContextDep = NewType("ContextDep", str)
+UserDep = NewType("UserDep", str)
 
 AppDep = NewType("AppDep", str)
 APP_DEP_VALUE = "APP"
@@ -17,11 +21,14 @@ AppMock = NewType("AppMock", Mock)
 
 
 class AppProvider(Provider):
+    context = from_context(provides=ContextDep, scope=Scope.REQUEST)
+
     def __init__(self) -> None:
         super().__init__()
         self.app_released = Mock()
         self.request_released = Mock()
         self.websocket_released = Mock()
+        self.db_session_released = Mock()
         self.mock = Mock()
         self.app_mock = AppMock(Mock())
 
@@ -40,6 +47,11 @@ class AppProvider(Provider):
         yield WS_DEP_VALUE
         self.websocket_released()
 
+    @provide(scope=Scope.REQUEST)
+    def user(self, context: FromDishka[ContextDep]) -> Iterable[UserDep]:
+        yield f"user_id.from_context({context})"
+        self.db_session_released()
+
     @provide(scope=Scope.REQUEST)
     def get_mock(self) -> Mock:
         return self.mock