diff --git a/mypy/checker.py b/mypy/checker.py index db229e274ab7..851f23185f4f 100644 --- a/mypy/checker.py +++ b/mypy/checker.py @@ -1838,6 +1838,7 @@ def visit_class_def(self, defn: ClassDef) -> None: if base.is_enum and base.fullname not in ENUM_BASES: self.check_final_enum(defn, base) self.check_enum_bases(defn) + self.check_enum_new(defn) def check_final_deletable(self, typ: TypeInfo) -> None: # These checks are only for mypyc. Only perform some checks that are easier @@ -1934,9 +1935,8 @@ def is_final_enum_value(self, sym: SymbolTableNode) -> bool: def check_enum_bases(self, defn: ClassDef) -> None: enum_base: Optional[Instance] = None - data_base: Optional[Instance] = None for base in defn.info.bases: - if enum_base is None and base.type.fullname in ENUM_BASES: + if enum_base is None and base.type.is_enum: enum_base = base continue elif enum_base is not None: @@ -1946,23 +1946,36 @@ def check_enum_bases(self, defn: ClassDef) -> None: ) break - # This might not be 100% correct, because runtime `__new__` - # and `__new__` from `typeshed` are sometimes different, - # but it is good enough. - new_method = base.type.get('__new__') - if (data_base is None - and enum_base is None # data type is always before `Enum` - and new_method - and new_method.node - and new_method.node.fullname != 'builtins.object.__new__'): - data_base = base - continue - elif data_base is not None: + def check_enum_new(self, defn: ClassDef) -> None: + def has_new_method(info: TypeInfo) -> bool: + new_method = info.get('__new__') + return bool( + new_method + and new_method.node + and new_method.node.fullname != 'builtins.object.__new__' + ) + + has_new = False + for base in defn.info.bases: + candidate = False + + if base.type.is_enum: + # If we have an `Enum`, then we need to check all its bases. + candidate = any( + not b.is_enum and has_new_method(b) + for b in base.type.mro[1:-1] + ) + else: + candidate = has_new_method(base.type) + + if candidate and has_new: self.fail( 'Only a single data type mixin is allowed for Enum subtypes, ' 'found extra "{}"'.format(base), defn, ) + elif candidate: + has_new = True def check_protocol_variance(self, defn: ClassDef) -> None: """Check that protocol definition is compatible with declared diff --git a/test-data/unit/check-enum.test b/test-data/unit/check-enum.test index e1dec9aec86f..78460909cd6a 100644 --- a/test-data/unit/check-enum.test +++ b/test-data/unit/check-enum.test @@ -1905,6 +1905,9 @@ class Mixin: # Correct Enums: +class Correct0(enum.Enum): + pass + class Correct1(Mixin, First, enum.Enum): pass @@ -1917,8 +1920,29 @@ class Correct3(Mixin, enum.Enum): class RegularClass(Mixin, First, Second): pass +# Correct inheritance: + +class _InheritingDataAndMixin(Correct1): + pass + +class _CorrectWithData(First, Correct0): + pass + +class _CorrectWithDataAndMixin(Mixin, First, Correct0): + pass + +class _CorrectWithMixin(Mixin, Correct2): + pass + # Wrong Enums: +class TwoDataTypesViaInheritance(Second, Correct2): # E: Only a single data type mixin is allowed for Enum subtypes, found extra "__main__.Correct2" + pass + +class TwoDataTypesViaInheritanceAndMixin(Second, Correct2, Mixin): # E: No base classes are allowed after "__main__.Correct2" \ + # E: Only a single data type mixin is allowed for Enum subtypes, found extra "__main__.Correct2" + pass + class MixinAfterEnum1(enum.Enum, Mixin): # E: No base classes are allowed after "enum.Enum" pass @@ -1928,27 +1952,56 @@ class MixinAfterEnum2(First, enum.Enum, Mixin): # E: No base classes are allowe class TwoDataTypes(First, Second, enum.Enum): # E: Only a single data type mixin is allowed for Enum subtypes, found extra "__main__.Second" pass -class TwoDataTypesAndIntEnumMixin(First, Second, enum.IntEnum, Mixin): # E: Only a single data type mixin is allowed for Enum subtypes, found extra "__main__.Second" \ - # E: No base classes are allowed after "enum.IntEnum" +class TwoDataTypesAndIntEnumMixin(First, Second, enum.IntEnum, Mixin): # E: No base classes are allowed after "enum.IntEnum" \ + # E: Only a single data type mixin is allowed for Enum subtypes, found extra "__main__.Second" pass class ThreeDataTypes(First, Second, Third, enum.Enum): # E: Only a single data type mixin is allowed for Enum subtypes, found extra "__main__.Second" \ - # E: # E: Only a single data type mixin is allowed for Enum subtypes, found extra "__main__.Third" + # E: Only a single data type mixin is allowed for Enum subtypes, found extra "__main__.Third" pass -class ThreeDataTypesAndMixin(First, Second, Third, enum.Enum, Mixin): # E: Only a single data type mixin is allowed for Enum subtypes, found extra "__main__.Second" \ - # E: # E: Only a single data type mixin is allowed for Enum subtypes, found extra "__main__.Third" \ - # E: No base classes are allowed after "enum.Enum" +class ThreeDataTypesAndMixin(First, Second, Third, enum.Enum, Mixin): # E: No base classes are allowed after "enum.Enum" \ + # E: Only a single data type mixin is allowed for Enum subtypes, found extra "__main__.Second" \ + # E: Only a single data type mixin is allowed for Enum subtypes, found extra "__main__.Third" pass -class FromEnumAndOther1(Correct2, Second, enum.Enum): # E: Only a single data type mixin is allowed for Enum subtypes, found extra "__main__.Second" +class FromEnumAndOther1(Correct2, Second, enum.Enum): # E: No base classes are allowed after "__main__.Correct2" \ + # E: Only a single data type mixin is allowed for Enum subtypes, found extra "__main__.Second" pass -class FromEnumAndOther2(Correct2, Second): # E: Only a single data type mixin is allowed for Enum subtypes, found extra "__main__.Second" +class FromEnumAndOther2(Correct2, Second): # E: No base classes are allowed after "__main__.Correct2" \ + # E: Only a single data type mixin is allowed for Enum subtypes, found extra "__main__.Second" pass [builtins fixtures/tuple.pyi] -[case testEnumtValueUnionSimplification] +[case testRegression12258] +from enum import Enum + +class MyEnum(Enum): ... + +class BytesEnum(bytes, MyEnum): ... # Should be ok +[builtins fixtures/tuple.pyi] + +[case testEnumWithNewHierarchy] +import enum + +class A: + def __new__(cls, val): ... +class B(A): + def __new__(cls, val): ... +class C: + def __new__(cls, val): ... + +class E1(A, enum.Enum): ... +class E2(B, enum.Enum): ... + +# Errors: + +class W1(C, E1): ... # E: Only a single data type mixin is allowed for Enum subtypes, found extra "__main__.E1" +class W2(C, E2): ... # E: Only a single data type mixin is allowed for Enum subtypes, found extra "__main__.E2" +[builtins fixtures/tuple.pyi] + +[case testEnumValueUnionSimplification] from enum import IntEnum from typing import Any