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

Question

考虑下面的代码

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 错了？

Answer 1

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

您需要从抽象方法中删除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 将函数的返回类型推断为Coroutine。但是，如果你还放了一个yield，它就会推断出返回类型是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]

我们看到返回类型现在是 AsyncIterator 并且现在与 v1 兼容，而不是 v2。换句话说，像async for i in v1():这样使用的函数

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

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

虽然语法不同，但 v3 和 v1 都是函数，调用时将返回 AsyncIterator，考虑到我们实际上返回的是 v1() 的结果，这一点应该很明显。