Support overriding dunder attributes in Enum subclass (#12138)

Allows any dunder (`__name__`) attributes except `__members__` (due to it being 
read-only) to be overridden in enum subclasses.

Fixes #12132.

Author: Petter Friberg

mypy/checker.py

@@ -1829,11 +1829,7 @@ def visit_class_def(self, defn: ClassDef) -> None:
         if typ.is_protocol and typ.defn.type_vars:
             self.check_protocol_variance(defn)
         if not defn.has_incompatible_baseclass and defn.info.is_enum:
-            for base in defn.info.mro[1:-1]:  # we don't need self and `object`
-                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)
+            self.check_enum(defn)
 
     def check_final_deletable(self, typ: TypeInfo) -> None:
         # These checks are only for mypyc. Only perform some checks that are easier
@@ -1891,6 +1887,24 @@ def check_init_subclass(self, defn: ClassDef) -> None:
             # all other bases have already been checked.
             break
 
+    def check_enum(self, defn: ClassDef) -> None:
+        assert defn.info.is_enum
+        if defn.info.fullname not in ENUM_BASES:
+            for sym in defn.info.names.values():
+                if (isinstance(sym.node, Var) and sym.node.has_explicit_value and
+                        sym.node.name == '__members__'):
+                    # `__members__` will always be overwritten by `Enum` and is considered
+                    # read-only so we disallow assigning a value to it
+                    self.fail(
+                        message_registry.ENUM_MEMBERS_ATTR_WILL_BE_OVERRIDEN, sym.node
+                    )
+        for base in defn.info.mro[1:-1]:  # we don't need self and `object`
+            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_enum(self, defn: ClassDef, base: TypeInfo) -> None:
         for sym in base.names.values():
             if self.is_final_enum_value(sym):

mypy/message_registry.py

@@ -204,6 +204,11 @@ def format(self, *args: object, **kwargs: object) -> "ErrorMessage":
 )
 CANNOT_MAKE_DELETABLE_FINAL: Final = ErrorMessage("Deletable attribute cannot be final")
 
+# Enum
+ENUM_MEMBERS_ATTR_WILL_BE_OVERRIDEN: Final = ErrorMessage(
+    'Assigned "__members__" will be overriden by "Enum" internally'
+)
+
 # ClassVar
 CANNOT_OVERRIDE_INSTANCE_VAR: Final = ErrorMessage(
     'Cannot override instance variable (previously declared on base class "{}") with class '

mypy/semanal.py

@@ -116,6 +116,7 @@
 )
 from mypy.util import (
     correct_relative_import, unmangle, module_prefix, is_typeshed_file, unnamed_function,
+    is_dunder,
 )
 from mypy.scope import Scope
 from mypy.semanal_shared import (
@@ -2473,8 +2474,9 @@ def store_final_status(self, s: AssignmentStmt) -> None:
                     cur_node = self.type.names.get(lval.name, None)
                     if (cur_node and isinstance(cur_node.node, Var) and
                             not (isinstance(s.rvalue, TempNode) and s.rvalue.no_rhs)):
-                        cur_node.node.is_final = True
-                        s.is_final_def = True
+                        # Double underscored members are writable on an `Enum`.
+                        # (Except read-only `__members__` but that is handled in type checker)
+                        cur_node.node.is_final = s.is_final_def = not is_dunder(cur_node.node.name)
 
                 # Special case: deferred initialization of a final attribute in __init__.
                 # In this case we just pretend this is a valid final definition to suppress

test-data/unit/check-enum.test

@@ -2029,3 +2029,54 @@ class C(Enum):
 
 C._ignore_ # E: "Type[C]" has no attribute "_ignore_"
 [typing fixtures/typing-medium.pyi]
+
+[case testCanOverrideDunderAttributes]
+import typing
+from enum import Enum, Flag
+
+class BaseEnum(Enum):
+    __dunder__ = 1
+    __labels__: typing.Dict[int, str]
+
+class Override(BaseEnum):
+    __dunder__ = 2
+    __labels__ = {1: "1"}
+
+Override.__dunder__ = 3
+BaseEnum.__dunder__ = 3
+Override.__labels__ = {2: "2"}
+
+class FlagBase(Flag):
+    __dunder__ = 1
+    __labels__: typing.Dict[int, str]
+
+class FlagOverride(FlagBase):
+    __dunder__ = 2
+    __labels = {1: "1"}
+
+FlagOverride.__dunder__ = 3
+FlagBase.__dunder__ = 3
+FlagOverride.__labels__ = {2: "2"}
+[builtins fixtures/dict.pyi]
+
+[case testCanNotInitialize__members__]
+import typing
+from enum import Enum
+
+class WritingMembers(Enum):
+    __members__: typing.Dict[Enum, Enum] = {}  # E: Assigned "__members__" will be overriden by "Enum" internally
+
+class OnlyAnnotatedMembers(Enum):
+    __members__: typing.Dict[Enum, Enum]
+[builtins fixtures/dict.pyi]
+
+[case testCanOverrideDunderOnNonFirstBaseEnum]
+import typing
+from enum import Enum
+
+class Some:
+    __labels__: typing.Dict[int, str]
+
+class A(Some, Enum):
+    __labels__ = {1: "1"}
+[builtins fixtures/dict.pyi]