Type annotations for Enum attribute

Question

I have this piece of code:

import enum


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


def get_color_return_something(some_color):
    pass

How do I properly add type annotations to the some_color variable in this function, if I suppose that I'll receive an enum attribute from the Color enum (for example: Color.RED )?

Answer 1

Type hinting the Color class should work:

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

Answer 2

def get_color_return_something(some_color: Color):
    pass

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

Answer 4

Another strange syntactic workaround is to specify Enum members as the type of the Enum class using the forward-referencing syntax of quoting (per PEP 484 ):

from enum import Enum


class ETest(Enum):
    EXAMPLE: 'ETest' = "example"  <--- forward referenced type


def example() -> ETest:
    return ETest.EXAMPLE


print(type(ETest.EXAMPLE.value))

<class 'str'>

In the image below it's evident that the warnings highlighted in PyCharm are no longer present.

For reference, here is a screenshot of PyCharm's grievance with specifying the EXAMPLE member as a <str> type as makes sense:

I am not a fan or this approach but it does get rid of the warning.

Answer 5

The following will work with Pyton 3.9/PyCharm

from enum import Enum
from typing import Optional, Union


class Color(Enum):
    RED: int = 1
    GREEN: int = 2


def guess_color(x: Union[Color.RED, Color.GREEN]) -> Optional[ValueError]:
    if x == Color.RED:
        print("Gotcha!")
    else:
        return ValueError(f"It's not {Color.RED}")


guess_color(Color.RED)