python type hints for decorators using TypeVar/generics and intermediary type variables

user_366221 0
Viewed 265

I'm trying to figure out a nice way of doing type hints for a decorator. Here's a simplified/minimal example. The decorator is called @stringify_input and it converts the inputs of the decorated function to strings:

@stringify_input
def two_things(x: str) -> float:
    return float(x + x)

assert two_things(1) == 11.0


# implementation with no hints:
def stringify_input(fn):
    def inner(thing: int):
        return fn(str(thing))
    return inner

Say I don't know in advance what the decorated function returns, but I know my decorator doesn't change it. currently the undecorated functon returns a float, so the decorated one should too. I've figured out how to do that with TypeVar:

# types check
T = TypeVar("T")

def stringify_input(fn: Callable[[str], T]) -> Callable[[int], T]:
    def inner(thing: int) -> T:
        return fn(str(thing))
    return inner

here's my question: those Callable[....] hints are hard to read, they'd be easier to understand if they had names (my IRL decorators are more complex). so i try this:

StringFn = Callable[[str], T]
IntFn = Callable[[int], T]

def stringify_input(fn: StringFn) -> IntFn:
    def inner(thing: int) -> T:
        return fn(str(thing))   # Returning Any from function declared to return 'T'
    return inner

but that doesn't type check:

Returning Any from function declared to return 'T'

why is mypy unhappy with that? i literally just extracted the two expressions! something about how TypeVar works?

mypy python
1 Answers

The TypeVar has been split to several type variables of same name in different scopes. Declare all expressions to use T of the same scope.

# TypeVar scope equals function scope
def stringify_input(fn: StringFn[T]) -> IntFn[T]:
    def inner(thing: int) -> T:
       ...

While TypeVars are usually declared globally, their substitution happens in a fixed scope. For example, one can use the same TypeVar in two functions to mean separate actual type substitutions.

T = TypeVar('T')

def foo(a: T): ...
def bar(b: T): ...

foo("bar!")   # T -> str
bar(b"foo!")  # T -> bytes

By moving the types outside of the function, there are three scopes with separate meaning for T.

StringFn = Callable[[str], T]    # T1
IntFn = Callable[[int], T]       # T2

def stringify_input(fn: StringFn) -> IntFn:
    def inner(thing: int) -> T:  # T3
        ...
    return inner

Since both StringFn and IntFn have a free TypeVar, they can be parameterised by a TypeVar of another context. Parameterise them by the T of the function scope to show that all of StringFn, IntFn and inner are parameterised by the same T:

def stringify_input(fn: StringFn[T]) -> IntFn[T]:
    def inner(thing: int) -> T:
       ...
user_53499160
Related