Type annotations for Enum attribute

Solution 1:

Type hinting the Color class should work:

def get_color_return_something(some_color: Color):
    print(some_color.value)

Solution 2:

def get_color_return_something(some_color: Color):
    pass

Solution 3:

You can try to use an option with type hint Literal.

From official PEP8 documentation we know that:

Literal it's type that can be used to indicate to type checkers that the corresponding variable or function parameter has a value equivalent to the provided literal (or one of several literals)

So in case if you need to use some specific values for a function argument it will be one of the best options. But this approach will not work fully as we expected, because of the type of the Enum values. Each value will have a type of Enum class. This means that for the code example below we will be able to put Color.GREEN as the function argument. So such a solution will be just information for developers, but not a mandatory rule for a function argument.

class Color(enum.Enum):
    RED = '1'
    BLUE = '2'
    GREEN = '3'

print(type(Color.RED)  # will return <enum 'Color'>

Code example:

from enum import Enum
from typing import Literal


class Color(Enum):
    RED = '1'
    BLUE = '2'
    GREEN = '3'

def some_function(some_color: Literal[Color.RED, Color.BLUE]) -> None:
    pass

The second option it's fully correct solution provided by @ibarrond from the post above with just a class type hint.

some_color: Color

So here you can choose the option to work with depending on your needs.

From my point of view we can try to specify possible Enum values for developers, to be more clear in our requirements for a function.