如何使用 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()

的结果，这一点应该是显而易见的。