Add support for context managers

Author: cofob

.flake8

@@ -92,6 +92,8 @@ ignore =
   WPS424,
   ; Found a too complex `f` string
   WPS237,
+  ; Found `no cover` comments overuse: 7 > 5
+  WPS403,
 
 per-file-ignores =
   ; all tests

taskiq_dependencies/ctx.py

@@ -5,7 +5,7 @@
 from logging import getLogger
 from typing import TYPE_CHECKING, Any, DefaultDict, Dict, Generator, List, Optional
 
-from taskiq_dependencies.utils import ParamInfo
+from taskiq_dependencies.utils import ParamInfo, isasynccontextmanager, iscontextmanager
 
 if TYPE_CHECKING:
     from taskiq_dependencies.graph import DependencyGraph  # pragma: no cover
@@ -59,7 +59,7 @@ def traverse_deps(  # noqa: C901, WPS210
             # later.
             if not dep.use_cache:
                 continue
-            # If somehow we have dependency with unknwon function.
+            # If somehow we have dependency with unknown function.
             if dep.dependency is None:
                 continue
             # If dependency is already calculated.
@@ -180,6 +180,8 @@ def close(self, *args: Any) -> None:  # noqa: C901
                     continue
                 for _ in dep:  # noqa: WPS328
                     pass  # noqa: WPS420
+            elif iscontextmanager(dep):
+                dep.__exit__(*args)  # noqa: WPS609
 
     def resolver(self, executed_func: Any, initial_cache: Dict[Any, Any]) -> Any:
         """
@@ -194,7 +196,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):
+        if getattr(executed_func, "dep_graph", False):  # noqa: WPS223
             ctx = SyncResolveContext(executed_func, initial_cache)
             self.sub_contexts.append(ctx)
             sub_result = ctx.resolve_kwargs()
@@ -206,7 +208,10 @@ def resolver(self, executed_func: Any, initial_cache: Dict[Any, Any]) -> Any:
                 "Coroutines cannot be used in sync context. "
                 "Please use async context instead.",
             )
-        elif inspect.isasyncgen(executed_func):
+        elif iscontextmanager(executed_func):
+            sub_result = executed_func.__enter__()  # noqa: WPS609
+            self.opened_dependencies.append(executed_func)
+        elif inspect.isasyncgen(executed_func) or isasynccontextmanager(executed_func):
             raise RuntimeError(
                 "Coroutines cannot be used in sync context. "
                 "Please use async context instead.",
@@ -299,8 +304,16 @@ async def close(self, *args: Any) -> None:  # noqa: C901
                     continue
                 async for _ in dep:  # noqa: WPS328
                     pass  # noqa: WPS420
+            elif iscontextmanager(dep):
+                dep.__exit__(*args)  # noqa: WPS609
+            elif isasynccontextmanager(dep):
+                await dep.__aexit__(*args)  # noqa: WPS609
 
-    async def resolver(self, executed_func: Any, initial_cache: Dict[Any, Any]) -> Any:
+    async def resolver(  # noqa: C901
+        self,
+        executed_func: Any,
+        initial_cache: Dict[Any, Any],
+    ) -> Any:
         """
         Async resolver.
 
@@ -311,7 +324,7 @@ async def resolver(self, executed_func: Any, initial_cache: Dict[Any, Any]) -> A
         :param initial_cache: cache to build a context if graph was passed.
         :return: dict with resolved kwargs.
         """
-        if getattr(executed_func, "dep_graph", False):
+        if getattr(executed_func, "dep_graph", False):  # noqa: WPS223
             ctx = AsyncResolveContext(executed_func, initial_cache)  # type: ignore
             self.sub_contexts.append(ctx)
             sub_result = await ctx.resolve_kwargs()
@@ -323,6 +336,12 @@ async def resolver(self, executed_func: Any, initial_cache: Dict[Any, Any]) -> A
         elif inspect.isasyncgen(executed_func):
             sub_result = await executed_func.__anext__()  # noqa: WPS609
             self.opened_dependencies.append(executed_func)
+        elif iscontextmanager(executed_func):
+            sub_result = executed_func.__enter__()  # noqa: WPS609
+            self.opened_dependencies.append(executed_func)
+        elif isasynccontextmanager(executed_func):
+            sub_result = await executed_func.__aenter__()  # noqa: WPS609
+            self.opened_dependencies.append(executed_func)
         else:
             sub_result = executed_func
         return sub_result

taskiq_dependencies/dependency.py

@@ -2,8 +2,10 @@
 import uuid
 from typing import (  # noqa: WPS235
     Any,
+    AsyncContextManager,
     AsyncGenerator,
     Callable,
+    ContextManager,
     Coroutine,
     Dict,
     Generator,
@@ -17,6 +19,26 @@
 _T = TypeVar("_T")  # noqa: WPS111
 
 
+@overload
+def Depends(  # noqa: WPS234
+    dependency: Optional[Callable[..., ContextManager[_T]]] = None,
+    *,
+    use_cache: bool = True,
+    kwargs: Optional[Dict[str, Any]] = None,
+) -> _T:  # pragma: no cover
+    ...
+
+
+@overload
+def Depends(  # noqa: WPS234
+    dependency: Optional[Callable[..., AsyncContextManager[_T]]] = None,
+    *,
+    use_cache: bool = True,
+    kwargs: Optional[Dict[str, Any]] = None,
+) -> _T:  # pragma: no cover
+    ...
+
+
 @overload
 def Depends(  # noqa: WPS234
     dependency: Optional[Callable[..., AsyncGenerator[_T, None]]] = None,

taskiq_dependencies/utils.py

@@ -1,5 +1,5 @@
 import inspect
-from typing import Optional
+from typing import Any, AsyncContextManager, ContextManager, Optional, TypeGuard
 
 
 class ParamInfo:
@@ -23,3 +23,31 @@ def __init__(
 
     def __repr__(self) -> str:
         return f"ParamInfo<name={self.name}>"
+
+
+def iscontextmanager(obj: Any) -> TypeGuard[ContextManager[Any]]:
+    """
+    Return true if the object is a sync context manager.
+
+    :param obj: object to check.
+    :return: bool that indicates whether the object is a context manager or not.
+    """
+    if not hasattr(obj, "__enter__"):  # noqa: WPS421
+        return False
+    elif not hasattr(obj, "__exit__"):  # noqa: WPS421
+        return False
+    return True
+
+
+def isasynccontextmanager(obj: Any) -> TypeGuard[AsyncContextManager[Any]]:
+    """
+    Return true if the object is a async context manager.
+
+    :param obj: object to check.
+    :return: bool that indicates whether the object is a async context manager or not.
+    """
+    if not hasattr(obj, "__aenter__"):  # noqa: WPS421
+        return False
+    elif not hasattr(obj, "__aexit__"):  # noqa: WPS421
+        return False
+    return True

tests/test_annotated.py

@@ -1,11 +1,12 @@
 import sys
+from contextlib import asynccontextmanager, contextmanager
 
 import pytest
 
 if sys.version_info < (3, 10):
     pytest.skip("Annotated is available only for python 3.10+", allow_module_level=True)
 
-from typing import Annotated, AsyncGenerator, Generic, Tuple, TypeVar
+from typing import Annotated, AsyncGenerator, Generator, Generic, Tuple, TypeVar
 
 from taskiq_dependencies import DependencyGraph, Depends
 
@@ -53,6 +54,38 @@ def test_func(a: Annotated[MainClass[MyClass], Depends()]) -> MyClass:
     assert isinstance(value, MyClass)
 
 
+@pytest.mark.anyio
+async def test_annotated_gen() -> None:
+    opened = False
+    closed = False
+
+    def my_gen() -> Generator[int, None, None]:
+        nonlocal opened, closed
+        opened = True
+
+        yield 1
+
+        closed = True
+
+    def test_func(dep: Annotated[int, Depends(my_gen)]) -> int:
+        return dep
+
+    with DependencyGraph(target=test_func).sync_ctx() as sctx:
+        value = test_func(**sctx.resolve_kwargs())
+        assert value == 1
+
+    assert opened and closed
+
+    opened = False
+    closed = False
+
+    async with DependencyGraph(target=test_func).async_ctx() as actx:
+        value = test_func(**(await actx.resolve_kwargs()))
+        assert value == 1
+
+    assert opened and closed
+
+
 @pytest.mark.anyio
 async def test_annotated_asyncgen() -> None:
     opened = False
@@ -76,6 +109,65 @@ def test_func(dep: Annotated[int, Depends(my_gen)]) -> int:
     assert opened and closed
 
 
+@pytest.mark.anyio
+async def test_annotated_manager() -> None:
+    opened = False
+    closed = False
+
+    @contextmanager
+    def my_gen() -> Generator[int, None, None]:
+        nonlocal opened, closed
+        opened = True
+
+        try:
+            yield 1
+        finally:
+            closed = True
+
+    def test_func(dep: Annotated[int, Depends(my_gen)]) -> int:
+        return dep
+
+    with DependencyGraph(target=test_func).sync_ctx() as sctx:
+        value = test_func(**sctx.resolve_kwargs())
+        assert value == 1
+
+    assert opened and closed
+
+    opened = False
+    closed = False
+
+    async with DependencyGraph(target=test_func).async_ctx() as actx:
+        value = test_func(**(await actx.resolve_kwargs()))
+        assert value == 1
+
+    assert opened and closed
+
+
+@pytest.mark.anyio
+async def test_annotated_asyncmanager() -> None:
+    opened = False
+    closed = False
+
+    @asynccontextmanager
+    async def my_gen() -> AsyncGenerator[int, None]:
+        nonlocal opened, closed
+        opened = True
+
+        try:
+            yield 1
+        finally:
+            closed = True
+
+    def test_func(dep: Annotated[int, Depends(my_gen)]) -> int:
+        return dep
+
+    async with DependencyGraph(target=test_func).async_ctx() as g:
+        value = test_func(**(await g.resolve_kwargs()))
+        assert value == 1
+
+    assert opened and closed
+
+
 def test_multiple() -> None:
     class TestClass:
         pass

tests/test_graph.py

@@ -1,6 +1,7 @@
 import asyncio
 import re
 import uuid
+from contextlib import asynccontextmanager, contextmanager
 from typing import Any, AsyncGenerator, Generator, Generic, Tuple, TypeVar
 
 import pytest
@@ -111,6 +112,76 @@ def testfunc(a: int = Depends(dep1)) -> int:
     assert closes == 1
 
 
+@pytest.mark.anyio
+async def test_dependency_contextmanager_successful() -> None:
+    """Tests that contextmanagers work as expected."""
+    starts = 0
+    closes = 0
+
+    @contextmanager
+    def dep1() -> Generator[int, None, None]:
+        nonlocal starts  # noqa: WPS420
+        nonlocal closes  # noqa: WPS420
+
+        starts += 1
+
+        try:
+            yield 1
+        finally:
+            closes += 1
+
+    def testfunc(a: int = Depends(dep1)) -> int:
+        return a
+
+    with DependencyGraph(testfunc).sync_ctx({}) as sctx:
+        assert sctx.resolve_kwargs() == {"a": 1}
+        assert starts == 1
+        assert closes == 0
+        starts = 0
+    assert closes == 1
+    closes = 0
+
+    async with DependencyGraph(testfunc).async_ctx({}) as actx:
+        assert await actx.resolve_kwargs() == {"a": 1}
+        assert starts == 1
+        assert closes == 0
+    assert closes == 1
+
+
+@pytest.mark.anyio
+async def test_dependency_async_manager_successful() -> None:
+    """This test checks that async contextmanagers work."""
+    starts = 0
+    closes = 0
+
+    @asynccontextmanager
+    async def dep1() -> AsyncGenerator[int, None]:
+        nonlocal starts  # noqa: WPS420
+        nonlocal closes  # noqa: WPS420
+
+        await asyncio.sleep(0.001)
+        starts += 1
+
+        try:
+            yield 1
+        finally:
+            await asyncio.sleep(0.001)
+            closes += 1
+
+    def testfunc(a: int = Depends(dep1)) -> int:
+        return a
+
+    with DependencyGraph(testfunc).sync_ctx({}) as sctx:
+        with pytest.raises(RuntimeError):
+            assert sctx.resolve_kwargs() == {"a": 1}
+
+    async with DependencyGraph(testfunc).async_ctx({}) as actx:
+        assert await actx.resolve_kwargs() == {"a": 1}
+        assert starts == 1
+        assert closes == 0
+    assert closes == 1
+
+
 @pytest.mark.anyio
 async def test_dependency_subdeps() -> None:
     """Tests how subdependencies work."""