How can I type-hint a function where the return type depends on the input type of an argument?

Question

Assume that I have a function which converts Python data-types to Postgres data-types like this:

def map_type(input):
    if isinstance(input, int):
                 


        

Answer 1

This is exactly what function overloads are for.

In short, you do the following:

from typing import overload

# ...snip...

@overload
def map_type(value: int) -> MyEnum: ...

@overload
def map_type(value: str) -> MyCustomClass: ...

def map_type(value: Union[int, str]) -> Union[MyEnum, MyCustomClass]:
    if isinstance(value, int):
        return MyEnum(value)
    elif isinstance(value, str):
        return MyCustomClass(value)
    raise TypeError('Invalid input type')

Now, when you do map_type(3), mypy will understand that the return type is MyEnum.

And at runtime, the only function to actually run is the final one -- the first two are completely overridden and ignored.