Python type hints for generic *args (specifically zip or zipWith)

Question

I am writing a function called zip_with with the following signature:

_A = TypeVar("_A")
_B = TypeVar("_B")
_C = TypeVar("_C")


def zip_with(zipper: Callable[[_A, _B], _C], a_vals: Iterable[_A], b_vals: Iterable[_B]) -> Generator[_C, None, None]: ...

It's like zip, but allows you to aggregate with any arbitrary function. This works fine for an implementation of zip_with that only allows 2 arguments.

Is there support for adding type hints for a variable number of arguments? Specifically, I want an arbitrary list of generic types, and I want the type checker to be able to match the types of the arguments to the arguments of the zipper. Here's how I can do it without specific types:

def zip_with(zipper: Callable[..., _C], *vals: Iterable) -> Generator[_C, None, None]: ...

In other words, I want the type checker to be able to match the types of *vals to the input arguments of zipper.

Answer 1

Unfortunately, there is not a clean way of expressing this kind of type signature. In order to do so, we need a feature called variadic generics. While there is general interest in adding this concept to PEP 484, it's probably not going to happen in the short-term.

For the mypy core team in particular, I'd roughly estimate this work on this feature might tentatively start later this year, but probably will not be available for general use until early to mid 2020 at the absolute earliest. (This is based on some in-person conversations with various members of their team.)

---

The current workaround is to abuse overloads like so:

from typing import TypeVar, overload, Callable, Iterable, Any, Generator

_T1 = TypeVar("_T1")
_T2 = TypeVar("_T2")
_T3 = TypeVar("_T3")
_T4 = TypeVar("_T4")
_T5 = TypeVar("_T5")

_TRet = TypeVar("_TRet")

@overload
def zip_with(zipper: Callable[[_T1, _T2], _TRet], 
             __vals1: Iterable[_T1],
             __vals2: Iterable[_T2],
             ) -> Generator[_TRet, None, None]: ...
@overload
def zip_with(zipper: Callable[[_T1, _T2, _T3], _TRet], 
             __vals1: Iterable[_T1],
             __vals2: Iterable[_T2],
             __vals3: Iterable[_T3],
             ) -> Generator[_TRet, None, None]: ...
@overload
def zip_with(zipper: Callable[[_T1, _T2, _T3, _T4], _TRet], 
             __vals1: Iterable[_T1],
             __vals2: Iterable[_T2],
             __vals3: Iterable[_T3],
             __vals4: Iterable[_T4],
             ) -> Generator[_TRet, None, None]: ...
@overload
def zip_with(zipper: Callable[[_T1, _T2, _T3, _T4, _T5], _TRet], 
             __vals1: Iterable[_T1],
             __vals2: Iterable[_T2],
             __vals3: Iterable[_T3],
             __vals4: Iterable[_T4],
             __vals5: Iterable[_T5],
             ) -> Generator[_TRet, None, None]: ...

# One final fallback overload if we want to handle callables with more than
# 5 args more gracefully. (We can omit this if we want to bias towards
# full precision at the cost of usability.)
@overload
def zip_with(zipper: Callable[..., _TRet],
             *__vals: Iterable[Any],
             ) -> Generator[_TRet, None, None]: ...

def zip_with(zipper: Callable[..., _TRet],
             *__vals: Iterable[Any],
             ) -> Generator[_TRet, None, None]:
    pass

This approach is obviously pretty inelegant -- it's clunky to write, and performs precise type-checking for only callables that accept up to 5 args.

But in practice, this is usually good enough. Pragmatically, most callables aren't too long, and we can always tack on more overloads to handle more special cases if needed.

And in fact, this technique is actually what's being used to define the types for zip: https://github.com/python/typeshed/blob/master/stdlib/2and3/builtins.pyi#L1403