Fix stubgen output for async context managers - #21870
Conversation
|
Thanks for picking this up — I filed #21869. I applied this patch to a v2.3.0 checkout and tested it. It fixes the issue, and the output is not just error-free but type-equivalent to the source. Running the same consumer against the generated stub and against the original module:
So dropping I also checked the failure mode this design could have — @asynccontextmanager
async def ctx() -> AsyncIterator[int]: ... # -> async dropped
async def plain_coro() -> int: ... # -> async kept
@contextmanager
def sync_ctx() -> Iterator[int]: ... # -> unchanged
class A:
@asynccontextmanager
async def actx(self) -> AsyncGenerator[str, None]: ... # -> dropped
async def method_coro(self) -> int: ... # -> kept
@staticmethod
@asynccontextmanager
async def sctx() -> AsyncIterator[int]: ... # -> dropped, both decorators keptAll correct. That holds because the flag is set in One suggestion: that non-interference isn't pinned by a test. Since it's the single property this approach could regress on, and the harness compares whole-file output, it costs four lines folded into the case you already added rather than a new one: [case testAsyncContextManager]
from collections.abc import AsyncGenerator, AsyncIterator
from contextlib import asynccontextmanager
@asynccontextmanager
async def ctx() -> AsyncIterator[int]:
yield 1
+async def coro() -> int:
+ return 1
+
class A:
@asynccontextmanager
async def ctx(self) -> AsyncGenerator[str, None]:
yield "value"
+
+ async def coro(self) -> int:
+ return 1
[out]
from collections.abc import AsyncGenerator, AsyncIterator
from contextlib import asynccontextmanager
@asynccontextmanager
def ctx() -> AsyncIterator[int]: ...
+async def coro() -> int: ...
class A:
@asynccontextmanager
def ctx(self) -> AsyncGenerator[str, None]: ...
+ async def coro(self) -> int: ...I generated that |
Fixes #21869
Stubgen preserved both the asynccontextmanager decorator and the async keyword, producing a stub that mypy rejects because a stub body cannot establish an async generator.
This tracks when contextlib.asynccontextmanager is present and omits async only from the emitted stub signature. The decorator, return annotation, and source AST behavior remain otherwise unchanged. Regression coverage includes both a module function and a method.
Tests: