pythonpython-typinghigher-order-functions

How to annotate interdependent functional parameter and return types with a varying number of parameters in Python (< 3.10)?


EDIT NOTE 1: By now I've found PEP 612 which solves this issue - starting with Python 3.10 - by introducing typing.ParamSpec. So this question is specifically about Python 3.9 or earlier.

EDIT NOTE 2: The original example was too narrow in that it had the return type match the parameter type exactly, but the question is in fact about a more generic case, in which the parameter signatures of the functions are the same but the return type is different. (More generic solutions that allow also for different parameter signatures are also welcome.)

The question: I'm using a transforming function which receives a function as an argument, and returns another function as its result. The passed function can have any number and types of parameters (for simplicity we can stick to positional parameters), and the returned function is invoked the same way as the original one (with the same number and types of parameters), and has a return type which depends on the return type of the passed function (but is not necessarily equal to it; for example it can be a tuple containing a value of the original return type and another value of some given type).

How can I annotate the transforming function in a way that will reflect the dependency of the signatures of the passed and returned functions?

A simple example:

from typing import Callable


def transform(original_function: Callable) -> Callable:
    def new_function(*args):
        extra_payload = <do some calculation>
        return original_function(*args), extra_payload
    return new_function


def f(x: int, s: str) -> bool:
    ...


f(3, 'abc')  # this is a valid call
f('abc', 3)  # PyCharm warns about wrong argument types

# The goal is to have a warning on the following line:
transform(f)('abc', 3)  # no warning with the mere Callable annotations above

Is there any way to let PyCharm know transform(f) has the same parameter signature as f?

If there was a fixed number of parameters for the transformed functions I could do it, e.g. (assuming two parameters):

from typing import TypeVar, Callable, Tuple


X = TypeVar('X')
Y = TypeVar('Y')
Z = TypeVar('Z')


def transform(original_function: Callable[[X, Y], Z]) -> Callable[[X, Y], Tuple[Z, <some type>]]:
   ...

, but my transform function is more generic than that, and I use it on functions with varying numbers of arguments, and then I don't know how to specify that in the first argument to Callable (where there's [X, Y] above).

(How) can this be done before the introduction of typing.ParamSpec?


Solution

  • It seems this can't be done elegantly and fully before Python 3.10, and that exactly for this reason Python 3.10 introduces ParamSpec (as per MisterMiyagi's comment).

    However, one can get close by using overloads (see relevant part of PEP 484) for different numbers of parameters:

    from typing import TypeVar, overload, Callable, Tuple
    
    
    T = TypeVar('T')
    T1 = TypeVar('T1')
    T2 = TypeVar('T2')
    T3 = TypeVar('T3')
    
    Payload = <some type>
    
    
    @overload
    def transform(function: Callable[[], T]) -> Callable[[], Tuple[T, Payload]]:
        ...
    
    
    @overload
    def transform(function: Callable[[T1], T]) -> Callable[[T1], Tuple[T, Payload]]:
        ...
    
    
    @overload
    def transform(function: Callable[[T1, T2], T]) -> Callable[[T1, T2], Tuple[T, Payload]]:
        ...
    
    
    @overload
    def transform(function: Callable[[T1, T2, T3], T]) -> Callable[[T1, T2, T3], Tuple[T, Payload]]:
        ...
    
    
    def transform(original_function):
        def new_function(*args):
            extra_payload = <do some calculation>
            return original_function(*args), extra_payload
        return new_function