Skip to content

Support _value_ as a fallback for ellipsis Enum members #19352

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 2 commits into
base: master
Choose a base branch
from
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
16 changes: 15 additions & 1 deletion mypy/plugins/enums.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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",
}:
Comment on lines +93 to +96
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

There's at least one other use of this set elsewhere. Could consider moving it to a ELLIPSIS_TYPE_NAMES constant

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]
Expand Down
64 changes: 55 additions & 9 deletions test-data/unit/check-enum.test
Original file line number Diff line number Diff line change
Expand Up @@ -619,8 +619,8 @@ 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]
Expand Down Expand Up @@ -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"
Expand All @@ -781,8 +779,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"
Expand All @@ -800,8 +798,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"
Expand Down Expand Up @@ -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]
5 changes: 5 additions & 0 deletions test-data/unit/lib-stub/enum.pyi
Original file line number Diff line number Diff line change
Expand Up @@ -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


Expand All @@ -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:
Expand Down