Computed types in mypy

Question

Aside: The title of this questions is not ideal. What I'm trying to do could be achieved via computed types, but also by other means.

I'm writing some code that validates and sometimes converts JSON data, dynamically typed, to static Python types. Here are a few functions:

def from_str(x: Any) -> str:
    assert isinstance(x, str)
    return x


def from_int(x: Any) -> int:
    assert isinstance(x, int)
    return x

def from_list(f: Callable[[Any], T], x: Any) -> List[T]:
    assert isinstance(x, list)
    return [f(y) for y in x]

These work great. I'd also like to be able to combine them to convert union types. Ideally like this:

union = from_union([from_str, from_int], json)

The problem is how to type the from_union function. My first approach was this:

def from_union(fs: Iterable[Callable[[Any], T]], x: Any) -> T:
    for f in fs:
        try:
            return f(x)
        except AssertionError:
            pass
    assert False

Technically this is correct. If we substitute Union[str,int] for T the above expression is correctly typed, since from_str, by virtue of returning a str also returns a Union[str,int] (any value of type str is a value of type Union[str,int]). However, mypy doesn't want to do this substitution:

test/fixtures/python/quicktype.py:59: error: Argument 1 to "from_union" has incompatible type "List[Callable[[Any], object]]"; expected "Iterable[Callable[[Any], <nothing>]]"

It seems to go right to object instead of inferring Union[str,int].

Ideally, the type I'd like to give to from_union is something like

def from_union(fs: Iterable[Union[[Callable[[Any], S], Callable[[Any], T], ...]], x: Any) -> Union[S, T, ...]):

That's not supported in Python's typings. Another option would be to be able to specify a function that can compute either the type of fs from the actual return type for a specific invocation, or the other way around. Is anything like that possible? Are there any other options to do this without having to resort to cast?

Answer 1

As you deduced, this is unfortunately not something that's possible to express within Python's type system. The best available workaround (which is the same workaround Typeshed uses to type builtins like map, filter, and zip) is to abuse overloads, like so:

from typing import Iterable, Callable, Any, Union, TypeVar, overload, List

T1 = TypeVar('T1')
T2 = TypeVar('T2')
T3 = TypeVar('T3')

# Note: the two underscores tell mypy that the argument is positional-only
# and that doing things like `from_union(blob, f1=from_str)` is not legal

@overload
def from_union(x: Any, 
               __f1: Callable[[Any], T1],
               ) -> T1: ...

@overload
def from_union(x: Any, 
               __f1: Callable[[Any], T1],
               __f2: Callable[[Any], T2],
               ) -> Union[T1, T2]: ...

@overload
def from_union(x: Any, 
               __f1: Callable[[Any], T1],
               __f2: Callable[[Any], T2],
               __f3: Callable[[Any], T3],
               ) -> Union[T1, T2, T3]: ...

# The fallback: give up on the remaining callables
@overload
def from_union(x: Any, 
               __f1: Callable[[Any], T1],
               __f2: Callable[[Any], T2],
               __f3: Callable[[Any], T3],
               *fs: Callable[[Any], Any]
               ) -> Union[T1, T2, T3, Any]: ...

def from_union(x: Any, *fs: Callable[[Any], Any]) -> Any:
    for f in fs:
        try:
            return f(x)
        except AssertionError:
            pass
    assert False

What this function basically does is hard-code in support for up to three callables, and gives up if you try passing in any more. Naturally, to support accepting even more callables, add a few more overloads.

This new function's API does change slightly: it needs to be called like so:

my_union = from_union(json_blob, from_str, from_int)

If you want an API that's more similar to the original and have the functions come first, you'd either need to convert x into a keyword-only argument (e.g. from_union(*fs: Callable[[Any], Any], *, x: Any) -> Any) or store the functions in a tuple, which would look like this:

@overload
def from_union(fs: Tuple[Callable[[Any], T1]], x: Any) -> T1: ...

@overload
def from_union(fs: Tuple[Callable[[Any], T1], Callable[[Any], T2]], x: Any) -> Union[T1, T2]: ...

# etc...

# The final fallback: have the tuple accept any number of callables
@overload
def from_union(fs: Tuple[Callable[[Any], Any], ...], x: Any) -> Any: ...

def from_union(fs: Tuple[Callable[[Any], Any], ...], x: Any) -> Any:
    for f in fs:
        try:
            return f(x)
        except AssertionError:
            pass
    assert False

In both cases, the "fallback" if the user passes in too many args will introduce some dynamism in the output. If you don't like this, just delete the final fallback.