What happens if overloaded higher-order functions (e.g. decorators) have overloaded arguments?

[wlritchi]

Function overloads can be used to inform the type checker that there are multiple valid parameter typings for the function, possibly with distinct return types for each. This can be applied to higher-order functions, such as decorators. For example:

@overload
def flip_args(
    func: Callable[[str, int], _R],
) -> Callable[[int, str], _R]:
    pass

@overload
def flip_args(
    func: Callable[[int, int], _R],
) -> Callable[[int, int], _R]:
    pass

def flip_args(func: Callable[..., _R]) -> Callable[..., _R]:
    """Contrived example, a real flip_args would use TypeVars"""

    def wrapper(p1: int, p2: str | int) -> _R:
        return func(p2, p1)

    return wrapper

This example decorator transforms two-parameter functions into new functions that accept the same positional parameters, in the opposite order. For the purposes of this example, we use overloads rather than TypeVars. This imposes unnecessary constraints on the input types, but it can still decorate a few example functions:

def takes_string_and_int(p1: str, p2: int) -> None:
    print(f"A {p1} has {p2} vertices")

def takes_two_ints(p1: int, p2: int) -> None:
    print(f"A {p1}-sided die has {p2} vertices")

takes_string_and_int_flipped = flip_args(takes_string_and_int)
takes_two_ints_flipped = flip_args(takes_two_ints)

But what happens if the input function is itself overloaded, and the possible overloads of the input function match different overloads of the decorator? At first, it seems to work:

@overload
def takes_something_and_int(p1: str, p2: int) -> None:
    pass

@overload
def takes_something_and_int(p1: int, p2: int) -> None:
    pass

def takes_something_and_int(p1: str | int, p2: int) -> None:
    if isinstance(p1, str):
        takes_string_and_int(p1, p2)
    else:
        takes_two_ints(p1, p2)

takes_something_and_int_flipped = flip_args(takes_something_and_int)

However, existing typecheckers select the first overload of the decorator that is able to match any overload of the input function. As a result, they infer only one type for takes_something_and_int_flipped, with no overloading: (int, str) -> int. I would have expected the resulting type to be equivalent to overloading (int, str) -> int and (int, int) -> int.

takes_string_and_int("cube", 8)
takes_two_ints(20, 12)
takes_something_and_int("cube", 8)
takes_something_and_int(20, 12)

takes_string_and_int_flipped(8, "cube")
takes_two_ints_flipped(12, 20)
takes_something_and_int_flipped(8, "cube")
takes_something_and_int_flipped(12, 20)  # does not typecheck

Should a decorator's overloading "collapse" when it is applied, forcing the typechecker to commit to exactly one type for the resulting decorated function? Or should it distribute over every overload that matches any overload of the input function?

[erictraut]

Let's consider the implications of "distributing over every overload that matches the input function". If a call expression targets a function that accepts a single callback parameter and there are n overloads that match that callback parameter, the type checker will need to evaluate the call expression n times, once for each overload. What would the result of this call be? A union of all of the evaluated results? That's probably not what you want. You probably want the results to somehow be combined back into another overload, but that operation (as opposed to a union) is not well defined in the type system. Combining the results back into an overload might be possible in the simple case where the function returns a single callable type, but what if it doesn't? What if it returns a wrapper object? A union of types?

Another consideration... What if the function has m parameters that are all callbacks, and each one matches n overloads? Now the type checker would need to analyze the function n ** m times, once for each combination of the overloads, and combine all of the results. Exponential explosions such as this are generally a bad idea in static analyzers, and it's best to avoid creating constructs that result in such explosions.

I think what you're suggesting here could work in the most simple and contrived cases, but it breaks down as soon as you start to generalize it. That is likely why all current type checkers simply choose the first overload that matches.