枚举值的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

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

从官方 PEP8 文档我们知道：

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

因此，如果您需要对函数参数使用某些特定值，那么它将是最好的选择之一。但由于枚举值的类型，这种方法不会完全按照我们的预期工作。每个值都有一个 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