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

Solution 1:

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.

Solution 2:

You can't specify that your function returns only a subset of a type's values using type hinting alone. As the name implies, type hinting is all about types not values.

However, you can create a new enum.Enum subtype that only has the values you're going to return and use it in the function. Then you can type hint that you're returning the enum type.

import enum

class cmp_results(enum.IntEnum):
    less = -1
    equal = 0
    greater = 1

def my_cmp_function(x, y) -> cmp_results:
    if x < y: return cmp_results.less
    elif x == y: return cmp_results.equal
    else: return cmp_results.greater

This may be overkill. Just hinting int as the return type (and documenting the specific values) is probably good enough.