Python Type Hinting – Function Returning Specific Values

literalspythonpython-3.xpython-typingtype-hinting

I have a function that can only return a, b or c, all of them are of type T. I want to include this fact in the signature because of the special meaning they carry in the context of the function. How do I do that?

Currently, I use this

def fun(...) -> "a or b or c":
    #briefly explain the meaning of a, b and c in its docstring

Is that the correct one?

I know that I can do this

def fun(...) -> T:
    # briefly explain the meaning of a, b and c in its docstring

but as I said, I want to express in the signature that the function only returns those specific values.

Best Answer

You can do that with literal types.

from typing_extensions import Literal
# from typing import Literal  # Python 3.8 or higher

def fun(b: int) -> Literal["a", "b", "c"]:
    if b == 0:
        return "a"
    if b == 1:
        return "b"
    return "d"

mypy is able to detect the return "d" as a invalid statement:

error: Incompatible return value type (got "Literal['d']",
expected "Union[Literal['a'], Literal['b'], Literal['c']]")

Python 3.8

Thanks to the PEP 586, the Literal is already included by default in the Python 3.8 typing module.

Related Solutions

Python Type Hints – Type Hints for Function Returning Multiple Values

EDIT: Since Python 3.9 and the acceptance of PEP 585, you should use the built-in tuple class to typehint tuples.

You can use a typing.Tuple type hint (to specify the type of the content of the tuple, if it is not necessary, the built-in class tuple can be used instead):

from typing import Tuple, List

def greeting(name: str) -> Tuple[str, List[float], int]:
    # do something
    return a, b, c

Python Typing: Declare Return Value Type Based on Function Argument – Python, Type Hinting

TLDR: You need a TypeVar for the return type of calling t:

def fun(t: Callable[[int], R]) -> R:
    ...

Constraining on a type is too restrictive here. The function accepts any Callable that takes an integer, and the return type of the function is that of the Callable. This can be specified using a TypeVar for the return type:

from typing import Callable, TypeVar


R = TypeVar('R')  # the variable return type


def fun(t: Callable[[int], R]) -> R:
    return t(42)

fun(int)                            # Revealed type is 'builtins.int*'
fun(float)                          # Revealed type is 'builtins.float*'
reveal_type(fun(lambda x: str(x)))  # Revealed type is 'builtins.str*'

This works for types as well, because type instantiation is a call.

If a more complex signature, e.g. with keyword arguments, is needed, use Protocol (from typing or typing_extensions).

Note that if one explicitly wants to pass only 42 to the Callable, Literal (from typing or typing_extensions) can be used to specify that.

R = TypeVar('R')


def fun(t: Callable[[Literal[42]], R]) -> R:
    return t(42)

Note that any function of the type Callable[[int], R] also satisfies Callable[[Literal[42]], R].