Fix contextlib.asynccontextmanager to work with coroutine functions

sobolevn · Sep 23, 2023 · Sep 23, 2023 · Sep 23, 2023 · Sep 23, 2023 · Sep 23, 2023
commit bf423b38704ad93d39a1b4b45625f43d74c62d89
diff --git a/stdlib/contextlib.pyi b/stdlib/contextlib.pyi
@@ -2,7 +2,7 @@ import abc
 import sys
 from _typeshed import FileDescriptorOrPath, Unused
 from abc import abstractmethod
-from collections.abc import AsyncGenerator, AsyncIterator, Awaitable, Callable, Generator, Iterator
+from collections.abc import AsyncGenerator, AsyncGenerator, Awaitable, Callable, Coroutine, Generator, Iterator
 from types import TracebackType
 from typing import IO, Any, Generic, Protocol, TypeVar, overload, runtime_checkable
 from typing_extensions import ParamSpec, Self, TypeAlias
@@ -84,7 +84,7 @@ if sys.version_info >= (3, 10):
     class _AsyncGeneratorContextManager(AbstractAsyncContextManager[_T_co], AsyncContextDecorator, Generic[_T_co]):
         # __init__ and these attributes are actually defined in the base class _GeneratorContextManagerBase,
         # which is more trouble than it's worth to include in the stub (see #6676)
-        def __init__(self, func: Callable[..., AsyncIterator[_T_co]], args: tuple[Any, ...], kwds: dict[str, Any]) -> None: ...
+        def __init__(self, func: Callable[..., AsyncGenerator[_T_co, Any]], args: tuple[Any, ...], kwds: dict[str, Any]) -> None: ...
         gen: AsyncGenerator[_T_co, Any]
         func: Callable[..., AsyncGenerator[_T_co, Any]]
         args: tuple[Any, ...]
@@ -95,7 +95,7 @@ if sys.version_info >= (3, 10):
 
 else:
     class _AsyncGeneratorContextManager(AbstractAsyncContextManager[_T_co], Generic[_T_co]):
-        def __init__(self, func: Callable[..., AsyncIterator[_T_co]], args: tuple[Any, ...], kwds: dict[str, Any]) -> None: ...
+        def __init__(self, func: Callable[..., AsyncGenerator[_T_co, Any]], args: tuple[Any, ...], kwds: dict[str, Any]) -> None: ...
         gen: AsyncGenerator[_T_co, Any]
         func: Callable[..., AsyncGenerator[_T_co, Any]]
         args: tuple[Any, ...]
@@ -104,7 +104,10 @@ else:
             self, typ: type[BaseException] | None, value: BaseException | None, traceback: TracebackType | None
         ) -> bool | None: ...
 
-def asynccontextmanager(func: Callable[_P, AsyncIterator[_T_co]]) -> Callable[_P, _AsyncGeneratorContextManager[_T_co]]: ...
+@overload
+def asynccontextmanager(func: Callable[_P, Coroutine[Any, Any, AsyncGenerator[_T_co, Any]]]) -> Callable[_P, _AsyncGeneratorContextManager[_T_co]]: ...
+@overload
+def asynccontextmanager(func: Callable[_P, AsyncGenerator[_T_co, Any]]) -> Callable[_P, _AsyncGeneratorContextManager[_T_co]]: ...
 
 class _SupportsClose(Protocol):
     def close(self) -> object: ...

diff --git a/test_cases/stdlib/check_contextlib.py b/test_cases/stdlib/check_contextlib.py
@@ -1,6 +1,7 @@
 from __future__ import annotations
 
-from contextlib import ExitStack
+from contextlib import ExitStack, asynccontextmanager
+from typing import AsyncIterator, Awaitable
 from typing_extensions import assert_type
 
 
@@ -18,3 +19,17 @@ class Thing(ExitStack):
     assert_type(cm, ExitStack)
 with thing as cm2:
     assert_type(cm2, Thing)
+
+
+@asynccontextmanager
+async def async_context() -> AsyncGenerator[str, None]:
+    yield 'example'
+
+
+async def async_gen() -> AsyncGenerator[str, None]:
+    yield 'async gen'
+
+
+@asynccontextmanager
+def async_cm_func() -> AsyncGenerator[str, None]:
+    return async_gen()