Typehint function *args -> tuple[*args] with constraint on the args

Question

We want to type hint a function

def f(*args:float)->tuple[float,...]:
   ...
   return tuple(args)

such that it is specified that the number of elements in the tuple matches the number of args. Of course, the return here is a placeholder for more complicated logic.

We would like to use mypy or pylance to check if we always return a) the correct number of elements and b) the correct tyoe of all elements.

Using TypeVarTuple (https://peps.python.org/pep-0646/#args-as-a-type-variable-tuple) would allow to specify that we return the same number of elements, but not the type.

Is there in current python (3.12) way to do it besides writing many overloads for 1-parameter, 2-parameter, 3-parameters etc?

Answer 1

Yes, you can write a no-op decorator to make the signature of f reject attempts at passing a type that you don't want. The following example makes f reject any attempts at passing values which aren't compatible with float.

Demo: mypy Playground, Pyright Playground

import typing_extensions as t

if t.TYPE_CHECKING:
    import collections.abc as cx

    F = t.TypeVar("F", bound=cx.Callable[..., t.Any])
    Ts = t.TypeVarTuple("Ts")
 
    class _FloatOnlyCallable(t.Protocol):
        def __call__(self, /, *args: float) -> t.Any: ...

def asFloatOnlyCallable(f: F, /) -> F | _FloatOnlyCallable:
    """Decorate a function to make it only accept variadic positional float arguments"""
    return f

@asFloatOnlyCallable
def f(*args: *Ts) -> tuple[*Ts]:
    return args

>>> a, b = f(1.0, 2.0)  # OK
>>> c, d = f("3.0", 4.0)  # Error: Incompatible type "str", expected "float"
>>> e, g, h = f(5.0, 6.0)  # Error: Need more values to unpack