diff --git a/mypy/plugins/enums.py b/mypy/plugins/enums.py index 8b7c5df6f51f..86c0899ff9ee 100644 --- a/mypy/plugins/enums.py +++ b/mypy/plugins/enums.py @@ -17,11 +17,13 @@ from typing import Final, 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.semanal_enum import ENUM_BASES 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, @@ -87,6 +89,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] @@ -114,6 +129,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 diff --git a/mypy/stubtest.py b/mypy/stubtest.py index 8ea9d786be22..d16e491fb1ab 100644 --- a/mypy/stubtest.py +++ b/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 diff --git a/mypy/types.py b/mypy/types.py index 8ecd2ccf52d9..05b02acc68c0 100644 --- a/mypy/types.py +++ b/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 ) diff --git a/test-data/unit/check-enum.test b/test-data/unit/check-enum.test index 1ab8109eda75..d034fe1a6f5f 100644 --- a/test-data/unit/check-enum.test +++ b/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] diff --git a/test-data/unit/lib-stub/enum.pyi b/test-data/unit/lib-stub/enum.pyi index ccb3818b9d25..5047f7083804 100644 --- a/test-data/unit/lib-stub/enum.pyi +++ b/test-data/unit/lib-stub/enum.pyi @@ -29,6 +29,7 @@ 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 @@ -36,6 +37,8 @@ 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: