Repository: python/mypy
Pull request: Fix inference of IntEnum value attribute type
If the enum value initializer has a NewType type, it should be
reflected in the type of the `value` attribute. It was broken because
the special casing for `__new__` introduced in #10057 didn't consider
the `__new__` in `IntEnum` as special.

Fixes #10411.

Problem statement:
IntEnum with NewType value type regression
The type of `E.A.value` in this program has changed since the last release:

```py
from typing import NewType
from enum import IntEnum

N = NewType("N", int)

class E(IntEnum):
    A = N(0)

# mypy 0.812: N
# master:     int
reveal_type(E.A.value)
```

The original inferred type `N` is more precise, so this looks like a regression. This causes false positives with mypy-protobuf.

Hints from issue discussion:
This was introduced in #10057.

Commit message: Fix inference of IntEnum value attribute type

If the enum value initializer has a NewType type, it should be
reflected in the type of the `value` attribute. It was broken because
the special casing for `__new__` introduced in #10057 didn't consider
the `__new__` in `IntEnum` as special.

Fixes #10411.

Patch to review:
diff --git a/mypy/plugins/enums.py b/mypy/plugins/enums.py
index bd531461d0fc..25936cb6a9a7 100644
--- a/mypy/plugins/enums.py
+++ b/mypy/plugins/enums.py
@@ -112,7 +112,7 @@ def _implements_new(info: TypeInfo) -> bool:
     type_with_new = _first(ti for ti in info.mro if ti.names.get('__new__'))
     if type_with_new is None:
         return False
-    return type_with_new.fullname != 'enum.Enum'
+    return type_with_new.fullname not in ('enum.Enum', 'enum.IntEnum')
 
 
 def enum_value_callback(ctx: 'mypy.plugin.AttributeContext') -> Type:
diff --git a/test-data/unit/check-enum.test b/test-data/unit/check-enum.test
index 06f79776f77f..59a5e54acee5 100644
--- a/test-data/unit/check-enum.test
+++ b/test-data/unit/check-enum.test
@@ -1349,3 +1349,14 @@ b: Literal[E.A, E.B] = e  # E: Incompatible types in assignment (expression has
 c: Literal[E.A, E.C] = e  # E: Incompatible types in assignment (expression has type "E", variable has type "Union[Literal[E.A], Literal[E.C]]")
 b = a  # E: Incompatible types in assignment (expression has type "Union[Literal[E.A], Literal[E.B], Literal[E.C]]", variable has type "Union[Literal[E.A], Literal[E.B]]")
 [builtins fixtures/bool.pyi]
+
+[case testIntEnumWithNewTypeValue]
+from typing import NewType
+from enum import IntEnum
+
+N = NewType("N", int)
+
+class E(IntEnum):
+    A = N(0)
+
+reveal_type(E.A.value) # N: Revealed type is "__main__.N"
diff --git a/test-data/unit/lib-stub/enum.pyi b/test-data/unit/lib-stub/enum.pyi
index 8d0e5fce291a..2c1b03532a5b 100644
--- a/test-data/unit/lib-stub/enum.pyi
+++ b/test-data/unit/lib-stub/enum.pyi
@@ -27,6 +27,7 @@ class Enum(metaclass=EnumMeta):
 
 class IntEnum(int, Enum):
     value: int
+    def __new__(cls: Type[_T], value: Union[int, _T]) -> _T: ...
 
 def unique(enumeration: _T) -> _T: pass
 


Code review comments:
diff_hunk: @@ -112,7 +112,7 @@ def _implements_new(info: TypeInfo) -> bool:
     type_with_new = _first(ti for ti in info.mro if ti.names.get('__new__'))
     if type_with_new is None:
         return False
-    return type_with_new.fullname != 'enum.Enum'
+    return type_with_new.fullname not in ('enum.Enum', 'enum.IntEnum')
path: mypy/plugins/enums.py
text: @user1:
Maybe also consider StrEnum, which is being added in 3.10
```suggestion
    return type_with_new.fullname not in ('enum.Enum', 'enum.IntEnum', 'enum.StrEnum')
```