python - 枚举属性的类型注释

Question

我有这段代码：

import enum


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


def get_color_return_something(some_color):
    pass

some_color如果我想我会从 Color 枚举中接收一个枚举属性（例如：），我该如何正确地向这个函数中的变量添加类型注释Color.RED？

Answer 1

类型提示 Color 类应该可以工作：

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

您可以尝试使用带有类型提示文字的选项。

从官方 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 4

另一个奇怪的语法解决方法是使用引用的前向引用语法（根据PEP 484）将 Enum 成员指定为 Enum 类的类型：

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 5

以下将适用于 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)