TLDR

I’m trying to make my type checker happy by specifying overloading of my functions. There is a specific signature I have many of, and I’d like to keep things dry.

Example and current code

I have many functions that take, as their first argument, an int or a float, and return the same type. To make sure my editor understands the one-to-one mapping, I add overloading:

from typing import overload

@overload
def add_two(x: int) -> int: ...

@overload
def add_two(x: float) -> float: ...

def add_two(x: int | float) -> int | float:
    return x + 2

This works, great.

Now, I have many such functions, e.g.

def add_two(x: int | float) -> int | float:
    return x + 2

def square(x: int | float) -> int | float:
    return x ** 2

def multiply_by_int_add_int(x: int | float, y: int, z: int) -> int | float:
    return x * y + z

(...)

Adding the overloading in the way it’s described, to each of the functions, will add a lot of almost identical code.

Question

Is there a way to write a decorator such, that I can add it to each function definition?

Note that:

• The functions have in common, that the type of the first argument determines the type of the return value.

• The number and type of the remaining parameters is not determined.

Current attempt

import functools
from typing import Callable, ParamSpec, TypeVar, overload

T = TypeVar("T")
P = ParamSpec("P")


def int2int_and_float2float(fn: Callable[P, T]) -> Callable[P, T]:

    @overload
    def wrapped(x: int, *args, **kwargs) -> int: ...

    @overload
    def wrapped(x: float, *args, **kwargs) -> float: ...

    @functools.wraps(fn)
    def wrapped(*args: P.args, **kwargs: P.kwargs):
        return fn(*args, **kwargs)

    return wrapped

@int2int_and_float2float
def add_two(x: int | float) -> int | float:
    """Add two"""
    return x * 2

This is not working; add_two(3) is not recognised to give an integer:

What am I missing?