How does one map Python type hints to value order in a non-homogenous iterable without manually writing out every @overload permutation?

36 views Asked by At

I need to map the argument order to different permutations of a set of passed in arguments.

For example (not exactly the code I'm using, this is a minimal illustrative example), suppose I define a couple of literals:

_S = Literal["str"]
_I = Literal["int"]

And I have a function that looks up some callables :


@overload
def get_func(name: _K) -> Callable[[int], str]:
    ...


@overload
def get_func(name: _L) -> Callable[[str], str]:
    ...


def get_func(name: _K | _L) -> Callable[[int], str] | Callable[[str], str]:
    match name:
        case "int":
            return from_int
        case "str":
            return from_str
        case _:
            # Handle the possible runtime error
            raise ValueError("OMG error or something.")

def process_args(args: tuple[_K | _L, ...]) -> tuple[Callable[[int], str] | Callable[[str], str], ...]:
    base: tuple[Callable[[int], str] | Callable[[str], str], ...] = tuple()
    for arg in args:
        base += (get_func(arg), )
    return base


funcs = process_args(("str", "str", "int"))
reveal_type(funcs)

The revealed type for funcs is:

Revealed type is "builtins.tuple[Union[def (builtins.int) -> builtins.str, def (builtins.str) -> builtins.str], ...]"

Because of course that's what it is. What I would like is to find a way to have the return type vary (much like the overload on get_func ).

I could get around that to some degree by making more overloads like this:



@overload
def get_func_tuple(args: tuple[_K, ...]) -> tuple[Callable[[int], str], ...]:
    ...

@overload
def get_func_tuple(args: tuple[_L, ...]) -> tuple[Callable[[str], str], ...]:
    ...

@overload
def get_func_tuple(args: tuple[_L, _K]) -> tuple[Callable[[str], str], Callable[[int], str]]:
    ...

@overload
def get_func_tuple(args: tuple[_K, _L]) -> tuple[Callable[[int], str], Callable[[str], str]]:
    ...

@overload
def get_func_tuple(args: tuple[_K, _L, _K]) -> tuple[Callable[[int], str], Callable[[str], str], Callable[[int], str]]:
    ...

@overload
def get_func_tuple(args: tuple[_L, _K, _L]) -> tuple[Callable[[str], str], Callable[[int], str], Callable[[str], str]]:
    ...

@overload
def get_func_tuple(args: tuple[_K, _L, _L, _K]) -> tuple[Callable[[str], str], Callable[[int], str], Callable[[int], str], Callable[[str], str]]:
    ...

@overload
def get_func_tuple(args: tuple[_L, _K, _K, _L]) -> tuple[Callable[[int], str], Callable[[str], str], Callable[[str], str], Callable[[int], str]]:
    ...

def get_func_tuple(args: tuple[_K | _L, ...]) -> tuple[Callable[[int], str] | Callable[[str], str], ...]:
    return tuple((get_func(x) for x in args))


func2 = get_func_tuple(("int", "str", "str", "int"))
reveal_type(func2)

Here the revealed type for func2 is:

Revealed type is "tuple[def (builtins.int) -> builtins.str, def (builtins.str) -> builtins.str, def (builtins.str) -> builtins.str, def (builtins.int) -> builtins.str]"

Which IS technically what I would like. What I don't want to do is have to describe every possible overload permutation. Since the input values are a finite set of literal values it seems to me that it ought to be possible to have the returned type be mapped to the input values.

i.e.

A --> X
B --> Z

(A, B)           ==> (X, Z)
(A, )            ==> (X, )
(A, A, A)        ==> (X, X, X)
(A, B, B, A, B)  ==> (X, Z, Z, X, Z)

The type system should have enough information to figure this all out since it's all literals, but I cannot find a non-annoying way to do it without limiting the number of input values and manually writing out all of the possible permutations.

0

There are 0 answers