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

Question

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.

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

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