Type hint for a function that returns only a specific set of values
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.
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.