pythonpython-typingliterals

Type hint for a function that returns only a specific set of values


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.


Solution

  • 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.