Skip to content

Commit

Permalink
Fixes enum value inference in cases where the value type is of a
Browse files Browse the repository at this point in the history
user-defined data type class (with __new__).

This fixes a regression introduced by python#10057 to fix python#10000.  The
`not ti.fullname.startswith("builtins.") clause seemed to be intended
to catch enums with a built-in data type like int or bytes, but this is
overly broad.  It should allow any type so long as it is not itself an
enum.Enum subclass.
  • Loading branch information
embray committed Oct 24, 2023
1 parent 8236c93 commit 87d1569
Show file tree
Hide file tree
Showing 2 changed files with 20 additions and 2 deletions.
8 changes: 6 additions & 2 deletions mypy/plugins/enums.py
Original file line number Diff line number Diff line change
Expand Up @@ -104,13 +104,17 @@ def _infer_value_type_with_auto_fallback(

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
subclass of enum.Enum. In the latter case, we must infer Any as long as mypy can't infer
the type of _value_ from assignments in __new__.
If, however, __new__ comes from a user-defined class that is not an Enum subclass (i.e.
the data type) this is allowed, because we should in general infer that an enum entry's
value has that type.
"""
type_with_new = _first(
ti
for ti in info.mro
if ti.names.get("__new__") and not ti.fullname.startswith("builtins.")
if ti.is_enum and ti.names.get("__new__")
)
if type_with_new is None:
return False
Expand Down
14 changes: 14 additions & 0 deletions test-data/unit/check-enum.test
Original file line number Diff line number Diff line change
Expand Up @@ -1366,6 +1366,20 @@ reveal_type(a._value_) # N: Revealed type is "Any"
[builtins fixtures/primitives.pyi]
[typing fixtures/typing-medium.pyi]

[case testValueTypeWithUserDataType]
from enum import Enum
from typing import Any

class Data:
def __new__(cls, value: Any) -> Data: pass

class DataEnum(Data, Enum):
A = Data(1)

reveal_type(DataEnum.A) # N: Revealed type is "Literal[__main__.DataEnum.A]?"
reveal_type(DataEnum.A.value) # N: Revealed type is "__main__.Data"
reveal_type(DataEnum.A._value_) # N: Revealed type is "__main__.Data"

[case testEnumNarrowedToTwoLiterals]
# Regression test: two literals of an enum would be joined
# as the full type, regardless of the amount of elements
Expand Down

0 comments on commit 87d1569

Please sign in to comment.