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.