如何使用 AsyncGenerator 和 AsyncContextManager 正确指定类型提示

How to correctly specify type hints with AsyncGenerator and AsyncContextManager

考虑以下代码

import contextlib
import abc
import asyncio

from typing import AsyncContextManager, AsyncGenerator, AsyncIterator


class Base:

    @abc.abstractmethod
    async def subscribe(self) -> AsyncContextManager[AsyncGenerator[int, None]]:
        pass

class Impl1(Base):

    @contextlib.asynccontextmanager
    async def subscribe(self) ->  AsyncIterator[ AsyncGenerator[int, None] ]: <-- mypy error here

        async def _generator():
            for i in range(5):
                await asyncio.sleep(1)
                yield i
                    
        yield _generator()

对于Impl1.subscribe mypy给出错误

Signature of "subscribe" incompatible with supertype "Base"

在上述情况下指定类型提示的正确方法是什么?还是mypy这里错了?

我刚好遇到同样的问题,当天就发现了这个问题,但也很快找到了答案。

您需要从抽象方法中删除 async

为了解释原因,我将案例简化为一个简单的异步迭代器:

@abc.abstractmethod
async def foo(self) -> AsyncIterator[int]:
    pass

async def v1(self) -> AsyncIterator[int]:
    yield 0

async def v2(self) -> AsyncIterator[int]:
    return v1()

如果比较 v1 和 v2,您会发现函数签名看起来相同,但实际上它们做的事情却截然不同。 v2兼容抽象方法,v1不兼容

当您添加 async 关键字时,mypy 将函数的 return 类型推断为 Coroutine。但是,如果您还放入 yield,它会推断 return 类型为 AsyncIterator:

reveal_type(foo)
# -> typing.Coroutine[Any, Any, typing.AsyncIterator[builtins.int]]
reveal_type(v1)
# -> typing.AsyncIterator[builtins.int]
reveal_type(v2)
# -> typing.Coroutine[Any, Any, typing.AsyncIterator[builtins.int]]

正如你所见,抽象方法中缺少yield意味着这被推断为Coroutine[..., AsyncIterator[int]]。换句话说,像 async for i in await v2():.

这样使用的函数

删除 async:

@abc.abstractmethod
def foo(self) -> AsyncIterator[int]:
    pass
reveal_type(foo)
# -> typing.AsyncIterator[builtins.int]

我们看到 return 类型现在是 AsyncIterator 并且现在与 v1 兼容,而不是 v2。换句话说,像 async for i in v1():

这样使用的函数

你也可以看到这和v1基本一样:

def v3(self) -> AsyncIterator[int]:
    return v1()

虽然语法不同,但 v3 和 v1 都是将 return 和 AsyncIterator 调用的函数,这应该是显而易见的,因为我们实际上是 returning 结果v1().