Search code examples
pythongeneratormypypython-typing

Proper type annotation of Python functions with yield


After reading Eli Bendersky's article on implementing state machines via Python coroutines I wanted to...

  • see his example run under Python3
  • and also add the appropriate type annotations for the generators

I succeeded in doing the first part (but without using async defs or yield froms, I basically just ported the code - so any improvements there are most welcome).

But I need some help with the type annotations of the coroutines:

#!/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)

This runs properly...

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

...and also typechecks:

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

But I am pretty sure I can do better than just use the Generator base class in the type specs (just as I did for the Callable). I know it takes 3 type parameters (Generator[A,B,C]), but I am not sure how exactly they'd be specified here.

Any help most welcome.


Solution

  • I figured out the answer on my own.

    I searched, but found no documentation for the 3 type parameters of Generator in the official typing documentation for Python 3.5.2 - beyond a truly cryptic mention of...

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

    Luckily, the original PEP484 (that started all this) was far more helpful:

    "The return type of generator functions can be annotated by the generic type Generator[yield_type, send_type, return_type] provided by typing.py module:

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

    Based on this, I was able to annotate my Generators, and saw mypy confirm my assignments:

    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:
        ...
    

    I tested my assignments by e.g. swaping the order of the types - and then as expected, mypy complained and asked for the proper ones (as seen above).

    The complete code is accessible from here.

    I will leave the question open for a couple of days, in case anyone wants to chime in - especially in terms of using the new coroutine styles of Python 3.5 (async def, etc) - I would appreciate a hint on exactly how they'd be used here.