rankN type equivalent for mypy in python

Question

In Haskell we can use rankN types like so:

rankN :: (forall n. Num n => n -> n) -> (Int, Double)
rankN f = (f 1, f 1.0)

Is the same thing possible in python with mypy?

I tried the following code in python 3.10.2 with mypy 1.7.1:

I = TypeVar("I", int, float)


def rankN(f: Callable[[I], I]) -> tuple[int, float]:
    return (f(1), f(1.0))

This produces the following errors, implying that f is specializing to float:

Incompatible return value type (got "tuple[float, float]", expected "tuple[int, float]")  [return-value]
Argument 1 has incompatible type "float"; expected "int"  [arg-type]

I'm not necessarily expecting this to work since the magic syntax in the Haskell case is the nested forall, but I don't know if there is a similar way to convey this to mypy if it is possible at all.

Answer 1

I’m not familiar with mypy, but my guess is that you can (and must) represent this as a protocol with a generic method, in order to scope the type variable to be per method call, rather than per invocation of rankN.

from typing import Protocol, TypeVar

class UnaryNumeric(Protocol):
    I = TypeVar("I", int, float)
    def __call__(self, input: I) -> I:
        pass

def rankN(f: UnaryNumeric) -> tuple[int, float]:
    return (f(1), f(1.0))

Now rankN(lambda x: abs(x)) is accepted, while rankN(lambda x: str(x)) is rejected. I was surprised that rankN(lambda x: x / 2) is also accepted, since int isn’t really a subtype of float, but I guess it floats like a duck.