8000 Check `Enum` definition for invalid base classes (#12026) · python/mypy@1321e9e · GitHub
[go: up one dir, main page]
Skip to content

Commit 1321e9e

Browse files
sobolevnsobolevn
authored
Check Enum definition for invalid base classes (#12026)
Closes #11948
1 parent a8b6d6f commit 1321e9e

File tree

2 files changed

+98
-0
lines changed

2 files changed

+98
-0
lines changed

‎mypy/checker.py

Lines changed: 33 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1837,6 +1837,7 @@ def visit_class_def(self, defn: ClassDef) -> None:
18371837
for base in defn.info.mro[1:-1]: # we don't need self and `object`
18381838
if base.is_enum and base.fullname not in ENUM_BASES:
18391839
self.check_final_enum(defn, base)
1840+
self.check_enum_bases(defn)
18401841

18411842
def check_final_deletable(self, typ: TypeInfo) -> None:
18421843
# These checks are only for mypyc. Only perform some checks that are easier
@@ -1931,6 +1932,38 @@ def is_final_enum_value(self, sym: SymbolTableNode) -> bool:
19311932
return True
19321933
return False
19331934

1935+
def check_enum_bases(self, defn: ClassDef) -> None:
1936+
enum_base: Optional[Instance] = None
1937+
data_base: Optional[Instance] = None
1938+
for base in defn.info.bases:
1939+
if enum_base is None and base.type.fullname in ENUM_BASES:
1940+
enum_base = base
1941+
continue
1942+
elif enum_base is not None:
1943+
self.fail(
1944+
'No base classes are allowed after "{}"'.format(enum_base),
1945+
defn,
1946+
)
1947+
break
1948+
1949+
# This might not be 100% correct, because runtime `__new__`
1950+
# and `__new__` from `typeshed` are sometimes different,
1951+
# but it is good enough.
1952+
new_method = base.type.get('__new__')
1953+
if (data_base is None
1954+
and enum_base is None # data type is always before `Enum`
1955+
and new_method
1956+
and new_method.node
1957+
and new_method.node.fullname != 'builtins.object.__new__'):
1958+
data_base = base
1959+
continue
1960+
elif data_base is not None:
1961+
self.fail(
1962+
'Only a single data type mixin is allowed for Enum subtypes, '
1963+
'found extra "{}"'.format(base),
1964+
defn,
1965+
)
1966+
19341967
def check_protocol_variance(self, defn: ClassDef) -> None:
19351968
"""Check that protocol definition is compatible with declared
19361969
variances of type variables.

‎test-data/unit/check-enum.test

Lines changed: 65 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1883,6 +1883,71 @@ class SubWithOverload(WithOverload): # Should pass
18831883
pass
18841884
[builtins fixtures/tuple.pyi]
18851885

1886+
[case testEnumBaseClassesOrder]
1887+
import enum
1888+
1889+
# Base types:
1890+
1891+
class First:
1892+
def __new__(cls, val):
1893+
pass
1894+
1895+
class Second:
1896+
def __new__(cls, val):
1897+
pass
1898+
1899+
class Third:
1900+
def __new__(cls, val):
1901+
pass
1902+
1903+
class Mixin:
1904+
pass
1905+
1906+
# Correct Enums:
1907+
1908+
class Correct1(Mixin, First, enum.Enum):
1909+
pass
1910+
1911+
class Correct2(First, enum.Enum):
1912+
pass
1913+
1914+
class Correct3(Mixin, enum.Enum):
1915+
pass
1916+
1917+
class RegularClass(Mixin, First, Second):
1918+
pass
1919+
1920+
# Wrong Enums:
1921+
1922+
class MixinAfterEnum1(enum.Enum, Mixin): # E: No base classes are allowed after "enum.Enum"
1923+
pass
1924+
1925+
class MixinAfterEnum2(First, enum.Enum, Mixin): # E: No base classes are allowed after "enum.Enum"
1926+
pass
1927+
1928+
class TwoDataTypes(First, Second, enum.Enum): # E: Only a single data type mixin is allowed for Enum subtypes, found extra "__main__.Second"
1929+
pass
1930+
1931+
class TwoDataTypesAndIntEnumMixin(First, Second, enum.IntEnum, Mixin): # E: Only a single data type mixin is allowed for Enum subtypes, found extra "__main__.Second" \
1932+
# E: No base classes are allowed after "enum.IntEnum"
1933+
pass
1934+
1935+
class ThreeDataTypes(First, Second, Third, enum.Enum): # E: Only a single data type mixin is allowed for Enum subtypes, found extra "__main__.Second" \
1936+
# E: # E: Only a single data type mixin is allowed for Enum subtypes, found extra "__main__.Third"
1937+
pass
1938+
1939+
class ThreeDataTypesAndMixin(First, Second, Third, enum.Enum, Mixin): # E: Only a single data type mixin is allowed for Enum subtypes, found extra "__main__.Second" \
1940+
# E: # E: Only a single data type mixin is allowed for Enum subtypes, found extra "__main__.Third" \
1941+
# E: No base classes are allowed after "enum.Enum"
1942+
pass
1943+
1944+
class FromEnumAndOther1(Correct2, Second, enum.Enum): # E: Only a single data type mixin is allowed for Enum subtypes, found extra "__main__.Second"
1945+
pass
1946+
1947+
class FromEnumAndOther2(Correct2, Second): # E: Only a single data type mixin is allowed for Enum subtypes, found extra "__main__.Second"
1948+
pass
1949+
[builtins fixtures/tuple.pyi]
1950+
18861951
[case testEnumtValueUnionSimplification]
18871952
from enum import IntEnum
18881953
from typing import Any

0 commit comments

Comments
 (0)
0