Skip to content

Commit d9c589f

Browse files
committed
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.
1 parent 8bc1115 commit d9c589f

File tree

3 files changed

+13
-1
lines changed

3 files changed

+13
-1
lines changed

mypy/plugins/enums.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -112,7 +112,7 @@ def _implements_new(info: TypeInfo) -> bool:
112112
type_with_new = _first(ti for ti in info.mro if ti.names.get('__new__'))
113113
if type_with_new is None:
114114
return False
115-
return type_with_new.fullname != 'enum.Enum'
115+
return type_with_new.fullname not in ('enum.Enum', 'enum.IntEnum')
116116

117117

118118
def enum_value_callback(ctx: 'mypy.plugin.AttributeContext') -> Type:

test-data/unit/check-enum.test

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1349,3 +1349,14 @@ b: Literal[E.A, E.B] = e # E: Incompatible types in assignment (expression has
13491349
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]]")
13501350
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]]")
13511351
[builtins fixtures/bool.pyi]
1352+
1353+
[case testIntEnumWithNewTypeValue]
1354+
from typing import NewType
1355+
from enum import IntEnum
1356+
1357+
N = NewType("N", int)
1358+
1359+
class E(IntEnum):
1360+
A = N(0)
1361+
1362+
reveal_type(E.A.value) # N: Revealed type is "__main__.N"

test-data/unit/lib-stub/enum.pyi

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,7 @@ class Enum(metaclass=EnumMeta):
2727

2828
class IntEnum(int, Enum):
2929
value: int
30+
def __new__(cls: Type[_T], value: Union[int, _T]) -> _T: ...
3031

3132
def unique(enumeration: _T) -> _T: pass
3233

0 commit comments

Comments
 (0)