diff --git a/mypy/plugins/enums.py b/mypy/plugins/enums.py index 7869a8b5cdfa..388d72f25001 100644 --- a/mypy/plugins/enums.py +++ b/mypy/plugins/enums.py @@ -104,14 +104,14 @@ 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.") - ) + type_with_new = _first(ti for ti in info.mro if ti.is_enum and ti.names.get("__new__")) if type_with_new is None: return False return type_with_new.fullname not in ("enum.Enum", "enum.IntEnum", "enum.StrEnum") diff --git a/test-data/unit/check-enum.test b/test-data/unit/check-enum.test index 6779ae266454..9563a7ec88c8 100644 --- a/test-data/unit/check-enum.test +++ b/test-data/unit/check-enum.test @@ -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