Expand variadic type variables

expander.py

LIMIT = 5
BOUND = 'object'

def prelude(limit: int, bound: str) -> None:
    print('from typing import Callable, Iterable, Iterator, Tuple, TypeVar, overload')
    print('Ts = TypeVar(\'Ts\', bound={bound})'.format(bound=bound))
    print('R = TypeVar(\'R\')')
    for i in range(LIMIT):
        print('T{i} = TypeVar(\'T{i}\', bound={bound})'.format(i=i+1, bound=bound))

def expand_template(template: str,
                    arg_template: str = 'arg{i}: {Ts}',
                    lower: int = 0,
                    limit: int = LIMIT) -> None:
    print()
    for i in range(lower, limit):
        tvs = ', '.join('T{i}'.format(i=j+1) for j in range(i))
        args = ', '.join(arg_template.format(i=j+1, Ts='T{}'.format(j+1))
                         for j in range(i))
        print('@overload')
        s = template.format(Ts=tvs, argsTs=args)
        s = s.replace('Tuple[]', 'Tuple[()]')
        print(s)
    args_l = [arg_template.format(i=j+1, Ts='Ts') for j in range(limit)]
    args_l.append('*' + (arg_template.format(i='s', Ts='Ts')))
    args = ', '.join(args_l)
    s = template.format(Ts='Ts, ...', argsTs=args)
    s = s.replace('Callable[[Ts, ...]', 'Callable[...')
    print('@overload')
    print(s)

def main():
    prelude(LIMIT, BOUND)

    # map()
    expand_template('def map(func: Callable[[{Ts}], R], {argsTs}) -> R: ...',
                    lower=1)
    # zip()
    expand_template('def zip({argsTs}) -> Tuple[{Ts}]: ...')

    # Naomi's examples
    expand_template('def my_zip({argsTs}) -> Iterator[Tuple[{Ts}]]: ...',
                    'arg{i}: Iterable[{Ts}]')
    expand_template('def make_check({argsTs}) -> Callable[[{Ts}], bool]: ...')
    expand_template('def my_map(f: Callable[[{Ts}], R], {argsTs}) -> Iterator[R]: ...',
                    'arg{i}: Iterable[{Ts}]')
                    

main()

output.txt

from typing import Callable, Iterable, Iterator, Tuple, TypeVar, overload
Ts = TypeVar('Ts', bound=object)
R = TypeVar('R')
T1 = TypeVar('T1', bound=object)
T2 = TypeVar('T2', bound=object)
T3 = TypeVar('T3', bound=object)
T4 = TypeVar('T4', bound=object)
T5 = TypeVar('T5', bound=object)

@overload
def map(func: Callable[[T1], R], arg1: T1) -> R: ...
@overload
def map(func: Callable[[T1, T2], R], arg1: T1, arg2: T2) -> R: ...
@overload
def map(func: Callable[[T1, T2, T3], R], arg1: T1, arg2: T2, arg3: T3) -> R: ...
@overload
def map(func: Callable[[T1, T2, T3, T4], R], arg1: T1, arg2: T2, arg3: T3, arg4: T4) -> R: ...
@overload
def map(func: Callable[..., R], arg1: Ts, arg2: Ts, arg3: Ts, arg4: Ts, arg5: Ts, *args: Ts) -> R: ...

@overload
def zip() -> Tuple[()]: ...
@overload
def zip(arg1: T1) -> Tuple[T1]: ...
@overload
def zip(arg1: T1, arg2: T2) -> Tuple[T1, T2]: ...
@overload
def zip(arg1: T1, arg2: T2, arg3: T3) -> Tuple[T1, T2, T3]: ...
@overload
def zip(arg1: T1, arg2: T2, arg3: T3, arg4: T4) -> Tuple[T1, T2, T3, T4]: ...
@overload
def zip(arg1: Ts, arg2: Ts, arg3: Ts, arg4: Ts, arg5: Ts, *args: Ts) -> Tuple[Ts, ...]: ...

@overload
def my_zip() -> Iterator[Tuple[()]]: ...
@overload
def my_zip(arg1: Iterable[T1]) -> Iterator[Tuple[T1]]: ...
@overload
def my_zip(arg1: Iterable[T1], arg2: Iterable[T2]) -> Iterator[Tuple[T1, T2]]: ...
@overload
def my_zip(arg1: Iterable[T1], arg2: Iterable[T2], arg3: Iterable[T3]) -> Iterator[Tuple[T1, T2, T3]]: ...
@overload
def my_zip(arg1: Iterable[T1], arg2: Iterable[T2], arg3: Iterable[T3], arg4: Iterable[T4]) -> Iterator[Tuple[T1, T2, T3, T4]]: ...
@overload
def my_zip(arg1: Iterable[Ts], arg2: Iterable[Ts], arg3: Iterable[Ts], arg4: Iterable[Ts], arg5: Iterable[Ts], *args: Iterable[Ts]) -> Iterator[Tuple[Ts, ...]]: ...

@overload
def make_check() -> Callable[[], bool]: ...
@overload
def make_check(arg1: T1) -> Callable[[T1], bool]: ...
@overload
def make_check(arg1: T1, arg2: T2) -> Callable[[T1, T2], bool]: ...
@overload
def make_check(arg1: T1, arg2: T2, arg3: T3) -> Callable[[T1, T2, T3], bool]: ...
@overload
def make_check(arg1: T1, arg2: T2, arg3: T3, arg4: T4) -> Callable[[T1, T2, T3, T4], bool]: ...
@overload
def make_check(arg1: Ts, arg2: Ts, arg3: Ts, arg4: Ts, arg5: Ts, *args: Ts) -> Callable[..., bool]: ...

@overload
def my_map(f: Callable[[], R], ) -> Iterator[R]: ...
@overload
def my_map(f: Callable[[T1], R], arg1: Iterable[T1]) -> Iterator[R]: ...
@overload
def my_map(f: Callable[[T1, T2], R], arg1: Iterable[T1], arg2: Iterable[T2]) -> Iterator[R]: ...
@overload
def my_map(f: Callable[[T1, T2, T3], R], arg1: Iterable[T1], arg2: Iterable[T2], arg3: Iterable[T3]) -> Iterator[R]: ...
@overload
def my_map(f: Callable[[T1, T2, T3, T4], R], arg1: Iterable[T1], arg2: Iterable[T2], arg3: Iterable[T3], arg4: Iterable[T4]) -> Iterator[R]: ...
@overload
def my_map(f: Callable[..., R], arg1: Iterable[Ts], arg2: Iterable[Ts], arg3: Iterable[Ts], arg4: Iterable[Ts], arg5: Iterable[Ts], *args: Iterable[Ts]) -> Iterator[R]: ...

Hm, the overflow expansion with *args is wrong -- it should use the upper bound of Ts, not Ts itself.

Actually I'm not so sure about that. This:

@overload
def my_zip(arg1: Iterable[Ts], arg2: Iterable[Ts], arg3: Iterable[Ts], arg4: Iterable[Ts], arg5: Iterable[Ts], *args: Iterable[Ts]) -> Iterator[Tuple[Ts, ...]]: ...

looks better than this:

@overload
def my_zip(arg1: Iterable[object], arg2: Iterable[object], arg3: Iterable[object], arg4: Iterable[oject], arg5: Iterable[object], *args: Iterable[object]) -> Iterator[Tuple[object, ...]]: ...

or this:

@overload
def my_zip(arg1: Iterable[Any], arg2: Iterable[Any], arg3: Iterable[Any], arg4: Iterable[Any], arg5: Iterable[Any], *args: Iterable[Any]) -> Iterator[Tuple[Any, ...]]: ...

The variant with the upper bound (here object) is suboptimal in case all the iterables actually derive from a more precise common base class -- the return type would still be tuples of objects. The variant with Any always returns tuples of Any which is too imprecise. So I like the variant with the typevar best after all.