pythonpython-asynciopython-typingcontextmanagerpyright

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


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:

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


Solution

  • 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 @asynccontextmanager:
    asynccontextmanager(func: Callable[_P, AsyncIterator[_T_co]]) -> Callable[_P, _AsyncGeneratorContextManager[_T_co]]

    That's a lot to parse but it says that the asynccontextmanager takes a function which returns AsyncIterator and transforms it into a new function that returns AsyncContextManager. The generic types _P and _T_co are preserved as well.

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

    enter image description here