枚举属性的类型注释

Question

我有这段代码：

import enum


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


def get_color_return_something(some_color):
    pass

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

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

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

从official PEP8 documentation 我们知道：

Literal 它是一种类型，可用于向类型检查器指示相应的变量或函数参数具有与提供的文字（或多个文字之一）等效的值

因此，如果您需要为函数参数使用某些特定值，它将是最佳选择之一。但是由于 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 从post above 提供的完全正确的解决方案，只有一个类类型提示。

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)