如何使用 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.subscribemypy 给出错误

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 会将函数的返回类型推断为 a Coroutine。但是，如果您还输入了 a 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]]

正如您所看到的，抽象方法中缺少 ayield意味着这被推断为 a 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 都是在调用时返回 an 的函数AsyncIterator，考虑到我们实际上返回的是 的结果，这一点应该是显而易见的v1()。