Typing function when decorator change generic return type

Viewed 764

This is similar to Typing function when decorator change return type but this time with a generic return type:

from typing import Generic, TypeVar, Generic, Callable, Any, cast

T = TypeVar('T')

class Container(Generic[T]):
    def __init__(self, item: T):
        self.item: T = item

def my_decorator(f: Callable[..., T]) -> Callable[..., Container[T]]:

    def wrapper(*args: Any, **kwargs: Any) -> Container[T]:
        return Container(f(*args, **kwargs))

    return cast(Callable[..., Container[T]], wrapper)

@my_decorator
def my_func(i: int, s: str) -> bool: ...

reveal_type(my_func) # Revealed type is 'def (*Any, **Any) -> file.Container[builtins.bool*]

Which mypy sorcery is required to keep the argument types intact for my_func?

Using typing.Protocol looks promising, but I don't see how to make it work.

1 Answers

Using Callable[..., T] is currently the best way to annotate this.

PEP 612 introduces ParamSpec, which can be used similar to a TypeVar and will solve your problem. It's currently planned for Python 3.10 and will be supported for older versions using typing_extensions

There you would write:

T = TypeVar('T')
P = ParamSpec('P')

def my_decorator(f: Callable[P, T]) -> Callable[P, Container[T]]:

    def wrapper(*args: Any, **kwargs: Any) -> Container[T]:
        return Container(f(*args, **kwargs))

    return cast(Callable[P, Container[T]], wrapper)

mypy support for PEP 612 isn't finished yet: https://github.com/python/mypy/issues/8645. Same with pytype (Google's python typechecker).

pyright (Microsoft's python typechecker) and pyre (facebook's python typechecker) do already support PEP 612

Related