Support a method of copying function signature

[KotlinIsland]

Feature
Capability to copy a functions signature, related to typing.ParamSpec.

def foo(a: str) -> None:
    ...


# pseudo code
def bar(*args: foo.args, **kwargs: foo.kwargs) -> foo.return_type:
    return foo(*args, **kwargs)

Pitch
If I want to forward a call to another method in a type safe way it requires duplicating type definitions which can rapidly become unwieldy and error prone(especially if the function is in a third party module). I would love a way to just say that one functions signature is the same as another functions signature (with support for typing.Concatenate)

Maybe

extend functionality of ParamSpec to be generic to a Callable?

def bar(*args: ParamSpec[foo].args, **kwargs: ParamSpec[foo].kwargs) -> None:
    return foo(*args, **kwargs)

• basedmypy Functionality to copy and enforce signatures of function declarations KotlinIsland/basedmypy#257

[erictraut]

Perhaps one of these approaches would meet your needs using the existing functionality?

from typing import Callable, TypeVar
from typing_extensions import ParamSpec

P = ParamSpec("P")
R = TypeVar("R")

def foo(a: str) -> None:
    ...

def wrap_func(fn: Callable[P, R]) -> Callable[P, R]:
    def inner(*args: P.args, **kwargs: P.kwargs) -> R:
        # Add code here as desired
        return fn(*args, **kwargs)
    return inner

bar = wrap_func(foo)

Or if you want to externalize the wrapper logic, wrap_func could take a wrapper function as an input parameter:

def wrap_func(
    wrapper: Callable[Concatenate[Callable[P, R], P], R], fn: Callable[P, R]
) -> Callable[P, R]:
    def inner(*args: P.args, **kwargs: P.kwargs) -> R:
        return wrapper(fn, *args, **kwargs)
    return inner

def bar_wrapper(fn: Callable[P, R], *args: P.args, **kwargs: P.kwargs) -> R:
    # Add code here as desired
    return fn(*args, *kwargs)

bar = wrap_func(bar_wrapper, foo)

[indigoviolet]

Both approaches above require passing the wrapped function fn or foo as an argument to the wrapping function. But a fairly common case is writing a particular function that invokes another function (say from a library, to create slightly better ergonomics or specialize some arguments). In that case, there doesn't appear to be a DRY method to describe the wrapper function's signature in terms of the wrapped function's signature (+/- some params).

See also #13617, #2003

[ligix]

Or if you want to externalize the wrapper logic, wrap_func could take a wrapper function as an input parameter:

def wrap_func(
    wrapper: Callable[Concatenate[Callable[P, R], P], R], fn: Callable[P, R]
) -> Callable[P, R]:
    def inner(*args: P.args, **kwargs: P.kwargs) -> R:
        return wrapper(fn, *args, **kwargs)
    return inner

def bar_wrapper(fn: Callable[P, R], *args: P.args, **kwargs: P.kwargs) -> R:
    # Add code here as desired
    return fn(*args, *kwargs)

bar = wrap_func(bar_wrapper, foo)

A slightly different version that works as a decorator:

def wraps_function(
    fun: Callable[P, T]
) -> Callable[[Callable[Concatenate[Callable[P, T], P], T]], Callable[P, T]]:
    def decorator(
        wrapper: Callable[Concatenate[Callable[P, T], P], T]
    ) -> Callable[P, T]:
        def decorated(*args: P.args, **kwargs: P.kwargs) -> T:
            return wrapper(fun, *args, **kwargs)
        return decorated
    return decorator

@wraps_function(myfun)
def bar(myfun: Callable[P, T], *args: P.args, **kwargs: P.kwargs) -> T:
    # Add code here as desired
    return myfun(*args, **kwargs) 

[zoranbosnjak]

@erictraut @ligix , is there any way to use the same function/method wrapping trick for object methods instead of regular functions? Or is there any other workaround? For example, in the code snippet below, the test2.get_arg().f(arg) is correctly checked. I would like to implement additional method Test2.f(...) with the same type signature as Test1.f(...), except for different self argument.

from typing import *

class Test1:
    @overload
    def f(self, key: Literal['A']) -> int: ...
    @overload
    def f(self, key: Literal['B']) -> str: ...
    def f(self, key: Any) -> Any:
        if key == 'A': return 1
        elif key == 'B': return 'x'
        assert_never(key)

class Test2:
    def __init__(self, arg: Test1):
        self._arg = arg

    def get_arg(self) -> Test1:
        return self._arg

    # what is a type signature of 'f'?
    def f(self, key):
        return self.get_arg().f(key)

# 'test2' contains 'test1'
test2 = Test2(Test1())

# test2.get_arg().f and test2.f are expected to behave the same

print(test2.get_arg().f('A'))
print(test2.get_arg().f('B'))
print(test2.get_arg().f('C')) # this is a type error as expected

print(test2.f('A'))
print(test2.f('B'))
print(test2.f('C')) # this is expected to be a type error too