What is the proper way to type hint the return value of an @asynccontextmanager?

Question

What is the proper way to add type hints for the return of a function with the @asynccontextmanager decorator? Here are two attempts I made that both fail.

from contextlib import asynccontextmanager
from typing import AsyncContextManager


async def caller():
    async with return_str() as i:
        print(i)

    async with return_AsyncContextManager() as j:
        print(j)


@asynccontextmanager
async def return_str() -> str:
    yield "hello"


@asynccontextmanager
async def return_AsyncContextManager() -> AsyncContextManager[str]:
    yield "world"

For both i and j Pylance in vscode shows type Any. Other ideas I've considered:

I thought maybe I could pass in the type info to the decorator itself (like @asynccontextmanager(cls=str), but I can't find any examples of that, or any description of the args I could pass in.
async with return_str() as i: # type: str doesn't work either. Even if it did, I'd rather hint at the function definition, not at every invocation. Type comments are not very DRY.
I tried to create an AsyncContextManager class with __aenter()__ and __aexit()__ functions, but was not successful. I would fall back to that if it worked, but I'd prefer to make the decorator work because it's much more concise.

Here's a screencap of me hovering the return_AsyncContextManager() function, and showing the Pylance popup saying it returns AsyncContextManager[_T]

user_596242 · Answer

You have to hint AsyncIterator as the return type, like this:

@asynccontextmanager
async def my_acm() -> AsyncIterator[str]:
    yield "this works"


async def caller():
    async with my_acm() as val:
        print(val)

This is because the yield keyword is used to create generators. Consider:

def a_number_generator():
    for x in range(10):  # type: int
        yield x

g = a_number_generator() # g is type Generator[int]

This makes sense given the type hints for @asynccontextgenerator:
asynccontextmanager(func: Callable[..., AsyncIterator[_T]]) -> Callable[..., AsyncContextManager[_T]]

That's a lot to parse but it says that the asynccontextgenerator takes a function which returns AsyncIterator and transforms it into a new function that returns AsyncContextManager A generic type _T is preserved as well.

Here is a screenshot showing the type hint transferring into the caller function.

What is the proper way to type hint the return value of an @asynccontextmanager?

1 Answers