Repository: reagento/dishka

Issue:
Support celery in threaded mode
Currently we store request container in app config. It works ok for worker processes, but would break for worker threads.

Discussion:

Solution patch:
diff --git a/src/dishka/factory_compiler.py b/src/dishka/factory_compiler.py
index da0b627..4c086f6 100644
--- a/src/dishka/factory_compiler.py
+++ b/src/dishka/factory_compiler.py
@@ -25,21 +25,12 @@ from .dependency_source import Factory
 from .exceptions import NoContextValueError, UnsupportedFactoryError
 
 
-def make_args(args: list[str], kwargs: list[str]) -> str:
-    res = ", ".join(
-        f"{{await}}getter({arg})"
-        for arg in args
+def make_args(args: list[str], kwargs: dict[str, str]) -> str:
+    """Format arguments for the factory function."""
+    return ", ".join(
+        [f"{{await}}getter({arg})" for arg in args] +
+        [f"{arg}={{await}}getter({var})" for arg, var in kwargs.items()],
     )
-    if not kwargs:
-        return res
-    if res:
-        res += ", "
-    res += ", ".join(
-        f"{arg}={{await}}getter({arg})"
-        for arg in kwargs
-    )
-    return res
-
 
 GENERATOR = """
 {async}def get(getter, exits, context):
@@ -115,7 +106,14 @@ def compile_factory(*, factory: Factory, is_async: bool) -> CompiledFactory:
         f"_dishka_arg{i}": dep
         for i, dep in enumerate(factory.dependencies)
     }
-    kwargs = factory.kw_dependencies
+    kwargs_to_globals = {
+        name: f"_dishka_kwarg_{name}"
+        for name in factory.kw_dependencies
+    }
+    kwargs = {
+        kwargs_to_globals[name]: dep
+        for name, dep in factory.kw_dependencies.items()
+    }
 
     if is_async:
         async_ = "async "
@@ -125,19 +123,15 @@ def compile_factory(*, factory: Factory, is_async: bool) -> CompiledFactory:
         async_ = ""
         await_ = ""
         body_template = SYNC_BODIES.get(factory.type, INVALID)
-    if factory.cache:
-        cache = CACHE
-    else:
-        cache = ""
 
-    args_str = make_args(list(args), list(kwargs)).format_map({
+    args_str = make_args(list(args), kwargs_to_globals).format_map({
         "await": await_,
     })
     body = body_template.format_map({
         "async": async_,
         "await": await_,
         "args": args_str,
-        "cache": cache,
+        "cache": CACHE if factory.cache else "",
     })
     func_globals = {
         "source": factory.source,
diff --git a/src/dishka/integrations/base.py b/src/dishka/integrations/base.py
index 94456e0..2a07a63 100644
--- a/src/dishka/integrations/base.py
+++ b/src/dishka/integrations/base.py
@@ -1,4 +1,11 @@
-from collections.abc import Awaitable, Callable, Sequence
+from collections.abc import (
+    AsyncIterator,
+    Awaitable,
+    Callable,
+    Iterator,
+    Sequence,
+)
+from dataclasses import dataclass
 from inspect import (
     Parameter,
     Signature,
@@ -37,9 +44,248 @@ DependsClass: TypeAlias = cast(
 )
 
 
+def _get_auto_injected_async_gen_scoped(
+    container_getter: ContainerGetter[AsyncContainer],
+    additional_params: Sequence[Parameter],
+    dependencies: dict[str, DependencyKey],
+    func: Callable[P, AsyncIterator[T]],
+) -> Callable[P, AsyncIterator[T]]:
+    async def auto_injected_generator(
+        *args: P.args,
+        **kwargs: P.kwargs,
+    ) -> AsyncIterator[T]:
+        for param in additional_params:
+            kwargs.pop(param.name)
+
+        container = container_getter(args, kwargs)
+        async with container() as container:
+            solved = {
+                name: await container.get(
+                    dep.type_hint,
+                    component=dep.component,
+                )
+                for name, dep in dependencies.items()
+            }
+            async for message in func(*args, **kwargs, **solved):
+                yield message
+
+    return auto_injected_generator
+
+
+def _get_auto_injected_async_gen(
+    container_getter: ContainerGetter[AsyncContainer],
+    additional_params: Sequence[Parameter],
+    dependencies: dict[str, DependencyKey],
+    func: Callable[P, AsyncIterator[T]],
+) -> Callable[P, AsyncIterator[T]]:
+    async def auto_injected_generator(
+        *args: P.args,
+        **kwargs: P.kwargs,
+    ) -> AsyncIterator[T]:
+        for param in additional_params:
+            kwargs.pop(param.name)
+
+        container = container_getter(args, kwargs)
+        solved = {
+            name: await container.get(
+                dep.type_hint,
+                component=dep.component,
+            )
+            for name, dep in dependencies.items()
+        }
+        async for message in func(*args, **kwargs, **solved):
+            yield message
+
+    return auto_injected_generator
+
+
+def _get_auto_injected_async_func_scoped(
+    container_getter: ContainerGetter[AsyncContainer],
+    additional_params: Sequence[Parameter],
+    dependencies: dict[str, DependencyKey],
+    func: Callable[P, Awaitable[T]],
+) -> Callable[P, Awaitable[T]]:
+    async def auto_injected_func(*args: P.args, **kwargs: P.kwargs) -> T:
+        for param in additional_params:
+            kwargs.pop(param.name)
+
+        container = container_getter(args, kwargs)
+        async with container() as container:
+            solved = {
+                name: await container.get(
+                    dep.type_hint,
+                    component=dep.component,
+                )
+                for name, dep in dependencies.items()
+            }
+            return await func(*args, **kwargs, **solved)
+
+    return auto_injected_func
+
+
+def _get_auto_injected_async_func(
+    container_getter: ContainerGetter[AsyncContainer],
+    additional_params: Sequence[Parameter],
+    dependencies: dict[str, DependencyKey],
+    func: Callable[P, Awaitable[T]],
+) -> 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)
+        solved = {
+            name: await container.get(
+                dep.type_hint,
+                component=dep.component,
+            )
+            for name, dep in dependencies.items()
+        }
+        return await func(*args, **kwargs, **solved)
+
+    return auto_injected_func
+
+
+def _get_auto_injected_sync_gen_scoped(
+    container_getter: ContainerGetter[Container],
+    additional_params: Sequence[Parameter],
+    dependencies: dict[str, DependencyKey],
+    func: Callable[P, Iterator[T]],
+) -> Callable[P, Iterator[T]]:
+    def auto_injected_generator(
+        *args: P.args,
+        **kwargs: P.kwargs,
+    ) -> Iterator[T]:
+        for param in additional_params:
+            kwargs.pop(param.name)
+
+        container = container_getter(args, kwargs)
+        with container() as container:
+            solved = {
+                name: container.get(dep.type_hint, component=dep.component)
+                for name, dep in dependencies.items()
+            }
+            yield from func(*args, **kwargs, **solved)
+
+    return auto_injected_generator
+
+
+def _get_auto_injected_sync_gen(
+    container_getter: ContainerGetter[Container],
+    additional_params: Sequence[Parameter],
+    dependencies: dict[str, DependencyKey],
+    func: Callable[P, Iterator[T]],
+) -> Callable[P, Iterator[T]]:
+    def auto_injected_generator(
+        *args: P.args,
+        **kwargs: P.kwargs,
+    ) -> Iterator[T]:
+        container = container_getter(args, kwargs)
+        for param in additional_params:
+            kwargs.pop(param.name)
+        solved = {
+            name: container.get(dep.type_hint, component=dep.component)
+            for name, dep in dependencies.items()
+        }
+        yield from func(*args, **kwargs, **solved)
+
+    return auto_injected_generator
+
+
+def _get_auto_injected_sync_func_scoped(
+    container_getter: ContainerGetter[Container],
+    additional_params: Sequence[Parameter],
+    dependencies: dict[str, DependencyKey],
+    func: Callable[P, T],
+) -> Callable[P, T]:
+    def auto_injected_func(*args: P.args, **kwargs: P.kwargs) -> T:
+        for param in additional_params:
+            kwargs.pop(param.name)
+
+        container = container_getter(args, kwargs)
+        with container() as container:
+            solved = {
+                name: container.get(dep.type_hint, component=dep.component)
+                for name, dep in dependencies.items()
+            }
+            return func(*args, **kwargs, **solved)
+
+    return auto_injected_func
+
+
+def _get_auto_injected_sync_func(
+    container_getter: ContainerGetter[Container],
+    additional_params: Sequence[Parameter],
+    dependencies: dict[str, DependencyKey],
+    func: Callable[P, T],
+) -> Callable[P, T]:
+    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)
+
+        solved = {
+            name: container.get(dep.type_hint, component=dep.component)
+            for name, dep in dependencies.items()
+        }
+        return func(*args, **kwargs, **solved)
+
+    return auto_injected_func
+
+
+@dataclass(frozen=True, slots=True, kw_only=True)
+class InjectedFuncType:
+    is_async: bool
+    is_generator: bool
+    is_scoped: bool
+
+
+_GET_AUTO_INJECTED_FUNC_DICT = {
+    InjectedFuncType(
+        is_async=True,
+        is_generator=True,
+        is_scoped=True,
+    ): _get_auto_injected_async_gen_scoped,
+    InjectedFuncType(
+        is_async=True,
+        is_generator=True,
+        is_scoped=False,
+    ): _get_auto_injected_async_gen,
+    InjectedFuncType(
+        is_async=True,
+        is_generator=False,
+        is_scoped=True,
+    ): _get_auto_injected_async_func_scoped,
+    InjectedFuncType(
+        is_async=True,
+        is_generator=False,
+        is_scoped=False,
+    ): _get_auto_injected_async_func,
+    InjectedFuncType(
+        is_async=False,
+        is_generator=True,
+        is_scoped=True,
+    ): _get_auto_injected_sync_gen_scoped,
+    InjectedFuncType(
+        is_async=False,
+        is_generator=True,
+        is_scoped=False,
+    ): _get_auto_injected_sync_gen,
+    InjectedFuncType(
+        is_async=False,
+        is_generator=False,
+        is_scoped=True,
+    ): _get_auto_injected_sync_func_scoped,
+    InjectedFuncType(
+        is_async=False,
+        is_generator=False,
+        is_scoped=False,
+    ): _get_auto_injected_sync_func,
+}
+
+
 def default_parse_dependency(
-        parameter: Parameter,
-        hint: Any,
+    parameter: Parameter,
+    hint: Any,
 ) -> DependencyKey | None:
     """Resolve dependency type or return None if it is not a dependency."""
     if get_origin(hint) is not Annotated:
@@ -59,38 +305,39 @@ def default_parse_dependency(
 
 @overload
 def wrap_injection(
-        *,
-        func: Callable[P, T],
-        container_getter: ContainerGetter[Container],
-        is_async: Literal[False] = False,
-        remove_depends: bool = True,
-        additional_params: Sequence[Parameter] = (),
-        parse_dependency: DependencyParser = default_parse_dependency,
-) -> Callable[P, T]:
-    ...
+    *,
+    func: Callable[P, T],
+    container_getter: ContainerGetter[Container],
+    is_async: Literal[False] = False,
+    remove_depends: bool = True,
+    additional_params: Sequence[Parameter] = (),
+    parse_dependency: DependencyParser = default_parse_dependency,
+    manage_scope: bool = False,
+) -> Callable[P, T]: ...
 
 
 @overload
 def wrap_injection(
-        *,
-        func: Callable[P, T],
-        container_getter: ContainerGetter[AsyncContainer],
-        is_async: Literal[True] = True,
-        remove_depends: bool = True,
-        additional_params: Sequence[Parameter] = (),
-        parse_dependency: DependencyParser = default_parse_dependency,
-) -> Callable[P, T]:
-    ...
+    *,
+    func: Callable[P, T],
+    container_getter: ContainerGetter[AsyncContainer],
+    is_async: Literal[True] = True,
+    remove_depends: bool = True,
+    additional_params: Sequence[Parameter] = (),
+    parse_dependency: DependencyParser = default_parse_dependency,
+    manage_scope: bool = False,
+) -> Callable[P, T]: ...
 
 
 def wrap_injection(
-        *,
-        func: Callable[P, T],
-        container_getter: ContainerGetter[Container | AsyncContainer],
-        is_async: bool = False,
-        remove_depends: bool = True,
-        additional_params: Sequence[Parameter] = (),
-        parse_dependency: DependencyParser = default_parse_dependency,
+    *,
+    func: Callable[P, T],
+    container_getter: ContainerGetter[Container | AsyncContainer],
+    is_async: bool = False,
+    remove_depends: bool = True,
+    additional_params: Sequence[Parameter] = (),
+    parse_dependency: DependencyParser = default_parse_dependency,
+    manage_scope: bool = False,
 ) -> Callable[P, T]:
     hints = get_type_hints(func, include_extras=True)
     func_signature = signature(func)
@@ -138,25 +385,33 @@ def wrap_injection(
             new_annotations[param.name] = param.annotation
 
     if is_async:
-        auto_injected_func = _async_injection_wrapper(
-            # typing.cast is needed because the function must be async
-            func=cast(Callable[P, Awaitable[T]],func),
-            dependencies=dependencies,
-            additional_params=additional_params,
-            container_getter=cast(
-                ContainerGetter[AsyncContainer], container_getter,
-            ),
+        func = cast(Callable[P, Awaitable[T]], func)
+        is_generator = isasyncgenfunction(func)
+        container_getter = cast(
+            ContainerGetter[AsyncContainer],
+            container_getter,
         )
     else:
-        auto_injected_func = _sync_injection_wrapper(
-            func=func,
-            dependencies=dependencies,
-            additional_params=additional_params,
-            container_getter=cast(
-                ContainerGetter[Container], container_getter,
-            ),
+        is_generator = isgeneratorfunction(func)
+        container_getter = cast(
+            ContainerGetter[Container],
+            container_getter,
         )
 
+    injected_func_type = InjectedFuncType(
+        is_async=is_async,
+        is_generator=is_generator,
+        is_scoped=manage_scope,
+    )
+    get_auto_injected_func = _GET_AUTO_INJECTED_FUNC_DICT[injected_func_type]
+
+    auto_injected_func = get_auto_injected_func(
+        func=func,
+        dependencies=dependencies,
+        additional_params=additional_params,
+        container_getter=container_getter,
+    )
+
     auto_injected_func.__dishka_orig_func__ = func
     auto_injected_func.__dishka_injected__ = True  # type: ignore[attr-defined]
     auto_injected_func.__name__ = func.__name__
@@ -175,71 +430,6 @@ def is_dishka_injected(func: Callable[..., Any]) -> bool:
     return hasattr(func, "__dishka_injected__")
 
 
-def _async_injection_wrapper(
-        container_getter: ContainerGetter[AsyncContainer],
-        additional_params: Sequence[Parameter],
-        dependencies: dict[str, DependencyKey],
-        func: Callable[P, Awaitable[T]],
-) -> Callable[P, Awaitable[T]]:
-    if isasyncgenfunction(func):
-        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)
-            solved = {
-                name: await container.get(
-                    dep.type_hint, component=dep.component,
-                )
-                for name, dep in dependencies.items()
-            }
-            async for message in func(*args, **kwargs, **solved):
-                yield message
-    else:
-        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)
-            solved = {
-                name: await container.get(
-                    dep.type_hint, component=dep.component,
-                )
-                for name, dep in dependencies.items()
-            }
-            return await func(*args, **kwargs, **solved)
-
-    return auto_injected_func
-
-
-def _sync_injection_wrapper(
-        container_getter: ContainerGetter[Container],
-        additional_params: Sequence[Parameter],
-        dependencies: dict[str, DependencyKey],
-        func: Callable[P, T],
-) -> Callable[P, T]:
-    if isgeneratorfunction(func):
-        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)
-            solved = {
-                name: container.get(dep.type_hint, component=dep.component)
-                for name, dep in dependencies.items()
-            }
-            yield from func(*args, **kwargs, **solved)
-    else:
-        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)
-            solved = {
-                name: container.get(dep.type_hint, component=dep.component)
-                for name, dep in dependencies.items()
-            }
-            return func(*args, **kwargs, **solved)
-
-    return auto_injected_func
-
-
 def _add_params(
     params: Sequence[Parameter],
     additional_params: Sequence[Parameter],
@@ -252,7 +442,6 @@ def _add_params(
     for param in additional_params:
         params_kind_dict.setdefault(param.kind, []).append(param)
 
-
     var_positional = params_kind_dict.get(Parameter.VAR_POSITIONAL, [])
     if len(var_positional) > 1:
         param_names = (param.name for param in var_positional)
diff --git a/src/dishka/integrations/celery.py b/src/dishka/integrations/celery.py
index ae39a0b..c21613e 100644
--- a/src/dishka/integrations/celery.py
+++ b/src/dishka/integrations/celery.py
@@ -6,17 +6,15 @@ __all__ = [
 ]
 
 from collections.abc import Callable
-from typing import Any, Final, ParamSpec, TypeVar
+from typing import Final, ParamSpec, TypeVar
 
 from celery import Celery, Task, current_app
-from celery.signals import task_postrun, task_prerun
 from celery.utils.functional import head_from_fun
 
 from dishka import Container, FromDishka
 from dishka.integrations.base import is_dishka_injected, wrap_injection
 
 CONTAINER_NAME: Final = "dishka_container"
-CONTAINER_NAME_REQ: Final = "dishka_container_req"
 
 
 T = TypeVar("T")
@@ -28,8 +26,9 @@ def inject(func: Callable[P, T]) -> Callable[P, T]:
         func=func,
         is_async=False,
         container_getter=lambda args, kwargs: current_app.conf[
-            CONTAINER_NAME_REQ
+            CONTAINER_NAME
         ],
+        manage_scope=True,
     )
 
 
@@ -37,33 +36,6 @@ def setup_dishka(container: Container, app: Celery):
     app.conf[CONTAINER_NAME] = container
 
 
-@task_prerun.connect()
-def enter_scope(
-    task_id,
-    task: Task,
-    args,
-    kwargs: dict[str, Any],
-    **other,
-) -> None:
-    if CONTAINER_NAME in task.app.conf:
-        container: Container = task.app.conf[CONTAINER_NAME]
-        task.app.conf[CONTAINER_NAME_REQ] = container().__enter__()
-
-
-@task_postrun.connect()
-def exit_scope(
-    task_id,
-    task,
-    args,
-    kwargs: dict[str, Any],
-    retval,
-    **other,
-) -> None:
-    if CONTAINER_NAME_REQ in task.app.conf:
-        container: Container = task.app.conf.pop(CONTAINER_NAME_REQ)
-        container.close()
-
-
 class DishkaTask(Task):
     def __init__(self) -> None:
         super().__init__()


Test patch:
diff --git a/tests/unit/container/test_decorator.py b/tests/unit/container/test_decorator.py
index 153f43e..24a8d5f 100644
--- a/tests/unit/container/test_decorator.py
+++ b/tests/unit/container/test_decorator.py
@@ -77,8 +77,8 @@ def test_decorator():
     assert isinstance(a.a, A)
 
 
-def kwarg_decorator(x: int, *, a: A) -> A:
-    return ADecorator(a)
+def kwarg_decorator(x: int, *, source: A) -> A:
+    return ADecorator(source)
 
 
 def test_kwargs():