How to type-hint a method that retrieves dynamic enum value?

user_2065017 0
Viewed 1429

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.

python-typing mypy type-hinting enums python
3 Answers

As you don't want to check-type for any Enum, I suggest to introduce a base type (say GrammaticalEnum) to mark all your Enums 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]
user_96526400

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
user_146549090

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 GrammaticalEnums 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.

user_16887860
Related