如何键入提示检索动态枚举值的方法？

Question

我有一个 Python 模块，它有许多简单的枚举，定义如下：

class WordType(Enum):
    ADJ = "Adjective"
    ADV = "Adverb"

class Number(Enum):
    S = "Singular"
    P = "Plural"

因为有很多这样的枚举，而我只在运行时决定哪些枚举来查询任何给定的输入，所以我想要一个 function 可以检索给定枚举类型和枚举值作为字符串的值。 我成功地做到了如下：

names = inspect.getmembers(sys.modules[__name__], inspect.isclass)

def get_enum_type(name: str):
    enum_class = [x[1] for x in names if x[0] == name]
    return enum_class[0]

def get_enum_value(object_name: str, value_name: str):
    return get_enum_type(object_name)[value_name]

这很好用，但是现在我正在添加类型提示，并且我正在努力为这些方法定义返回类型：我尝试了slice和Literal[] ，两者都由mypy建议，但都没有检查出来（可能是因为我不明白我可以给Literal[]什么类型参数）。

我愿意修改枚举定义，但我更愿意保持动态查询不变。 最坏的情况，我可以做# type: ignore或只是 return -> Any ，但我希望有更好的东西。

Answer 1

由于您不想检查任何Enum的类型，我建议引入一个基本类型（例如GrammaticalEnum ）来标记所有Enum并将它们分组到一个自己的模块中：

# module grammar_enums
import sys
import inspect
from enum import Enum


class GrammaticalEnum(Enum):
    """use as a base to mark all grammatical enums"""
    pass


class WordType(GrammaticalEnum):
    ADJ = "Adjective"
    ADV = "Adverb"


class Number(GrammaticalEnum):
    S = "Singular"
    P = "Plural"


# keep this statement at the end, as all enums must be known first
grammatical_enums = dict(
    m for m in inspect.getmembers(sys.modules[__name__], inspect.isclass)
    if issubclass(m[1], GrammaticalEnum))

# you might prefer the shorter alternative:
# grammatical_enums = {k: v for (k, v) in globals().items()
#                      if inspect.isclass(v) and issubclass(v, GrammaticalEnum)}

关于打字，yakir0 已经建议了正确的类型，
但有了共同的基础，你可以缩小它们。

如果你愿意，你甚至可以完全摆脱你的功能：

from grammar_enums import grammatical_enums as g_enums
from grammar_enums import GrammaticalEnum

# just use g_enums instead of get_enum_value like this
WordType_ADJ: GrammaticalEnum = g_enums['WordType']['ADJ']

# ...or use your old functions:

# as your grammatical enums are collected in a dict now,
# you don't need this function any more:
def get_enum_type(name: str) -> Type[GrammaticalEnum]:
    return g_enums[name]


def get_enum_value(enum_name: str, value_name: str) -> GrammaticalEnum:
    # return get_enum_type(enum_name)[value_name]
    return g_enums[enum_name][value_name]

Answer 2

您始终可以运行您的函数并打印 function 的结果，以了解它应该是什么。 请注意，您可以像任何其他 class 一样在类型提示中使用Enum 。

例如：

>>> result = get_enum_type('WordType')
... print(result)
... print(type(result))
<enum 'WordType'>
<class 'enum.EnumMeta'>

所以你实际上可以使用

get_enum_type(name: str) -> EnumMeta

但是你可以通过使用Type from typing使它更漂亮，因为EnumMeta是一般 Enum 的类型。

get_enum_type(name: str) -> Type[Enum]

对于与get_enum_value类似的过程，您会得到

>>> type(get_enum_value('WordType', 'ADJ'))
<enum 'WordType'>

显然，您不会总是返回WordType类型，因此您可以使用Enum来概括返回类型。

总结一下：

get_enum_type(name: str) -> Type[Enum]
get_enum_value(object_name: str, value_name: str) -> Enum

Answer 3

正如我在评论中所说，我认为不可能拥有您的动态代码并让 MyPy 预测输出。 例如，我认为不可能让 MyPy 知道get_enum_type("WordType")应该是WordType而get_enum_type("Number")应该是Number 。

正如其他人所说，您可以澄清它们将是枚举。 您可以添加一个基本类型，并说它们将专门作为基本类型之一。 部分问题在于，尽管您可以 promise 它，但 MyPy 无法确认。 它无法知道inspect.getmembers(sys.modules[__name__], inspect.isclass)在[1]中只会有Enums或GrammaticalEnum 。

如果您愿意更改查找的实现，那么我建议您可以有利可图地使用__init_subclass__ 。 就像是

GRAMMATICAL_ENUM_LOOKUP: "Mapping[str, GrammaticalEnum]" = {}
class GrammaticalEnum(Enum):
    def __init_subclass__(cls, **kwargs):
        GRAMMATICAL_ENUM_LOOKUP[cls.__name__] = cls
        super().__init_subclass__(**kwargs)

 def get_enum_type(name: str) -> Type[GrammaticalEnum]:
     return GRAMMATICAL_ENUM_LOOKUP[name]

这至少有一个优点，即 MyPy 可以看到正在发生的事情，并且应该对此感到满意。 它知道所有内容实际上都是有效的GrammaticalEnum ，因为这就是GRAMMATICAL_ENUM_LOOKUP填充的全部内容。