Fix enum value inference with user-defined data type mixin (#16320)

Fixes enum value inference in cases where the value type is of a
user-defined data type class (with `__new__`).

According to https://docs.python.org/3/howto/enum.html#others when
defining an `Enum` subclass with a custom data type:

> A data type is a mixin that defines `__new__()`, or a dataclass

This fixes a regression introduced by #10057 to fix #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.

Author: embray

mypy/plugins/enums.py

@@ -128,14 +128,14 @@ def _is_defined_in_stub(ctx: mypy.plugin.AttributeContext) -> bool:
 
 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")

test-data/unit/check-enum.test

@@ -1524,6 +1524,21 @@ 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"
+[builtins fixtures/tuple.pyi]
+
 [case testEnumNarrowedToTwoLiterals]
 # Regression test: two literals of an enum would be joined
 # as the full type, regardless of the amount of elements