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

by

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.

Leave a Comment