diff --git a/mypy/stubgen.py b/mypy/stubgen.py index fe90c3d256b7..a692d561ba62 100755 --- a/mypy/stubgen.py +++ b/mypy/stubgen.py @@ -530,6 +530,7 @@ def __init__( ) -> None: super().__init__(_all_, include_private, export_less, include_docstrings) self._decorators: list[str] = [] + self._has_async_contextmanager = False # Stack of defined variables (per scope). self._vars: list[list[str]] = [[]] # What was generated previously in the stub file. @@ -748,7 +749,10 @@ def visit_func_def(self, o: FuncDef) -> None: sigs = self.get_signatures(default_sig, self.sig_generators, ctx) for output in self.format_func_def( - sigs, is_coroutine=o.is_coroutine, decorators=self._decorators, docstring=ctx.docstring + sigs, + is_coroutine=o.is_coroutine and not self._has_async_contextmanager, + decorators=self._decorators, + docstring=ctx.docstring, ): self.add(output + "\n") @@ -809,6 +813,10 @@ def process_decorator(self, o: Decorator) -> None: o.func.is_overload = True elif qualname.endswith((".setter", ".deleter")): self.add_decorator(qualname, require_name=False) + elif fullname == "contextlib.asynccontextmanager": + p = AliasPrinter(self) + self._decorators.append(f"@{decorator.accept(p)}") + self._has_async_contextmanager = True elif fullname in DATACLASS_TRANSFORM_NAMES: p = AliasPrinter(self) self._decorators.append(f"@{decorator.accept(p)}") @@ -1365,6 +1373,7 @@ def add_decorator(self, name: str, require_name: bool = False) -> None: def clear_decorators(self) -> None: self._decorators.clear() + self._has_async_contextmanager = False def is_private_member(self, fullname: str) -> bool: parts = fullname.split(".") diff --git a/test-data/unit/stubgen.test b/test-data/unit/stubgen.test index 0c8b74ecf29a..ab4caf2fd2fc 100644 --- a/test-data/unit/stubgen.test +++ b/test-data/unit/stubgen.test @@ -355,6 +355,29 @@ def foo(x) -> None: ... def bar(x) -> None: ... def foo_bar(x) -> None: ... +[case testAsyncContextManager] +from collections.abc import AsyncGenerator, AsyncIterator +from contextlib import asynccontextmanager + +@asynccontextmanager +async def ctx() -> AsyncIterator[int]: + yield 1 + +class A: + @asynccontextmanager + async def ctx(self) -> AsyncGenerator[str, None]: + yield "value" +[out] +from collections.abc import AsyncGenerator, AsyncIterator +from contextlib import asynccontextmanager + +@asynccontextmanager +def ctx() -> AsyncIterator[int]: ... + +class A: + @asynccontextmanager + def ctx(self) -> AsyncGenerator[str, None]: ... + [case testMultipleAssignment] x, y = 1, 2 [out]