How do I type correctly a wrapper function?

Assume following declarations:

from typing import Callable, TypeVar
T = TypeVar('T')

def wrapper(fn: Callable[..., T]) -> Callable[..., T]:
    ...

def identity(a: T) -> T:
    ...

@wrapper
def w_wrapped(a: T) -> T:
    ...

@identity
def i_wrapped(a: T) -> T:
    ...

The two annotated functions can be used like this:

def apply(fn: Callable[[str], int]) -> int:
    # types fine:
    val1 = fn(i_wrapped(''))
    # mypy complains: error: Argument 1 has incompatible type "T"; expected "str"
    val2 = fn(w_wrapped(''))
    return val1 + val2

What's wrong with the Callable type? I can use Callable[..., Any] instead of Callable[..., T] in the wrapper declaration. But I feel like this partially defeats the purpose, I would like to declare that when you use the wrapper with str, the result would be str, not just anything. There may be other workarounds too, but is this mypy limitation or my misunderstanding?

Solution

I think there may be two things going on here:

Firstly, a mypy bug, see this and this

Secondly, Callable[..., T] is too loose. Specifically, there's no connection between its argument and its return value. As a result, with @wrapper, w_wrapped becomes a Callable[..., T], with no constraint on its argument, and the output of w_wrapped('') is a unbound T, which can't be passed to fn which expects a str.

You have a number of options depending on your use case, including

def wrapper(fn: Callable[[U], T]) -> Callable[[U], T]: for U = TypeVar('U'), though I believe the mypy bugs stops this working. U could also be T
def wrapper(fn: C) -> C: for C = TypeVar('C', bound=Callable). This doesn't have the problem of Any because you're bounding on Callable so you retain the type signature. However, it will limit your implementation of wrapper, short of type: ignore