Make MyPy properly detect type of the descriptor used within Enum class

Question

Problem description

I've been trying to add MyPy support for my implementation of a slightly modified Enum class, where it is possible to define a certain containers for enum values.

My enum and descriptor classes definitions (enums.py file):

from __future__ import annotations

import inspect
from enum import Enum, EnumMeta
from typing import Iterable, Any, Generic, TypeVar

T = TypeVar("T", bound=Iterable[Any])
K = TypeVar("K", bound=dict[Any, Any])


class EnumMapper(Generic[T]):
    """Simple descriptor class that enables definitions of iterable within Enum classes."""

    def __init__(self, container: T):
        self._container = container

    def __get__(self, instance: Any, instance_cls: Any) -> T:
        return self._container


class StrictEnumMapper(EnumMapper[K]):
    """
    Alias for `EnumMapper` descriptor class that requires all enum values
    to be present in definition of mapping before Enum class can be
    instantiated.
    Only dict-based mappings are supported!
    """

    def __init__(self, container: K):
        if not issubclass(type(container), dict):
            raise ValueError(
                f"{StrictEnumMapper.__name__} supports only dict-based containers, {type(container)} provided"
            )
        super().__init__(container)


class NamespacedEnumMeta(EnumMeta):
    """
    Metaclass checking Enum-based class for `StrictEnumMapper` fields
    and ensuring that all enum values are provided within such fields.
    E.g.
    >>> class Items(metaclass=NamespacedEnumMeta):
    >>>     spam = "spam"
    >>>     eggs = "eggs"
    >>>     foo = "foo"
    >>>
    >>>     food_preferences = StrictEnumMapper({
    >>>         spam: "I like spam!",
    >>>         eggs: "I really don't like eggs...",
    >>>     })
    will not instantiate and the `RuntimeError` informing about missing mapping
    for `Items.foo` will be raised.

    This class takes burden of remembering to add new enum values to mapping
    off programmers' shoulders by doing it automatically during runtime.
    The app will simply not start and inform about a mistake.
    """

    def __new__(
            mcs,
            cls,
            bases: tuple[type, ...],
            namespace: dict[str, Any],  # namespace is actually of _EnumDict type
    ):
        enum_values = [namespace[member] for member in namespace._member_names]  # type: ignore
        strict_enum_mappers_violated = [
            field_name
            for (field_name, field) in namespace.items()
            if (
                    inspect.ismethoddescriptor(field)
                    and issubclass(type(field), StrictEnumMapper)
                    and not all(enum in field._container for enum in enum_values)
            )
        ]
        if strict_enum_mappers_violated:
            raise RuntimeError(
                f"The following {cls} fields do not contain all possible "
                + f"enum values: {strict_enum_mappers_violated}"
            )
        return EnumMeta.__new__(mcs, cls, bases, namespace)


class NamespacedEnum(Enum, metaclass=NamespacedEnumMeta):
    """Extension of the basic Enum class, allowing for EnumMapper and StrictEnumMapper usages."""

Example usage of the classes (main.py file):

from enums import NamespacedEnum, StrictEnumMapper


class Food(NamespacedEnum):
    spam = "spam"
    eggs = "eggs"
    foo = "foo"

    reactions = StrictEnumMapper(
        {
            spam: "I like it",
            eggs: "I don't like it...",
            foo: "I like?",
        }
    )


if __name__ == "__main__":
    print(Food.reactions[Food.eggs.value])
    try:
        class ActionEnum(NamespacedEnum):
            action1 = "action1"
            action2 = "action2"

            func_for_action = StrictEnumMapper({
                action1: lambda: print("Doing action 1"),
            })
    except RuntimeError as e:
        print(f"Properly detected missing enum value in `func_for_action`. Error message: {e}"

Running above should result in:

$ python main.py
I don't like it...
Properly detected missing enum value in `func_for_action`. Error message: The following ActionEnum fields do not contain all possible enum values: ['func_for_action']

Running MyPy (version 0.910) returns following errors:

$ mypy --python-version 3.9 enums.py main.py                                                                                                                                        1 ✘  NamespacedEnums   11:53:05  
main.py:19: error: Value of type "Food" is not indexable
Found 1 error in 1 file (checked 2 source files)

Is there some kind of MyPy magic/annotations hack that would allow me to explicitly inform MyPy that the descriptors are not converted to enum values at runtime? I don't want to use # type: ignore on each line where the "mappers" are used.

EDIT (on 6th Oct 2021):

Just to rephrase my issue: I don't want to "teach" MyPy that it should detect any missing enum values defined in StrictEnumMapper - its metaclass already does it for me. I want MyPy to properly detect type of the "reaction" field in the Food class example above - it should determine that this field is not Food enum but actually a dict.

The rationale of my approach

(this part is optional, read if you want to understand my idea behind this approach and maybe suggest a better solution that I may not be aware of)

In my projects I often define enums that encapsulate set of possible values, e.g. Django ChoiceField (I'm aware of the TextChoices class but for the sake of this issue lets assume that it wasn't the tool I wanted to use), etc.

I've noticed that often there's a need to define some corresponding values/operations for each of the defined enum value. Usually such actions are defined using separate function, that has branching paths for enum values, e.g.

def perform_action(action_identifier: ActionEnum) -> None:
    if action_identifier == ActionEnum.action1:
        run_action1()
    elif action_identifier == ActionEnum.action2:
        run_action2()
    [...]
    else:
        raise ValueError("Unrecognized enum value")

The problems arise when we update the enum and add new possible values. Programmers have to go through all the code and update places where the enum is used. I have prepared my own solution for this called NamespacedEnum - a simple class inheriting from Enum class, with slightly extended metaclass. This metaclass, along with my descriptors, allows defining containers within class definition. What's more, the metaclass takes care of ensuring that all enum values exist in the mapper. To make it more clear, the above example of perform_action() function would be rewritten using my class like that:

class MyAction(NamespacedEnum):
    action1 = "action1"
    action2 = "action2"
    [...]

    action_for_value = StrictEnumMapper({
        action1: run_action1,
        action2: run_action2,
        [...]
    })


def perform_action(action_identifier: ActionEnum) -> None:
     ActionEnum.action_for_value[action_identifier.value]()

I find it a lot easier to use this approach in projects that rely heavily on enums. I also think that this approach is better than defining separate containers for enum values outside of enum class definition because of the namespace clutter.

Answer 1

If i understand you correctly, you want to statically check that a dictionary contains a key-value pair for every variant of a given enum.

This approach is probably overkill but you can write a custom mypy plugin to support this:

First, the runtime check:

# file: test.py
from enum import Enum
from typing import TypeVar, Dict, Type


E = TypeVar('E', bound='Enum')
T = TypeVar('T')


# check that all enum variants have an entry in the dict
def enum_extras(enum: Type[E], **kwargs: T) -> Dict[E, T]:
    missing = ', '.join(e.name for e in enum if e.name not in kwargs)
    assert not missing, f"Missing enum mappings: {missing}"

    extras = {
        enum[key]: value for key, value in kwargs.items()
    }

    return extras


# a dummy enum
class MyEnum(Enum):
    x = 1
    y = 2


# oups, misspelled the key 'y'
names = enum_extras(MyEnum, x="ex", zzz="why")

This raises AssertionError: Missing enum mappings: y at runtime but mypy does not find the error.

Here is the plugin:

# file: plugin.py
from typing import Optional, Callable, Type

from mypy.plugin import Plugin, FunctionContext
from mypy.types import Type as MypyType


# this is needed for mypy to discover the plugin
def plugin(_: str) -> Type[Plugin]:
    return MyPlugin


# we define a plugin
class MyPlugin(Plugin):
    # this callback looks at every function call
    def get_function_hook(self, fullname: str) -> Optional[Callable[[FunctionContext], MypyType]]:
        if fullname == 'test.enum_extras':
            # return a custom checker for our special function
            return enum_extras_plugin
        return None


# this is called for every function call of `enum_extras`
def enum_extras_plugin(ctx: FunctionContext) -> MypyType:
    # these are the dictionary keys
    _, dict_keys = ctx.arg_names

    # the first argument of the function call (a function that returns an enum class)
    first_argument = ctx.arg_types[0][0]

    try:
        # the names of the enum variants
        enum_keys = first_argument.ret_type.type.names.keys()
    except AttributeError:
        return ctx.default_return_type

    # check that every variant has a key in the dict
    missing = enum_keys - set(dict_keys)

    if missing:
        missing_keys = ', '.join(missing)
        # signal an error if a variant is missing
        ctx.api.fail(f"Missing value for enum variants: {missing_keys}", ctx.context)

    return ctx.default_return_type

To use the plugin we need a mypy.ini:

[mypy]
plugins = plugin.py

And now mypy will find missing keys:

test.py:35: error: Missing value for enum variants: y

Found 1 error in 1 file (checked 1 source file)

Note: I was using python=3.6 and mypy=0.910

Warning: when you pass something unexpected to enum_extras mypy will crash. Some more error handling could help.