Аннотации типа Python для значения Enum

Question

У меня есть этот кусок кода:

import enum


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


def get_color_return_something(some_color):
    pass

Как правильно добавить аннотации к some_color varaible в этой функции, если я предполагаю получить значение из перечисления Color (например: Color.RED)?

Answer 1

Тип, намекающий на класс Color, должен работать:

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

Answer 2

Вы можете попробовать использовать вариант с подсказкой типа Literal.

Из официальной документации PEP8 мы знаем, что:

Литерал - это тип, который можно использовать для указания средствам проверки типов, что соответствующая переменная или параметр функции имеет значение, эквивалентное предоставленному литералу (или одному из нескольких литералов).

Так что в случае, если вам нужно использовать какие-то конкретные значения для аргумента функции, это будет один из лучших вариантов. Но этот подход не будет работать полностью, как мы ожидали, из-за типа значений Enum. Каждое значение будет иметь тип класса Enum. Это означает, что для приведенного ниже примера кода мы сможем использовать Color.GREEN в качестве аргумента функции. Таким образом, такое решение будет просто информацией для разработчиков, но не обязательным правилом для аргумента функции.

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

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

Пример кода:

      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

Второй вариант - это полностью правильное решение, предоставленное @ibarrond из сообщения выше, с подсказкой типа класса.

      some_color: Color

Так что здесь вы можете выбрать вариант работы в зависимости от ваших потребностей.

С моей точки зрения, мы можем попытаться указать возможные значения Enum для разработчиков, чтобы быть более ясными в наших требованиях к функции.

Answer 3

def get_color_return_something(some_color: Color):
pass

Answer 4

Следующее будет работать с 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)

Answer 5

Другой странный синтаксический обходной путь - указать члены Enum как тип класса Enum с использованием синтаксиса цитирования с прямой ссылкой (согласно 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'>

На изображении ниже видно, что предупреждений, выделенных в PyCharm, больше нет.

Для справки, вот скриншот жалобы PyCharm с указанием EXAMPLE член как <str> введите как имеет смысл:

Я не поклонник этого подхода , но это действительно избавиться от предупреждения.

Answer 6

Вы можете напечатать имя и значение следующим образом:

import enum

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

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

get_color_return_something(Color.RED)
# RED
# 1