Use _value_ as a fallback for ellipsis Enum members

Author: sterliakov

mypy/plugins/enums.py

@@ -17,7 +17,7 @@
 from typing import Final, TypeVar, cast
 
 import mypy.plugin  # To avoid circular imports.
-from mypy.nodes import TypeInfo
+from mypy.nodes import TypeInfo, Var
 from mypy.semanal_enum import ENUM_BASES
 from mypy.subtypes import is_equivalent
 from mypy.typeops import fixup_partial_type, make_simplified_union
@@ -87,6 +87,20 @@ def _infer_value_type_with_auto_fallback(
     if proper_type is None:
         return None
     proper_type = get_proper_type(fixup_partial_type(proper_type))
+    # Enums in stubs may have ... instead of actual values. If `_value_` is annotated
+    # (manually or inherited from IntEnum, for example), it is a more reasonable guess
+    # than literal ellipsis type.
+    if isinstance(proper_type, Instance) and proper_type.type.fullname in {
+        "types.EllipsisType",
+        "builtins.ellipsis",
+    }:
+        if (
+            isinstance(ctx.type, Instance)
+            and (value_type := ctx.type.type.get("_value_"))
+            and isinstance(var := value_type.node, Var)
+        ):
+            return var.type
+        return proper_type
     if not (isinstance(proper_type, Instance) and proper_type.type.fullname == "enum.auto"):
         if is_named_instance(proper_type, "enum.member") and proper_type.args:
             return proper_type.args[0]

test-data/unit/check-enum.test

@@ -755,12 +755,10 @@ class B2(IntEnum):
 class B3(IntEnum):
     x = 1
 
-# TODO: getting B1.x._value_ and B2.x._value_ to have type 'int' requires a typeshed change
-
 is_x(reveal_type(B1.x.name))    # N: Revealed type is "Literal['x']"
 is_x(reveal_type(B1.x._name_))  # N: Revealed type is "Literal['x']"
 reveal_type(B1.x.value)         # N: Revealed type is "builtins.int"
-reveal_type(B1.x._value_)       # N: Revealed type is "Any"
+reveal_type(B1.x._value_)       # N: Revealed type is "builtins.int"
 is_x(reveal_type(B2.x.name))    # N: Revealed type is "Literal['x']"
 is_x(reveal_type(B2.x._name_))  # N: Revealed type is "Literal['x']"
 reveal_type(B2.x.value)         # N: Revealed type is "builtins.int"
@@ -2539,3 +2537,51 @@ def check(thing: Things) -> None:
         return None
     return None  # E: Statement is unreachable
 [builtins fixtures/enum.pyi]
+
+[case testSunderValueType]
+from enum import Enum, IntEnum, StrEnum, Flag, IntFlag
+
+class Basic(Enum):
+    _value_: int
+    FOO = 1
+
+reveal_type(Basic.FOO)  # N: Revealed type is "Literal[__main__.Basic.FOO]?"
+reveal_type(Basic.FOO.value)  # N: Revealed type is "Literal[1]?"
+reveal_type(Basic.FOO._value_)  # N: Revealed type is "builtins.int"
+
+class FromStub(Enum):
+    _value_: int
+    FOO = ...
+
+reveal_type(FromStub.FOO)  # N: Revealed type is "Literal[__main__.FromStub.FOO]?"
+reveal_type(FromStub.FOO.value)  # N: Revealed type is "builtins.int"
+reveal_type(FromStub.FOO._value_)  # N: Revealed type is "builtins.int"
+
+class InheritedInt(IntEnum):
+    FOO = ...
+
+reveal_type(InheritedInt.FOO)  # N: Revealed type is "Literal[__main__.InheritedInt.FOO]?"
+reveal_type(InheritedInt.FOO.value)  # N: Revealed type is "builtins.int"
+reveal_type(InheritedInt.FOO._value_)  # N: Revealed type is "builtins.int"
+
+class InheritedStr(StrEnum):
+    FOO = ...
+
+reveal_type(InheritedStr.FOO)  # N: Revealed type is "Literal[__main__.InheritedStr.FOO]?"
+reveal_type(InheritedStr.FOO.value)  # N: Revealed type is "builtins.str"
+reveal_type(InheritedStr.FOO._value_)  # N: Revealed type is "builtins.str"
+
+class InheritedFlag(Flag):
+    FOO = ...
+
+reveal_type(InheritedFlag.FOO)  # N: Revealed type is "Literal[__main__.InheritedFlag.FOO]?"
+reveal_type(InheritedFlag.FOO.value)  # N: Revealed type is "builtins.int"
+reveal_type(InheritedFlag.FOO._value_)  # N: Revealed type is "builtins.int"
+
+class InheritedIntFlag(IntFlag):
+    FOO = ...
+
+reveal_type(InheritedIntFlag.FOO)  # N: Revealed type is "Literal[__main__.InheritedIntFlag.FOO]?"
+reveal_type(InheritedIntFlag.FOO.value)  # N: Revealed type is "builtins.int"
+reveal_type(InheritedIntFlag.FOO._value_)  # N: Revealed type is "builtins.int"
+[builtins fixtures/enum.pyi]

test-data/unit/lib-stub/enum.pyi

@@ -29,13 +29,16 @@ class Enum(metaclass=EnumMeta):
 
 class IntEnum(int, Enum):
     value: int
+    _value_: int
     def __new__(cls: Type[_T], value: Union[int, _T]) -> _T: ...
 
 def unique(enumeration: _T) -> _T: pass
 
 # In reality Flag and IntFlag are 3.6 only
 
 class Flag(Enum):
+    value: int
+    _value_: int
     def __or__(self: _T, other: Union[int, _T]) -> _T: pass
 
 
@@ -49,6 +52,8 @@ class auto(IntFlag):
 
 # It is python-3.11+ only:
 class StrEnum(str, Enum):
+    _value_: str
+    value: str
     def __new__(cls: Type[_T], value: str | _T) -> _T: ...
 
 # It is python-3.11+ only: