如何让mypy正确处理EnumType/EnumMeta的子类化？

免责声明：下文将使用EnumType而非旧版的EnumMeta

以下代码在运行时是有效的（也是通常推荐的写法）：

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]

但这段代码会触发以下类型提示报错：

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]

如果去掉cls: Type[_E]类型提示，mypy会出现更多报错：

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]

疑问 #

1. 为什么EnumType.__getitem__的类型被揭示为同时返回Never和_EnumMemberT？这似乎是mypy的内部冲突。
2. 如何正确为MultipleEnumAccessMeta.__getitem__的cls参数添加类型提示，既保证mypy能正确解析TypeVar，又能让super().__getitem__（或EnumType.__getitem__）的调用类型正确？

解答 #

1. 关于EnumType.__getitem__的类型揭示问题

这是mypy对Enum元类的特殊处理逻辑导致的。EnumMeta（Python 3.11+中更名为EnumType）的__getitem__在mypy的内部类型定义里，基类方法的返回类型被标记为Never，但实际类型检查时，mypy会特殊推断枚举成员的类型，所以reveal_type会显示出实际有效的_EnumMemberT。这并非真正的类型冲突，而是mypy为了适配Enum的特殊行为而做的内部处理。

2. 正确的类型提示写法

需要调整元类方法的self（即cls）参数类型，并处理基类方法的调用类型，同时解决mypy对Enum索引的默认限制：

import sys
from enum import Enum, IntEnum
from typing import TypeVar, Union, cast, TYPE_CHECKING
if sys.version_info >= (3, 11):
    from enum import EnumType
    from typing import Self
else:
    from enum import EnumMeta as EnumType
    from typing_extensions import Self

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

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

    def __getitem__(self: Self, key: Union[str, tuple[str, ...]]) -> Union[_E, list[_E]]:
        if isinstance(key, tuple):
            return [cast(_E, super().__getitem__(name)) for name in key]
        return cast(_E, super().__getitem__(key))

if TYPE_CHECKING:
    reveal_type(EnumType.__getitem__)
    reveal_type(MultipleEnumAccessMeta.__getitem__)

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

# 处理类型检查报错
assert Names["Alice"] == Names.Alice
assert cast(list[Names], Names["Alice", "Bob"]) == [Names.Alice, Names.Bob]
# 或者使用类型忽略：assert Names["Alice", "Bob"] == [Names.Alice, Names.Bob]  # type: ignore[misc,comparison-overlap]

关键调整点：

• 使用Self类型标注元类方法的self参数，让mypy正确识别元类实例的类型
• 直接调用super().__getitem__，通过cast明确返回的枚举成员类型
• 对于tuple索引的断言，使用cast告诉mypy返回的是列表类型，或者添加类型忽略注释

内容的提问来源于stack exchange，提问作者kratsg