Support _value_ as a fallback for ellipsis Enum members (#19352)

Fixes #19334. 

This does not affect enums with explicit values different from ellipsis
and is limited to enums defined in stub files.

Author: sterliakov

mypy/plugins/enums.py

@@ -17,10 +17,12 @@
 from typing import TypeVar, cast
 
 import mypy.plugin  # To avoid circular imports.
-from mypy.nodes import TypeInfo
+from mypy.checker import TypeChecker
+from mypy.nodes import TypeInfo, Var
 from mypy.subtypes import is_equivalent
 from mypy.typeops import fixup_partial_type, make_simplified_union
 from mypy.types import (
+    ELLIPSIS_TYPE_NAMES,
     CallableType,
     Instance,
     LiteralType,
@@ -79,6 +81,19 @@ 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 (
+        _is_defined_in_stub(ctx)
+        and isinstance(proper_type, Instance)
+        and proper_type.type.fullname in ELLIPSIS_TYPE_NAMES
+        and isinstance(ctx.type, Instance)
+    ):
+        value_type = ctx.type.type.get("_value_")
+        if value_type is not None 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]
@@ -106,6 +121,11 @@ def _infer_value_type_with_auto_fallback(
     return ctx.default_attr_type
 
 
+def _is_defined_in_stub(ctx: mypy.plugin.AttributeContext) -> bool:
+    assert isinstance(ctx.api, TypeChecker)
+    return isinstance(ctx.type, Instance) and ctx.api.modules[ctx.type.type.module_name].is_stub
+
+
 def _implements_new(info: TypeInfo) -> bool:
     """Check whether __new__ comes from enum.Enum or was implemented in a
     subclass. In the latter case, we must infer Any as long as mypy can't infer

mypy/stubtest.py

@@ -1149,7 +1149,7 @@ def verify_var(
             proper_type = mypy.types.get_proper_type(stub.type)
             if (
                 isinstance(proper_type, mypy.types.Instance)
-                and proper_type.type.fullname == "builtins.ellipsis"
+                and proper_type.type.fullname in mypy.types.ELLIPSIS_TYPE_NAMES
             ):
                 should_error = False
 

mypy/types.py

@@ -181,6 +181,8 @@
 # Supported @override decorator names.
 OVERRIDE_DECORATOR_NAMES: Final = ("typing.override", "typing_extensions.override")
 
+ELLIPSIS_TYPE_NAMES: Final = ("builtins.ellipsis", "types.EllipsisType")
+
 # A placeholder used for Bogus[...] parameters
 _dummy: Final[Any] = object()
 
@@ -1574,7 +1576,7 @@ def is_singleton_type(self) -> bool:
         return (
             self.type.is_enum
             and len(self.type.enum_members) == 1
-            or self.type.fullname in {"builtins.ellipsis", "types.EllipsisType"}
+            or self.type.fullname in ELLIPSIS_TYPE_NAMES
         )
 
 

test-data/unit/check-enum.test

@@ -618,9 +618,8 @@ reveal_type(B.a)        # N: Revealed type is "Literal[__main__.B.a]?"
 reveal_type(A.x.name)   # N: Revealed type is "Literal['x']?"
 reveal_type(B.a.name)   # N: Revealed type is "Literal['a']?"
 
-# TODO: The revealed type should be 'int' here
-reveal_type(A.x.value)  # N: Revealed type is "Any"
-reveal_type(B.a.value)  # N: Revealed type is "Any"
+reveal_type(A.x.value)  # N: Revealed type is "builtins.int"
+reveal_type(B.a.value)  # N: Revealed type is "builtins.int"
 [builtins fixtures/enum.pyi]
 
 [case testAnonymousFunctionalEnum]
@@ -755,12 +754,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"
@@ -770,9 +767,6 @@ is_x(reveal_type(B3.x._name_))  # N: Revealed type is "Literal['x']"
 reveal_type(B3.x.value)         # N: Revealed type is "Literal[1]?"
 reveal_type(B3.x._value_)       # N: Revealed type is "Literal[1]?"
 
-# TODO: C1.x.value and C2.x.value should also be of type 'int'
-# This requires either a typeshed change or a plugin refinement
-
 C1 = IntFlag('C1', 'x')
 class C2(IntFlag):
     x = auto()
@@ -781,8 +775,8 @@ class C3(IntFlag):
 
 is_x(reveal_type(C1.x.name))    # N: Revealed type is "Literal['x']"
 is_x(reveal_type(C1.x._name_))  # N: Revealed type is "Literal['x']"
-reveal_type(C1.x.value)         # N: Revealed type is "Any"
-reveal_type(C1.x._value_)       # N: Revealed type is "Any"
+reveal_type(C1.x.value)         # N: Revealed type is "builtins.int"
+reveal_type(C1.x._value_)       # N: Revealed type is "builtins.int"
 is_x(reveal_type(C2.x.name))    # N: Revealed type is "Literal['x']"
 is_x(reveal_type(C2.x._name_))  # N: Revealed type is "Literal['x']"
 reveal_type(C2.x.value)         # N: Revealed type is "builtins.int"
@@ -800,8 +794,8 @@ class D3(Flag):
 
 is_x(reveal_type(D1.x.name))    # N: Revealed type is "Literal['x']"
 is_x(reveal_type(D1.x._name_))  # N: Revealed type is "Literal['x']"
-reveal_type(D1.x.value)         # N: Revealed type is "Any"
-reveal_type(D1.x._value_)       # N: Revealed type is "Any"
+reveal_type(D1.x.value)         # N: Revealed type is "builtins.int"
+reveal_type(D1.x._value_)       # N: Revealed type is "builtins.int"
 is_x(reveal_type(D2.x.name))    # N: Revealed type is "Literal['x']"
 is_x(reveal_type(D2.x._name_))  # N: Revealed type is "Literal['x']"
 reveal_type(D2.x.value)         # N: Revealed type is "builtins.int"
@@ -2539,3 +2533,105 @@ def check(thing: Things) -> None:
         return None
     return None  # E: Statement is unreachable
 [builtins fixtures/enum.pyi]
+
+[case testSunderValueTypeEllipsis]
+from foo.bar import (
+    Basic, FromStub, InheritedInt, InheritedStr, InheritedFlag,
+    InheritedIntFlag, Wrapper
+)
+
+reveal_type(Basic.FOO)  # N: Revealed type is "Literal[foo.bar.Basic.FOO]?"
+reveal_type(Basic.FOO.value)  # N: Revealed type is "Literal[1]?"
+reveal_type(Basic.FOO._value_)  # N: Revealed type is "builtins.int"
+
+reveal_type(FromStub.FOO)  # N: Revealed type is "Literal[foo.bar.FromStub.FOO]?"
+reveal_type(FromStub.FOO.value)  # N: Revealed type is "builtins.int"
+reveal_type(FromStub.FOO._value_)  # N: Revealed type is "builtins.int"
+
+reveal_type(Wrapper.Nested.FOO)  # N: Revealed type is "Literal[foo.bar.Wrapper.Nested.FOO]?"
+reveal_type(Wrapper.Nested.FOO.value)  # N: Revealed type is "builtins.int"
+reveal_type(Wrapper.Nested.FOO._value_)  # N: Revealed type is "builtins.int"
+
+reveal_type(InheritedInt.FOO)  # N: Revealed type is "Literal[foo.bar.InheritedInt.FOO]?"
+reveal_type(InheritedInt.FOO.value)  # N: Revealed type is "builtins.int"
+reveal_type(InheritedInt.FOO._value_)  # N: Revealed type is "builtins.int"
+
+reveal_type(InheritedStr.FOO)  # N: Revealed type is "Literal[foo.bar.InheritedStr.FOO]?"
+reveal_type(InheritedStr.FOO.value)  # N: Revealed type is "builtins.str"
+reveal_type(InheritedStr.FOO._value_)  # N: Revealed type is "builtins.str"
+
+reveal_type(InheritedFlag.FOO)  # N: Revealed type is "Literal[foo.bar.InheritedFlag.FOO]?"
+reveal_type(InheritedFlag.FOO.value)  # N: Revealed type is "builtins.int"
+reveal_type(InheritedFlag.FOO._value_)  # N: Revealed type is "builtins.int"
+
+reveal_type(InheritedIntFlag.FOO)  # N: Revealed type is "Literal[foo.bar.InheritedIntFlag.FOO]?"
+reveal_type(InheritedIntFlag.FOO.value)  # N: Revealed type is "builtins.int"
+reveal_type(InheritedIntFlag.FOO._value_)  # N: Revealed type is "builtins.int"
+
+[file foo/__init__.pyi]
+[file foo/bar/__init__.pyi]
+from enum import Enum, IntEnum, StrEnum, Flag, IntFlag
+
+class Basic(Enum):
+    _value_: int
+    FOO = 1
+
+class FromStub(Enum):
+    _value_: int
+    FOO = ...
+
+class Wrapper:
+    class Nested(Enum):
+        _value_: int
+        FOO = ...
+
+class InheritedInt(IntEnum):
+    FOO = ...
+
+class InheritedStr(StrEnum):
+    FOO = ...
+
+class InheritedFlag(Flag):
+    FOO = ...
+
+class InheritedIntFlag(IntFlag):
+    FOO = ...
+[builtins fixtures/enum.pyi]
+
+[case testSunderValueTypeEllipsisNonStub]
+from enum import Enum, StrEnum
+
+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"
+
+# TODO: this and below should produce diagnostics, Ellipsis is not assignable to int
+# Now we do not check members against _value_ at all.
+
+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.ellipsis"
+reveal_type(FromStub.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.ellipsis"
+reveal_type(InheritedStr.FOO._value_)  # N: Revealed type is "builtins.ellipsis"
+
+class Wrapper:
+    class Nested(StrEnum):
+        FOO = ...
+
+reveal_type(Wrapper.Nested.FOO)  # N: Revealed type is "Literal[__main__.Wrapper.Nested.FOO]?"
+reveal_type(Wrapper.Nested.FOO.value)  # N: Revealed type is "builtins.ellipsis"
+reveal_type(Wrapper.Nested.FOO._value_)  # N: Revealed type is "builtins.ellipsis"
+[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: