type friendly delegation in Python

Question

Consider the following code sample

def sum(a: int, b: int):
  return a + b

def wrap(*args, **kwargs):
  # delegate to sum
  return sum(*args, **kwargs)

The code works well except that type hint is lost. It's very common in Python to use *args, **kwargs to implement delegation pattern. It would be great to have a way to keep the type hint while using them, but I don't know if it is possible and how.

Answer 1

See https://github.com/python/typing/issues/270 for a long discussion of this problem. You can achieve this by decorating wrap with an appropriately typed identity function:

F = TypeVar("F", bound=Callable)
def copy_signature(_: F) -> Callable[..., F]:
    return lambda f: f

def s(x: int, y: int) -> int:
    return x + y

@copy_signature(s)
def wrap(*args, **kwargs):
    s(*args, **kwargs)

reveal_type(wrap)  # Revealed type is "def (x: int, y: int) -> int"

As far as I know, the decorator is necessary - it is still not possible to do this using type hints alone, even with PEP612. Since it was already good practice to use the functools.wraps decorator in this situation (which copies the runtime type information), this is not such a loss - you could instead define

def wraps(f: F) -> Callable[..., F]:
    return functools.wraps(f) # type: ignore

and then both the runtime and static type information should be correct so long as you use this decorator. (Sadly the typeshed stubs for functools.wraps included with mypy aren't quite restrictive enough to get this working out of the box.)

PEP612 adds the ability to add/remove arguments in your wrapper (by combining ParamSpec with Concatenate), but it doesn't remove the need for some kind of higher-order function (like a decorator) to let the type system infer the signature of wrap from that of s.