Notes on Type-Hinting compose

[mentalisttraceur]

I recently finally sat down and made compose-stubs for my compose, so here are some notes in case it helps here with the type hints for funcy's compose.

Unfortunately, there's no Good way to do it. But it's possible to brute-force it with overloads up to some finite arity:

from typing import Callable, ParamSpec, TypeVar, overload


P = ParamSpec('P')
R1 = TypeVar('R1')
R2 = TypeVar('R2')
R3 = TypeVar('R3')
# and so on up to whatever arity you want to support.


@overload
def compose(f1: Callable[P, R1], /) -> Callable[P, R1]:
    ...
@overload
def compose(f2: Callable[[R1], R2], f1: Callable[P, R1], /) -> Callable[P, R2]:
    ...
@overload
def compose(f3: Callable[[R2], R3], f2: Callable[[R1], R2], f1: Callable[P, R1], /) -> Callable[P, R3]:
    ...
# and so on up to whatever arity you want to support.

ParamSpec is only available on 3.10 and up. Personally I didn't bother with compatibility to older Pythons in my stubs, but it would be easy enough to gracefully degrade from ParamSpec(P) to f.e. ... if you want to support older Pythons:

try:
    from typing import ParamSpec
except ImportError:
    P = ...
else:
    P = ParamSpec('P')

[mentalisttraceur]

In my stubs, I type-hinted arities up to 16. That seemed like a good spot for reliably feels-instant type-checking speed.

I initially tried going all the way up to 256, but that was very slow, especially when there was a typing error, because modern type checkers seem to check every overload (they don't pre-separate/narrow by arity), so if there's a type mismatch and you have 256 overloads, all 256 get checked for a possible match before the type checker gives up.

[mentalisttraceur]

I also wrote a little script to generate compose overloads up to some arity, since it's obviously too much of a tedium to manually generate them. Here's a version of that simplified down to your needs:

#!/usr/bin/env python3


def _callable(arguments, return_):
    return f'Callable[{arguments}, {return_}]'


def _overload(name, arguments, return_):
    return f'@overload\ndef {name}({arguments}) -> {return_}:\n    ...'


def _overloads(name, up_to):
    argument = _callable('P', 'R1')
    arguments = f'f1: {argument}, /'
    yield _overload(name, arguments, _callable('P', 'R1'))
    for number in range(2, up_to):
        argument = _callable(f'[R{number-1}]', f'R{number}')
        arguments = f'f{number}: {argument}, {arguments}'
        yield _overload(name, arguments, _callable('P', f'R{number}'))
        

print('from typing import Callable, ParamSpec, TypeVar, overload')
print("P = ParamSpec('P')")

for number in range(1, 16):
    print(f"R{number} = TypeVar('R{number}')")

for overload in _overloads('compose', 16):
    print(overload)

If I remember right, funcy also has rcompose - you'll need to reverse the arguments build-up in _overloads for that. (And mind the trailing , / - assuming you don't just remove it, f.e. to support Python <=3.7.)

[mentalisttraceur]

Oh, my compose also raises an error instead of returning the identity function in the zero arguments case. I can't remember what funcy's does, but if it does the latter, you'll also want:

@overload
def compose() -> Callable[[R1], R1]:
    ...

[mentalisttraceur]

Note that the overloads have to be right after each other.

I had been tempted to interleave each additional Rn = TypeVar('Rn') right before the nth overload of compose instead of all TypeVars and then all overloads, but interleaving causes MyPy to ignore all overloads that aren't immediately after the first one. (Not sure if Pyright and Pyre have this issue - I only started testing those after I got it working in MyPy.)

[mentalisttraceur]

Of course, the other complication we have here that I didn't have with my compose is funcy's "extended function semantics".

I don't have a good answer to that, short of a combinatoric explosion of type hint overloads. And even that doesn't really work fully because the functions produced by re_finder have different return types depending on the string contents.