[Python] How can I get mypy to handle subclassing EnumType/EnumMeta correctly?

disclaimer: will refer to EnumType as opposed to the older EnumMeta

The following is valid at runtime (and is typically what is recommended to do):

import sys
from enum import Enum, IntEnum
from typing import Any, Type, TypeVar, Union, Callable, cast, TYPE_CHECKING

if sys.version_info >= (3, 11):
from enum import EnumType
else:
from enum import EnumMeta as EnumType

_E = TypeVar("_E", bound=Enum)

class MultipleEnumAccessMeta(EnumType):
"""
Enum Metaclass to provide a way to access multiple values all at once.
"""

def __getitem__(cls: Type[_E], key: Union[str, tuple[str, ...]]) -> Union[_E, list[_E]]:
getitem = cast(Callable[[Type[_E], str], _E], EnumType.__getitem__) # Ensure correct typing for __getitem__

if isinstance(key, tuple):
return [getitem(cls, name) for name in key] # Return list for tuple keys
return getitem(cls, key) # Return single value for single key

if TYPE_CHECKING:
reveal_type(EnumType.__getitem__) # Base method signature
reveal_type(MultipleEnumAccessMeta.__getitem__) # Overridden method signature

# Test Enum with metaclass
class Names(IntEnum, metaclass=MultipleEnumAccessMeta):
Alice = 0
Bob = 1
Charlie = 2

# Test cases
assert Names["Alice"] == Names.Alice
assert Names["Alice", "Bob"] == [Names.Alice, Names.Bob]


However, this gives the following typehint complaints

test.py:17: error: Self argument missing for a non-static method (or an invalid type for self) [misc]
test.py:17: error: Return type "list[Never]" of "__getitem__" incompatible with return type "Never" in supertype "EnumMeta" [override]
test.py:25: note: Revealed type is "def [_EnumMemberT] (type[_EnumMemberT`3], builtins.str) -> _EnumMemberT`3"
test.py:26: note: Revealed type is "def [_E <: enum.Enum] (type[_E`4], Union[builtins.str, tuple[builtins.str, ...]]) -> Union[_E`4, builtins.list[_E`4]]"
test.py:36: error: Enum index should be a string (actual index type "tuple[str, str]") [misc]
test.py:36: error: Non-overlapping equality check (left operand type: "Names", right operand type: "list[Names]") [comparison-overlap]


So to first order, you say "fine", if I drop the cls: Type[_E] portion type-hint, mypy gets more confused:

test.py:17: error: Return type "Union[_E, list[_E]]" of "__getitem__" incompatible with return type "Never" in supertype "EnumMeta" [override]
test.py:21: error: Argument 1 has incompatible type "MultipleEnumAccessMeta"; expected "type[_E]" [arg-type]
test.py:22: error: Argument 1 has incompatible type "MultipleEnumAccessMeta"; expected "type[_E]" [arg-type]
test.py:25: note: Revealed type is "def [_EnumMemberT] (type[_EnumMemberT`3], builtins.str) -> _EnumMemberT`3"
test.py:26: note: Revealed type is "def [_E <: enum.Enum] (atlas_schema.test.MultipleEnumAccessMeta, Union[builtins.str, tuple[builtins.str, ...]]) -> Union[_E`4, builtins.list[_E`4]]"
test.py:36: error: Enum index should be a string (actual index type "tuple[str, str]") [misc]
test.py:36: error: Non-overlapping equality check (left operand type: "Names", right operand type: "list[Names]") [comparison-overlap]


I have questions. Likely this is partially due to the metaclassing that's being very tricky and I'm doing my best to be careful here, but I'm apparently either not careful enough or mypy really has some bugs:

1. why is EnumType.__getitem__ revealed as both returning Never and _EnumMemberT`3 ? This seems like a weird internal conflict with mypy.
2. How do we typehint the cls parameter of MultipleEnumAccessMeta.__getitem__ correctly, both to ensure the TypeVar is evaluated consistently by mypy, but also so that the call to super().__getitem__ [or to EnumType.__getitem__] is correct?

Continue reading...