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.

Answers:

Thank you for visiting the Q&A section on Magenaut. Please note that all the answers may not help you solve the issue immediately. So please treat them as advisements. If you found the post helpful (or not), leave a comment & I’ll get back to you as soon as possible.

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

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

Method 3

If all are of the same exact type just add that as the return type:

def func(...) -> T: # or int or whatever else

I want to express in the signature that the function only return those specific values

Type hints don’t specify a name or a value they just specify a type; a type checker tries to act on the type that’s provided.

If you’re just doing this for documentation purposes, add 'a or b or c‘ to it if you want; users will understand it but type checkers won’t and they definitely won’t act on it.

All methods was sourced from stackoverflow.com or stackexchange.com, is licensed under cc by-sa 2.5, cc by-sa 3.0 and cc by-sa 4.0