8000 Support overriding dunder attributes in Enum subclass · python/mypy@360fd99 · GitHub
[go: up one dir, main page]
Skip to content

Commit 360fd99

Browse files
flaeppeflaeppe
committed
Support overriding dunder attributes in Enum subclass
1 parent 9135795 commit 360fd99

File tree

4 files changed

+79
-7
lines changed

4 files changed

+79
-7
lines changed

‎mypy/checker.py

Lines changed: 19 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1834,11 +1834,7 @@ def visit_class_def(self, defn: ClassDef) -> None:
18341834
if typ.is_protocol and typ.defn.type_vars:
18351835
self.check_protocol_variance(defn)
18361836
if not defn.has_incompatible_baseclass and defn.info.is_enum:
1837-
for base in defn.info.mro[1:-1]: # we don't need self and `object`
1838-
if base.is_enum and base.fullname not in ENUM_BASES:
1839-
self.check_final_enum(defn, base)
1840-
self.check_enum_bases(defn)
1841-
self.check_enum_new(defn)
1837+
self.check_enum(defn)
18421838

18431839
def check_final_deletable(self, typ: TypeInfo) -> None:
18441840
# These checks are only for mypyc. Only perform some checks that are easier
@@ -1896,6 +1892,24 @@ def check_init_subclass(self, defn: ClassDef) -> None:
18961892
# all other bases have already been checked.
18971893
break
18981894

1895+
def check_enum(self, defn: ClassDef) -> None:
1896+
assert defn.info.is_enum
1897+
if defn.info.fullname not in ENUM_BASES:
1898+
for sym in defn.info.names.values():
1899+
if (isinstance(sym.node, Var) and sym.node.has_explicit_value and
1900+
sym.node.name == '__members__'):
1901+
# `__members__` will always be overwritten by `Enum` and is considered
1902+
# read-only so we disallow assigning a value to it
1903+
self.fail(
1904+
message_registry.ENUM_MEMBERS_ATTR_WILL_BE_OVERRIDEN, sym.node
1905+
)
1906+
for base in defn.info.mro[1:-1]: # we don't need self and `object`
1907+
if base.is_enum and base.fullname not in ENUM_BASES:
1908+
self.check_final_enum(defn, base)
1909+
1910+
self.check_enum_bases(defn)
1911+
self.check_enum_new(defn)
1912+
18991913
def check_final_enum(self, defn: ClassDef, base: TypeInfo) -> None:
19001914
for sym in base.names.values():
19011915
if self.is_final_enum_value(sym):

‎mypy/message_registry.py

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -205,6 +205,11 @@ def format(self, *args: object, **kwargs: object) -> "ErrorMessage":
205205
)
206206
CANNOT_MAKE_DELETABLE_FINAL: Final = ErrorMessage("Deletable attribute cannot be final")
207207

208+
# Enum
209+
ENUM_MEMBERS_ATTR_WILL_BE_OVERRIDEN: Final = ErrorMessage(
210+
'Assigned "__members__" will be overriden by "Enum" internally'
211+
)
212+
208213
# ClassVar
209214
CANNOT_OVERRIDE_INSTANCE_VAR: Final = ErrorMessage(
210215
'Cannot override instance variable (previously declared on base class "{}") with class '

‎mypy/semanal.py

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -116,6 +116,7 @@
116116
)
117117
from mypy.util import (
118118
correct_relative_import, unmangle, module_prefix, is_typeshed_file, unnamed_function,
119+
is_dunder,
119120
)
120121
from mypy.scope import Scope
121122
from mypy.semanal_shared import (
@@ -2480,8 +2481,9 @@ def store_final_status(self, s: AssignmentStmt) -> None:
24802481
cur_node = self.type.names.get(lval.name, None)
24812482
if (cur_node and isinstance(cur_node.node, Var) and
24822483
not (isinstance(s.rvalue, TempNode) and s.rvalue.no_rhs)):
2483-
cur_node.node.is_final = True
2484-
s.is_final_def = True
2484+
# Double underscored members are writable on an `Enum`.
2485+
# (Except read-only `__members__` but that is handled in type checker)
2486+
cur_node.node.is_final = s.is_final_def = not is_dunder(cur_node.node.name)
24852487

24862488
# Special case: deferred initialization of a final attribute in __init__.
24872489
# In this case we just pretend this is a valid final definition to suppress

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

Lines changed: 51 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2029,3 +2029,54 @@ class C(Enum):
20292029

20302030
C._ignore_ # E: "Type[C]" has no attribute "_ignore_"
20312031
[typing fixtures/typing-medium.pyi]
2032+
2033+
[case testCanOverrideDunderAttributes]
2034+
import typing
2035+
from enum import Enum, Flag
2036+
2037+
class BaseEnum(Enum):
2038+
__dunder__ = 1
2039+
__labels__: typing.Dict[int, str]
2040+
2041+
class Override(BaseEnum):
2042+
__dunder__ = 2
2043+
__labels__ = {1: "1"}
2044+
2045+
Override.__dunder__ = 3
2046+
BaseEnum.__dunder__ = 3
2047+
Override.__labels__ = {2: "2"}
2048+
2049+
class FlagBase(Flag):
2050+
__dunder__ = 1
2051+
__labels__: typing.Dict[int, str]
2052+
2053+
class FlagOverride(FlagBase):
2054+
__dunder__ = 2
2055+
__labels = {1: "1"}
2056+
2057+
FlagOverride.__dunder__ = 3
2058+
FlagBase.__dunder__ = 3
2059+
FlagOverride.__labels__ = {2: "2"}
2060+
[builtins fixtures/dict.pyi]
2061+
2062+
[case testCanNotInitialize__members__]
2063+
import typing
2064+
from enum import Enum
2065+
2066+
class WritingMembers(Enum):
2067+
__members__: typing.Dict[Enum, Enum] = {} # E: Assigned "__members__" will be overriden by "Enum" internally
2068+
2069+
class OnlyAnnotatedMembers(Enum):
2070+
__members__: typing.Dict[Enum, Enum]
2071+
[builtins fixtures/dict.pyi]
2072+
2073+
[case testCanOverrideDunderOnNonFirstBaseEnum]
2074+
import typing
2075+
from enum import Enum
2076+
2077+
class Some:
2078+
__labels__: typing.Dict[int, str]
2079+
2080+
class A(Some, Enum):
2081+
__labels__ = {1: "1"}
2082+
[builtins fixtures/dict.pyi]

0 commit comments

Comments
 (0)
0