python python-3.x generator python-typing

How to annotate a decorator on a generator using type hinting?

I'm working with generators-as-coroutines as described in the excellent triplet of presentations by David Beazley (at http://www.dabeaz.com/coroutines/) and I can't figure out how to type the decorator consumer. Here's what I have so far:

from typing import Any, Callable, Generator, Iterable

ArbGenerator = Generator[Any, Any, Any]

def consumer(fn: ❓) -> ❓:
    @wraps(fn)
    def start(*args: Any) -> ArbGenerator:
        c = fn(*args)
        c.send(None)
        return c
return start

Example of use, kind of abridged:

@consumer
def identity(target: ArbGenerator) -> ArbGenerator:
    while True:
        item = yield
        target.send(item)

@consumer
def logeach(label: Any, target: ArbGenerator) -> ArbGenerator:
    while True:
        item = yield
        print(label, item)
        target.send(item)

pipeline = identity(logeach("EXAMPLE", some_coroutine_sink()))

Bold ❓ marks where I'm unsure - and I'm also unsure about the type I defined ArbGenerator. (Problem is, without the (decorator) function consumer itself typed, I'm not sure mypy is analyzing any generator function with that decorator so that's why I'm unsure about ArbGenerator.)

I'm interested in the tightest type, something better than Any, so that when I compose chains of these coroutines mypy will give me nice warnings if the chain isn't set up right.

(Python 3.5 if it matters.)

Solution

As a more specific way, here are few things you can do:

Use Callable type instead of question marks.
Use typing.Coroutine for targets and drop the ArbGenerator.
Coroutines return a generator and the return type could be a Generator or one of its supertypes

The reason that you should use callable instead of question marks is that fn is supposed to be a callable object at first and that's why you're wrapping it with a decorator. The Coroutine will be created after calling the object and the return type is/should be obviously a callable object as well.

from typing import Any, Callable,Generator, Coroutine
from functools import wraps


def consumer(fn: Callable) -> Callable:
    @wraps(fn)
    def start(*args: Any) -> Coroutine:
        c = fn(*args)  # type: Coroutine
        c.send(None)
        return c
    return start


@consumer
def identity(target: Coroutine) -> Generator:
    while True:
        item = yield
        target.send(item)

@consumer
def logeach(label: Any, target: Coroutine) -> Generator:
    while True:
        item = yield
        print(label, item)
        target.send(item)

Note: As it's also mentioned in documentation, if you want to use a more precise syntax for annotating the generator type you can use the following syntax:

Generator[YieldType, SendType, ReturnType]