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

Question

I have a Python module that has a number of simple enums that are defined as follows:

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

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

Because there are a lot of these enums and I only decide at runtime which enums to query for any given input, I wanted a function that can retrieve the value given the enum-type and the enum-value as strings. I succeeded in doing that as follows:

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]

This works well, but now I'm adding type hinting and I'm struggling with how to define the return types for these methods: I've tried slice and Literal[] , both suggested by mypy , but neither checks out (maybe because I don't understand what type parameter I can give to Literal[] ).

I am willing to modify the enum definitions, but I'd prefer to keep the dynamic querying as-is. Worst case scenario, I can do # type: ignore or just return -> Any , but I hope there's something better.

Answer 1

As you don't want to check-type for any Enum , I suggest to introduce a base type (say GrammaticalEnum ) to mark all your Enum s and to group them in an own module:

# 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)}

Regarding typing, yakir0 already suggested the right types,
but with the common base you can narrow them.

If you like, you even could get rid of your functions at all:

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

You can always run your functions and print the result of the function to get a sense of what it should be. Note that you can use Enum in type hinting like any other class.

For example:

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

So you can actually use

get_enum_type(name: str) -> EnumMeta

But you can make it prettier by using Type from typing since EnumMeta is the type of a general Enum.

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

For a similar process with get_enum_value you get

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

Obviously you won't always return the type WordType so you can use Enum to generalize the return type.

To sum it all up:

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

Answer 3

As I said in a comment, I don't think it's possible to have your dynamic code and have MyPy predict the outputs. For example, I don't think it would be possible to have MyPy know that get_enum_type("WordType") should be a WordType whereas get_enum_type("Number") should be a Number .

As others have said, you could clarify that they'll be Enums. You could add a base type, and say that they'll specifically be one of the base types. Part of the problem is that, although you could promise it, MyPy wouldn't be able to confirm. It can't know that inspect.getmembers(sys.modules[__name__], inspect.isclass) will just have Enums or GrammaticalEnum s in [1] .

If you're willing to change the implementation of your lookup, then I'd suggest you could make profitable use of __init_subclass__ . Something like

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]

This at least has the advantage that the MyPy can see what's going on, and should be broadly happy with it. It knows that everything will in fact be a valid GrammaticalEnum because that's all that GRAMMATICAL_ENUM_LOOKUP gets populated with.