Python函数中使用yield的正确类型注释

Question

Python函数中使用yield的正确类型注释

pythongeneratorcoroutinestatic-typingmypy

102

在阅读了Eli Bendersky的文章关于通过Python协程实现状态机后，我想...

看到他的示例在Python3下运行
并且为生成器添加适当的类型注释

我已经成功完成了第一部分(但没有使用async def或yield from，我基本上只是将代码进行移植 - 所以任何改进都非常欢迎)。

但是我需要一些帮助来为协程添加类型注释:

#!/usr/bin/env python3

from typing import Callable, Generator

def unwrap_protocol(header: int=0x61,
                    footer: int=0x62,
                    dle: int=0xAB,
                    after_dle_func: Callable[[int], int]=lambda x: x,
                    target: Generator=None) -> Generator:
    """ Simplified protocol unwrapping co-routine."""
    #
    # Outer loop looking for a frame header
    #
    while True:
        byte = (yield)
        frame = []  # type: List[int]

        if byte == header:
            #
            # Capture the full frame
            #
            while True:
                byte = (yield)
                if byte == footer:
                    target.send(frame)
                    break
                elif byte == dle:
                    byte = (yield)
                    frame.append(after_dle_func(byte))
                else:
                    frame.append(byte)


def frame_receiver() -> Generator:
    """ A simple co-routine "sink" for receiving full frames."""
    while True:
        frame = (yield)
        print('Got frame:', ''.join('%02x' % x for x in frame))

bytestream = bytes(
    bytearray((0x70, 0x24,
               0x61, 0x99, 0xAF, 0xD1, 0x62,
               0x56, 0x62,
               0x61, 0xAB, 0xAB, 0x14, 0x62,
               0x7)))

frame_consumer = frame_receiver()
next(frame_consumer)  # Get to the yield

unwrapper = unwrap_protocol(target=frame_consumer)
next(unwrapper)  # Get to the yield

for byte in bytestream:
    unwrapper.send(byte)

这段代码正常运行...

$ ./decoder.py 
Got frame: 99afd1
Got frame: ab14

...并进行类型检查：

$ mypy --disallow-untyped-defs decoder.py 
$

但我相信我可以做得比仅使用 Generator 基类在类型规范中更好（就像我为 Callable 所做的那样）。我知道它需要 3 个类型参数（Generator[A,B,C]），但我不确定它们在这里应该如何指定。

任何帮助都非常欢迎。

- ttsiodras

3个回答

75

如果您有一个使用yield的简单函数，那么您可以使用Iterator类型来注释其结果，而不是使用Generator。

from collections.abc import Iterator  # Python >=3.9

def count_up() -> Iterator[int]:
    for x in range(10):
        yield x

如果函数是 async，你需要使用 AsyncIterator 替代：

from collections.abc import AsyncIterator  # Python >=3.5

async def count_up() -> AsyncIterator[int]:
    for x in range(10):
        yield x

在Python 3.9之前，你必须以不同的方式导入Iterator模块。

from typing import Iterator  # Python <3.9

- David Foster

9

如果生成器只产生值而不接收或返回值，那么这是适用的。 - Peter Bašista

从Python 3.7开始，不再需要from typing import Iterator，而是使用collections.abc.Iterator替代Iterator。 - Wok

不过，人们必须要 import collections.abc。 - Wok

10

截至撰写本文时，Python 文档明确提到了在异步情况下该怎么做（非异步示例已在被接受的答案中提到）。

引用自该文档：

async def echo_round() -> AsyncGenerator[int, float]:
    sent = yield 0
    while sent >= 0.0:
        rounded = await round(sent)
        sent = yield rounded

（第一个参数是yield类型，第二个参数是send类型），或者对于简单情况（其中send类型为None）

async def infinite_stream(start: int) -> AsyncIterator[int]:
    while True:
        yield start
        start = await increment(start)

- phispi

网页内容由stack overflow 提供, 点击上面的

可以查看英文原文，
原文链接

- ttsiodras · Accepted Answer

我自己找到了答案。

我搜索了Python 3.5.2官方类型注释文档，但并没有找到关于Generator的三个类型参数的任何说明，除了一个非常晦涩的提及...

class typing.Generator(Iterator[T_co], Generic[T_co, T_contra, V_co])

幸运的是，引发所有这些的原始PEP484更有帮助：

“生成器函数的返回类型可以由 typing.py 模块提供的通用类型 Generator[yield_type, send_type, return_type] 进行注释:”

def echo_round() -> Generator[int, float, str]:
    res = yield
    while res:
        res = yield round(res)
    return 'OK'

基于此，我能够为我的生成器添加注释，并且看到 mypy 确认了我的分配：

from typing import Callable, Generator

# A protocol decoder:
#
# - yields Nothing
# - expects ints to be `send` in his yield waits
# - and doesn't return anything.
ProtocolDecodingCoroutine = Generator[None, int, None]

# A frame consumer (passed as an argument to a protocol decoder):
#
# - yields Nothing
# - expects List[int] to be `send` in his waiting yields
# - and doesn't return anything.
FrameConsumerCoroutine = Generator[None, List[int], None]


def unwrap_protocol(header: int=0x61,
                    footer: int=0x62,
                    dle :int=0xAB,
                    after_dle_func: Callable[[int], int]=lambda x: x,
                    target: FrameConsumerCoroutine=None) -> ProtocolDecodingCoroutine:
    ...

def frame_receiver() -> FrameConsumerCoroutine:
    ...

我通过交换类型的顺序等方式测试了我的作业 - 正如预期的那样，mypy会抱怨并要求使用正确的类型（如上所示）。

完整的代码可从此处访问。

我将在未来几天内保持问题开放，以防有人想加入 - 特别是在使用Python 3.5的新协程样式（async def等方面）方面 - 我希望能得到关于它们在这里的确切用法的提示。